From ab58c9ee8a6e9ace3a93198496b4d85e8cb2b5a9 Mon Sep 17 00:00:00 2001
From: Eduardo Caldas <ecaldas@google.com>
Date: Fri, 14 Aug 2020 09:43:20 +0000
Subject: [PATCH] [SyntaxTree] Implement annotation-based test infrastructure

We add the method `SyntaxTreeTest::treeDumpEqualOnAnnotations`, which
allows us to compare the treeDump of only annotated code. This will reduce a
lot of noise from our `BuildTreeTest` and make them short and easier to
read.
---
 .../unittests/Tooling/Syntax/TreeTestBase.cpp | 29 +++++++++++++++++++
 clang/unittests/Tooling/Syntax/TreeTestBase.h |  3 ++
 2 files changed, 32 insertions(+)

diff --git a/clang/unittests/Tooling/Syntax/TreeTestBase.cpp b/clang/unittests/Tooling/Syntax/TreeTestBase.cpp
index 6d2efeaaa8eb..05fbac4f47e1 100644
--- a/clang/unittests/Tooling/Syntax/TreeTestBase.cpp
+++ b/clang/unittests/Tooling/Syntax/TreeTestBase.cpp
@@ -180,6 +180,35 @@ SyntaxTreeTest::buildTree(StringRef Code, const TestClangConfig &ClangConfig) {
   return ::testing::AssertionSuccess();
 }
 
+::testing::AssertionResult
+SyntaxTreeTest::treeDumpEqualOnAnnotations(StringRef CodeWithAnnotations,
+                                           ArrayRef<StringRef> TreeDumps) {
+  SCOPED_TRACE(llvm::join(GetParam().getCommandLineArgs(), " "));
+
+  auto AnnotatedCode = llvm::Annotations(CodeWithAnnotations);
+  auto *Root = buildTree(AnnotatedCode.code(), GetParam());
+
+  if (Diags->getClient()->getNumErrors() != 0) {
+    return ::testing::AssertionFailure()
+           << "Source file has syntax errors, they were printed to the test "
+              "log";
+  }
+
+  bool failed = false;
+  auto AnnotatedRanges = AnnotatedCode.ranges();
+  assert(AnnotatedRanges.size() == TreeDumps.size());
+  for (auto i = 0ul; i < AnnotatedRanges.size(); i++) {
+    auto *AnnotatedNode = nodeByRange(AnnotatedRanges[i], Root);
+    assert(AnnotatedNode);
+    auto AnnotatedNodeDump =
+        std::string(StringRef(AnnotatedNode->dump(*Arena)).trim());
+    // EXPECT_EQ shows the diff between the two strings if they are different.
+    EXPECT_EQ(TreeDumps[i].trim().str(), AnnotatedNodeDump);
+    if (AnnotatedNodeDump != TreeDumps[i].trim().str())
+      failed = true;
+  }
+  return failed ? ::testing::AssertionFailure() : ::testing::AssertionSuccess();
+}
 syntax::Node *SyntaxTreeTest::nodeByRange(llvm::Annotations::Range R,
                                           syntax::Node *Root) {
   ArrayRef<syntax::Token> Toks = tokens(Root);
diff --git a/clang/unittests/Tooling/Syntax/TreeTestBase.h b/clang/unittests/Tooling/Syntax/TreeTestBase.h
index bfa6ecd7909f..c282bbf45fd3 100644
--- a/clang/unittests/Tooling/Syntax/TreeTestBase.h
+++ b/clang/unittests/Tooling/Syntax/TreeTestBase.h
@@ -34,6 +34,9 @@ protected:
 
   ::testing::AssertionResult treeDumpEqual(StringRef Code, StringRef Tree);
 
+  ::testing::AssertionResult
+  treeDumpEqualOnAnnotations(StringRef CodeWithAnnotations,
+                             ArrayRef<StringRef> TreeDumps);
   /// Finds the deepest node in the tree that covers exactly \p R.
   /// FIXME: implement this efficiently and move to public syntax tree API.
   syntax::Node *nodeByRange(llvm::Annotations::Range R, syntax::Node *Root);