1 files changed, 227 insertions, 0 deletions
diff --git a/unittests/Tooling/LexicallyOrderedRecursiveASTVisitorTest.cpp b/unittests/Tooling/LexicallyOrderedRecursiveASTVisitorTest.cpp
new file mode 100644
index 0000000000000..50727a55fc714
--- /dev/null
+++ b/unittests/Tooling/LexicallyOrderedRecursiveASTVisitorTest.cpp
@@ -0,0 +1,227 @@
+//===- unittest/Tooling/LexicallyOrderedRecursiveASTVisitorTest.cpp -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestVisitor.h"
+#include "clang/AST/LexicallyOrderedRecursiveASTVisitor.h"
+#include <stack>
+
+using namespace clang;
+
+namespace {
+
+class DummyMatchVisitor;
+
+class LexicallyOrderedDeclVisitor
+    : public LexicallyOrderedRecursiveASTVisitor<LexicallyOrderedDeclVisitor> {
+public:
+  LexicallyOrderedDeclVisitor(DummyMatchVisitor &Matcher,
+                              const SourceManager &SM, bool EmitDeclIndices,
+                              bool EmitStmtIndices)
+      : LexicallyOrderedRecursiveASTVisitor(SM), Matcher(Matcher),
+        EmitDeclIndices(EmitDeclIndices), EmitStmtIndices(EmitStmtIndices) {}
+
+  bool TraverseDecl(Decl *D) {
+    TraversalStack.push_back(D);
+    LexicallyOrderedRecursiveASTVisitor::TraverseDecl(D);
+    TraversalStack.pop_back();
+    return true;
+  }
+
+  bool TraverseStmt(Stmt *S);
+
+  bool VisitNamedDecl(const NamedDecl *D);
+  bool VisitDeclRefExpr(const DeclRefExpr *D);
+
+private:
+  DummyMatchVisitor &Matcher;
+  bool EmitDeclIndices, EmitStmtIndices;
+  unsigned Index = 0;
+  llvm::SmallVector<Decl *, 8> TraversalStack;
+};
+
+class DummyMatchVisitor : public ExpectedLocationVisitor<DummyMatchVisitor> {
+  bool EmitDeclIndices, EmitStmtIndices;
+
+public:
+  DummyMatchVisitor(bool EmitDeclIndices = false, bool EmitStmtIndices = false)
+      : EmitDeclIndices(EmitDeclIndices), EmitStmtIndices(EmitStmtIndices) {}
+  bool VisitTranslationUnitDecl(TranslationUnitDecl *TU) {
+    const ASTContext &Context = TU->getASTContext();
+    const SourceManager &SM = Context.getSourceManager();
+    LexicallyOrderedDeclVisitor SubVisitor(*this, SM, EmitDeclIndices,
+                                           EmitStmtIndices);
+    SubVisitor.TraverseDecl(TU);
+    return false;
+  }
+
+  template <class T> void match(StringRef Path, const T *D) {
+    Match(Path, D->getLocStart());
+  }
+};
+
+bool LexicallyOrderedDeclVisitor::TraverseStmt(Stmt *S) {
+  Matcher.match("overridden TraverseStmt", S);
+  return LexicallyOrderedRecursiveASTVisitor::TraverseStmt(S);
+}
+
+bool LexicallyOrderedDeclVisitor::VisitNamedDecl(const NamedDecl *D) {
+  std::string Path;
+  llvm::raw_string_ostream OS(Path);
+  assert(TraversalStack.back() == D);
+  for (const Decl *D : TraversalStack) {
+    if (isa<TranslationUnitDecl>(D)) {
+      OS << "/";
+      continue;
+    }
+    if (const auto *ND = dyn_cast<NamedDecl>(D))
+      OS << ND->getNameAsString();
+    else
+      OS << "???";
+    if (isa<DeclContext>(D) || isa<TemplateDecl>(D))
+      OS << "/";
+  }
+  if (EmitDeclIndices)
+    OS << "@" << Index++;
+  Matcher.match(OS.str(), D);
+  return true;
+}
+
+bool LexicallyOrderedDeclVisitor::VisitDeclRefExpr(const DeclRefExpr *D) {
+  std::string Name = D->getFoundDecl()->getNameAsString();
+  llvm::raw_string_ostream OS(Name);
+  if (EmitStmtIndices)
+    OS << "@" << Index++;
+  Matcher.match(OS.str(), D);
+  return true;
+}
+
+TEST(LexicallyOrderedRecursiveASTVisitor, VisitDeclsInImplementation) {
+  StringRef Source = R"(
+@interface I
+@end
+@implementation I
+
+int nestedFunction() { }
+
+- (void) method{ }
+
+int anotherNestedFunction(int x) {
+  return x;
+}
+
+int innerVariable = 0;
+
+@end
+
+int outerVariable = 0;
+
+@implementation I(Cat)
+
+void catF() { }
+
+@end
+
+void outerFunction() { }
+)";
+  DummyMatchVisitor Visitor;
+  Visitor.DisallowMatch("/nestedFunction/", 6, 1);
+  Visitor.ExpectMatch("/I/nestedFunction/", 6, 1);
+  Visitor.ExpectMatch("/I/method/", 8, 1);
+  Visitor.DisallowMatch("/anotherNestedFunction/", 10, 1);
+  Visitor.ExpectMatch("/I/anotherNestedFunction/", 10, 1);
+  Visitor.DisallowMatch("/innerVariable", 14, 1);
+  Visitor.ExpectMatch("/I/innerVariable", 14, 1);
+  Visitor.ExpectMatch("/outerVariable", 18, 1);
+  Visitor.DisallowMatch("/catF/", 22, 1);
+  Visitor.ExpectMatch("/Cat/catF/", 22, 1);
+  Visitor.ExpectMatch("/outerFunction/", 26, 1);
+  EXPECT_TRUE(Visitor.runOver(Source, DummyMatchVisitor::Lang_OBJC));
+}
+
+TEST(LexicallyOrderedRecursiveASTVisitor, VisitMacroDeclsInImplementation) {
+  StringRef Source = R"(
+@interface I
+@end
+
+void outerFunction() { }
+
+#define MACRO_F(x) void nestedFunction##x() { }
+
+@implementation I
+
+MACRO_F(1)
+
+@end
+
+MACRO_F(2)
+)";
+  DummyMatchVisitor Visitor;
+  Visitor.ExpectMatch("/outerFunction/", 5, 1);
+  Visitor.ExpectMatch("/I/nestedFunction1/", 7, 20);
+  Visitor.ExpectMatch("/nestedFunction2/", 7, 20);
+  EXPECT_TRUE(Visitor.runOver(Source, DummyMatchVisitor::Lang_OBJC));
+}
+
+TEST(LexicallyOrderedRecursiveASTVisitor, VisitTemplateDecl) {
+  StringRef Source = R"(
+template <class T> T f();
+template <class U, class = void> class Class {};
+)";
+  DummyMatchVisitor Visitor(/*EmitIndices=*/true);
+  Visitor.ExpectMatch("/f/T@1", 2, 11);
+  Visitor.ExpectMatch("/f/f/@2", 2, 20);
+  Visitor.ExpectMatch("/Class/U@4", 3, 11);
+  Visitor.ExpectMatch("/Class/@5", 3, 20);
+  Visitor.ExpectMatch("/Class/Class/@6", 3, 34);
+  EXPECT_TRUE(Visitor.runOver(Source));
+}
+
+TEST(LexicallyOrderedRecursiveASTVisitor, VisitCXXOperatorCallExpr) {
+  StringRef Source = R"(
+struct S {
+  S &operator+(S&);
+  S *operator->();
+  S &operator++();
+  S operator++(int);
+  void operator()(int, int);
+  void operator[](int);
+  void f();
+};
+S a, b, c;
+
+void test() {
+  a = b + c;
+  a->f();
+  a(1, 2);
+  b[0];
+  ++a;
+  b++;
+}
+)";
+  DummyMatchVisitor Visitor(/*EmitDeclIndices=*/false,
+                            /*EmitStmtIndices=*/true);
+  // There are two overloaded operators that start at this point
+  // This makes sure they are both traversed using the overridden
+  // TraverseStmt, as the traversal is implemented by us for
+  // CXXOperatorCallExpr.
+  Visitor.ExpectMatch("overridden TraverseStmt", 14, 3, 2);
+  Visitor.ExpectMatch("a@0", 14, 3);
+  Visitor.ExpectMatch("operator=@1", 14, 5);
+  Visitor.ExpectMatch("b@2", 14, 7);
+  Visitor.ExpectMatch("operator+@3", 14, 9);
+  Visitor.ExpectMatch("c@4", 14, 11);
+  Visitor.ExpectMatch("operator->@6", 15, 4);
+  Visitor.ExpectMatch("operator()@8", 16, 4);
+  Visitor.ExpectMatch("operator[]@10", 17, 4);
+  Visitor.ExpectMatch("operator++@11", 18, 3);
+  Visitor.ExpectMatch("operator++@14", 19, 4);
+  EXPECT_TRUE(Visitor.runOver(Source));
+}
+
+} // end anonymous namespace