1 files changed, 164 insertions, 0 deletions
diff --git a/include/clang/AST/LexicallyOrderedRecursiveASTVisitor.h b/include/clang/AST/LexicallyOrderedRecursiveASTVisitor.h
new file mode 100644
index 0000000000000..264f20f19ad5b
--- /dev/null
+++ b/include/clang/AST/LexicallyOrderedRecursiveASTVisitor.h
@@ -0,0 +1,164 @@
+//===--- LexicallyOrderedRecursiveASTVisitor.h - ----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines the LexicallyOrderedRecursiveASTVisitor interface, which
+//  recursively traverses the entire AST in a lexical order.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_AST_LEXICALLY_ORDERED_RECURSIVEASTVISITOR_H
+#define LLVM_CLANG_AST_LEXICALLY_ORDERED_RECURSIVEASTVISITOR_H
+
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Basic/LLVM.h"
+#include "clang/Basic/SourceManager.h"
+#include "llvm/Support/SaveAndRestore.h"
+
+namespace clang {
+
+/// A RecursiveASTVisitor subclass that guarantees that AST traversal is
+/// performed in a lexical order (i.e. the order in which declarations are
+/// written in the source).
+///
+/// RecursiveASTVisitor doesn't guarantee lexical ordering because there are
+/// some declarations, like Objective-C @implementation declarations
+/// that might be represented in the AST differently to how they were written
+/// in the source.
+/// In particular, Objective-C @implementation declarations may contain
+/// non-Objective-C declarations, like functions:
+///
+///   @implementation MyClass
+///
+///   - (void) method { }
+///   void normalFunction() { }
+///
+///   @end
+///
+/// Clang's AST stores these declarations outside of the @implementation
+/// declaration, so the example above would be represented using the following
+/// AST:
+///   |-ObjCImplementationDecl ... MyClass
+///   | `-ObjCMethodDecl ... method
+///   |    ...
+///   `-FunctionDecl ... normalFunction
+///       ...
+///
+/// This class ensures that these declarations are traversed before the
+/// corresponding TraverseDecl for the @implementation returns. This ensures
+/// that the lexical parent relationship between these declarations and the
+/// @implementation is preserved while traversing the AST. Note that the
+/// current implementation doesn't mix these declarations with the declarations
+/// contained in the @implementation, so the traversal of all of the
+/// declarations in the @implementation still doesn't follow the lexical order.
+template <typename Derived>
+class LexicallyOrderedRecursiveASTVisitor
+    : public RecursiveASTVisitor<Derived> {
+  using BaseType = RecursiveASTVisitor<Derived>;
+
+public:
+  LexicallyOrderedRecursiveASTVisitor(const SourceManager &SM) : SM(SM) {}
+
+  bool TraverseObjCImplementationDecl(ObjCImplementationDecl *D) {
+    // Objective-C @implementation declarations should not trigger early exit
+    // until the additional decls are traversed as their children are not
+    // lexically ordered.
+    bool Result = BaseType::TraverseObjCImplementationDecl(D);
+    return TraverseAdditionalLexicallyNestedDeclarations() ? Result : false;
+  }
+
+  bool TraverseObjCCategoryImplDecl(ObjCCategoryImplDecl *D) {
+    bool Result = BaseType::TraverseObjCCategoryImplDecl(D);
+    return TraverseAdditionalLexicallyNestedDeclarations() ? Result : false;
+  }
+
+  bool TraverseDeclContextHelper(DeclContext *DC) {
+    if (!DC)
+      return true;
+
+    for (auto I = DC->decls_begin(), E = DC->decls_end(); I != E;) {
+      Decl *Child = *I;
+      if (BaseType::canIgnoreChildDeclWhileTraversingDeclContext(Child)) {
+        ++I;
+        continue;
+      }
+      if (!isa<ObjCImplementationDecl>(Child) &&
+          !isa<ObjCCategoryImplDecl>(Child)) {
+        if (!BaseType::getDerived().TraverseDecl(Child))
+          return false;
+        ++I;
+        continue;
+      }
+      // Gather declarations that follow the Objective-C implementation
+      // declarations but are lexically contained in the implementation.
+      LexicallyNestedDeclarations.clear();
+      for (++I; I != E; ++I) {
+        Decl *Sibling = *I;
+        if (!SM.isBeforeInTranslationUnit(Sibling->getLocStart(),
+                                          Child->getLocEnd()))
+          break;
+        if (!BaseType::canIgnoreChildDeclWhileTraversingDeclContext(Sibling))
+          LexicallyNestedDeclarations.push_back(Sibling);
+      }
+      if (!BaseType::getDerived().TraverseDecl(Child))
+        return false;
+    }
+    return true;
+  }
+
+  Stmt::child_range getStmtChildren(Stmt *S) { return S->children(); }
+
+  SmallVector<Stmt *, 8> getStmtChildren(CXXOperatorCallExpr *CE) {
+    SmallVector<Stmt *, 8> Children(CE->children());
+    bool Swap;
+    // Switch the operator and the first operand for all infix and postfix
+    // operations.
+    switch (CE->getOperator()) {
+    case OO_Arrow:
+    case OO_Call:
+    case OO_Subscript:
+      Swap = true;
+      break;
+    case OO_PlusPlus:
+    case OO_MinusMinus:
+      // These are postfix unless there is exactly one argument.
+      Swap = Children.size() != 2;
+      break;
+    default:
+      Swap = CE->isInfixBinaryOp();
+      break;
+    }
+    if (Swap && Children.size() > 1)
+      std::swap(Children[0], Children[1]);
+    return Children;
+  }
+
+private:
+  bool TraverseAdditionalLexicallyNestedDeclarations() {
+    // FIXME: Ideally the gathered declarations and the declarations in the
+    // @implementation should be mixed and sorted to get a true lexical order,
+    // but right now we only care about getting the correct lexical parent, so
+    // we can traverse the gathered nested declarations after the declarations
+    // in the decl context.
+    assert(!BaseType::getDerived().shouldTraversePostOrder() &&
+           "post-order traversal is not supported for lexically ordered "
+           "recursive ast visitor");
+    for (Decl *D : LexicallyNestedDeclarations) {
+      if (!BaseType::getDerived().TraverseDecl(D))
+        return false;
+    }
+    return true;
+  }
+
+  const SourceManager &SM;
+  llvm::SmallVector<Decl *, 8> LexicallyNestedDeclarations;
+};
+
+} // end namespace clang
+
+#endif // LLVM_CLANG_AST_LEXICALLY_ORDERED_RECURSIVEASTVISITOR_H