[Testing] TestAST, a helper for writing straight-line AST tests

Tests that need ASTs have to deal with the awkward control flow of
FrontendAction in some way. There are a few idioms used:
 - don't bother with unit tests, use clang -dump-ast
 - create an ASTConsumer by hand, which is bulky
 - use ASTMatchFinder - works pretty well if matchers are actually
   needed, very strange if they are not
 - use ASTUnit - this yields nice straight-line code, but ASTUnit is a
   terrifically complicated library not designed for this purpose

TestAST provides a very simple way to write straight-line tests: specify
the code/flags and it provides an AST that is kept alive until the
object is destroyed.
It's loosely modeled after TestTU in clangd, which we've successfully
used for a variety of tests.

I've updated a couple of clang tests to use this helper, IMO they're clearer.

Differential Revision: https://reviews.llvm.org/D123668

clang/include/clang/Basic/Diagnostic.h

@@ -39,7 +39,8 @@
 
 namespace llvm {
 class Error;
-}
+class raw_ostream;
+} // namespace llvm
 
 namespace clang {
 
@@ -1717,6 +1718,9 @@ public:
   }
 };
 
+// Simple debug printing of StoredDiagnostic.
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const StoredDiagnostic &);
+
 /// Abstract interface, implemented by clients of the front-end, which
 /// formats and prints fully processed diagnostics.
 class DiagnosticConsumer {

clang/include/clang/Testing/CommandLineArgs.h

@@ -33,6 +33,7 @@ enum TestLanguage {
 };
 
 std::vector<std::string> getCommandLineArgsForTesting(TestLanguage Lang);
+std::vector<std::string> getCC1ArgsForTesting(TestLanguage Lang);
 
 StringRef getFilenameForTesting(TestLanguage Lang);
 

clang/include/clang/Testing/TestAST.h

@@ -0,0 +1,91 @@
+//===--- TestAST.h - Build clang ASTs for testing -------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// In normal operation of Clang, the FrontendAction's lifecycle both creates
+// and destroys the AST, and code should operate on it during callbacks in
+// between (e.g. via ASTConsumer).
+//
+// For tests it is often more convenient to parse an AST from code, and keep it
+// alive as a normal local object, with assertions as straight-line code.
+// TestAST provides such an interface.
+// (ASTUnit can be used for this purpose, but is a production library with
+// broad scope and complicated API).
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TESTING_TESTAST_H
+#define LLVM_CLANG_TESTING_TESTAST_H
+
+#include "clang/Basic/LLVM.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Testing/CommandLineArgs.h"
+#include "llvm/ADT/StringRef.h"
+#include <string>
+#include <vector>
+
+namespace clang {
+
+/// Specifies a virtual source file to be parsed as part of a test.
+struct TestInputs {
+  TestInputs() = default;
+  TestInputs(StringRef Code) : Code(Code) {}
+
+  /// The source code of the input file to be parsed.
+  std::string Code;
+
+  /// The language to parse as.
+  /// This affects the -x and -std flags used, and the filename.
+  TestLanguage Language = TestLanguage::Lang_OBJCXX;
+
+  /// Extra argv to pass to clang -cc1.
+  std::vector<std::string> ExtraArgs = {};
+
+  /// By default, error diagnostics during parsing are reported as gtest errors.
+  /// To suppress this, set ErrorOK or include "error-ok" in a comment in Code.
+  /// In either case, all diagnostics appear in TestAST::diagnostics().
+  bool ErrorOK = false;
+};
+
+/// The result of parsing a file specified by TestInputs.
+///
+/// The ASTContext, Sema etc are valid as long as this object is alive.
+class TestAST {
+public:
+  /// Constructing a TestAST parses the virtual file.
+  ///
+  /// To keep tests terse, critical errors (e.g. invalid flags) are reported as
+  /// unit test failures with ADD_FAILURE() and produce an empty ASTContext,
+  /// Sema etc. This frees the test code from handling these explicitly.
+  TestAST(const TestInputs &);
+  TestAST(StringRef Code) : TestAST(TestInputs(Code)) {}
+  TestAST(TestAST &&M);
+  TestAST &operator=(TestAST &&);
+  ~TestAST();
+
+  /// Provides access to the AST context and other parts of Clang.
+
+  ASTContext &context() { return Clang->getASTContext(); }
+  Sema &sema() { return Clang->getSema(); }
+  SourceManager &sourceManager() { return Clang->getSourceManager(); }
+  FileManager &fileManager() { return Clang->getFileManager(); }
+  Preprocessor &preprocessor() { return Clang->getPreprocessor(); }
+
+  /// Returns diagnostics emitted during parsing.
+  /// (By default, errors cause test failures, see TestInputs::ErrorOK).
+  llvm::ArrayRef<StoredDiagnostic> diagnostics() { return Diagnostics; }
+
+private:
+  void clear();
+  std::unique_ptr<FrontendAction> Action;
+  std::unique_ptr<CompilerInstance> Clang;
+  std::vector<StoredDiagnostic> Diagnostics;
+};
+
+} // end namespace clang
+
+#endif

clang/lib/Basic/Diagnostic.cpp

@@ -1138,6 +1138,14 @@ StoredDiagnostic::StoredDiagnostic(DiagnosticsEngine::Level Level, unsigned ID,
 {
 }
 
+llvm::raw_ostream &clang::operator<<(llvm::raw_ostream &OS,
+                                     const StoredDiagnostic &SD) {
+  if (SD.getLocation().hasManager())
+    OS << SD.getLocation().printToString(SD.getLocation().getManager()) << ": ";
+  OS << SD.getMessage();
+  return OS;
+}
+
 /// IncludeInDiagnosticCounts - This method (whose default implementation
 ///  returns true) indicates whether the diagnostics handled by this
 ///  DiagnosticConsumer should be included in the number of diagnostics

clang/lib/Testing/CMakeLists.txt

@@ -1,14 +1,19 @@
-set(LLVM_LINK_COMPONENTS
-  Support
-  )
-
 # Not add_clang_library: this is not part of clang's public library interface.
 # Unit tests should depend on this with target_link_libraries(), rather
 # than with clang_target_link_libraries().
 add_llvm_library(clangTesting
   CommandLineArgs.cpp
+  TestAST.cpp
+
   BUILDTREE_ONLY
 
   LINK_COMPONENTS
   Support
   )
+
+target_link_libraries(clangTesting
+  PRIVATE
+  llvm_gtest
+  clangBasic
+  clangFrontend
+  )

clang/lib/Testing/CommandLineArgs.cpp

@@ -45,6 +45,39 @@ std::vector<std::string> getCommandLineArgsForTesting(TestLanguage Lang) {
   return Args;
 }
 
+std::vector<std::string> getCC1ArgsForTesting(TestLanguage Lang) {
+  std::vector<std::string> Args;
+  switch (Lang) {
+  case Lang_C89:
+    Args = {"-xc", "-std=c89"};
+    break;
+  case Lang_C99:
+    Args = {"-xc", "-std=c99"};
+    break;
+  case Lang_CXX03:
+    Args = {"-std=c++03"};
+    break;
+  case Lang_CXX11:
+    Args = {"-std=c++11"};
+    break;
+  case Lang_CXX14:
+    Args = {"-std=c++14"};
+    break;
+  case Lang_CXX17:
+    Args = {"-std=c++17"};
+    break;
+  case Lang_CXX20:
+    Args = {"-std=c++20"};
+    break;
+  case Lang_OBJCXX:
+    Args = {"-xobjective-c++"};
+    break;
+  case Lang_OpenCL:
+    llvm_unreachable("Not implemented yet!");
+  }
+  return Args;
+}
+
 StringRef getFilenameForTesting(TestLanguage Lang) {
   switch (Lang) {
   case Lang_C89:

clang/lib/Testing/TestAST.cpp

@@ -0,0 +1,158 @@
+//===--- TestAST.cpp ------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Testing/TestAST.h"
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/LangOptions.h"
+#include "clang/Frontend/FrontendActions.h"
+#include "clang/Frontend/TextDiagnostic.h"
+#include "clang/Testing/CommandLineArgs.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/Support/VirtualFileSystem.h"
+
+#include "gtest/gtest.h"
+
+namespace clang {
+namespace {
+
+// Captures diagnostics into a vector, optionally reporting errors to gtest.
+class StoreDiagnostics : public DiagnosticConsumer {
+  std::vector<StoredDiagnostic> &Out;
+  bool ReportErrors;
+  LangOptions LangOpts;
+
+public:
+  StoreDiagnostics(std::vector<StoredDiagnostic> &Out, bool ReportErrors)
+      : Out(Out), ReportErrors(ReportErrors) {}
+
+  void BeginSourceFile(const LangOptions &LangOpts,
+                       const Preprocessor *) override {
+    this->LangOpts = LangOpts;
+  }
+
+  void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
+                        const Diagnostic &Info) override {
+    Out.emplace_back(DiagLevel, Info);
+    if (ReportErrors && DiagLevel >= DiagnosticsEngine::Error) {
+      std::string Text;
+      llvm::raw_string_ostream OS(Text);
+      TextDiagnostic Renderer(OS, LangOpts,
+                              &Info.getDiags()->getDiagnosticOptions());
+      Renderer.emitStoredDiagnostic(Out.back());
+      ADD_FAILURE() << Text;
+    }
+  }
+};
+
+// Fills in the bits of a CompilerInstance that weren't initialized yet.
+// Provides "empty" ASTContext etc if we fail before parsing gets started.
+void createMissingComponents(CompilerInstance &Clang) {
+  if (!Clang.hasDiagnostics())
+    Clang.createDiagnostics();
+  if (!Clang.hasFileManager())
+    Clang.createFileManager();
+  if (!Clang.hasSourceManager())
+    Clang.createSourceManager(Clang.getFileManager());
+  if (!Clang.hasTarget())
+    Clang.createTarget();
+  if (!Clang.hasPreprocessor())
+    Clang.createPreprocessor(TU_Complete);
+  if (!Clang.hasASTConsumer())
+    Clang.setASTConsumer(std::make_unique<ASTConsumer>());
+  if (!Clang.hasASTContext())
+    Clang.createASTContext();
+  if (!Clang.hasSema())
+    Clang.createSema(TU_Complete, /*CodeCompleteConsumer=*/nullptr);
+}
+
+} // namespace
+
+TestAST::TestAST(const TestInputs &In) {
+  Clang = std::make_unique<CompilerInstance>(
+      std::make_shared<PCHContainerOperations>());
+  // If we don't manage to finish parsing, create CompilerInstance components
+  // anyway so that the test will see an empty AST instead of crashing.
+  auto RecoverFromEarlyExit =
+      llvm::make_scope_exit([&] { createMissingComponents(*Clang); });
+
+  // Extra error conditions are reported through diagnostics, set that up first.
+  bool ErrorOK = In.ErrorOK || llvm::StringRef(In.Code).contains("error-ok");
+  Clang->createDiagnostics(new StoreDiagnostics(Diagnostics, !ErrorOK));
+
+  // Parse cc1 argv, (typically [-std=c++20 input.cc]) into CompilerInvocation.
+  std::vector<const char *> Argv;
+  std::vector<std::string> LangArgs = getCC1ArgsForTesting(In.Language);
+  for (const auto &S : LangArgs)
+    Argv.push_back(S.c_str());
+  for (const auto &S : In.ExtraArgs)
+    Argv.push_back(S.c_str());
+  std::string Filename = getFilenameForTesting(In.Language).str();
+  Argv.push_back(Filename.c_str());
+  Clang->setInvocation(std::make_unique<CompilerInvocation>());
+  if (!CompilerInvocation::CreateFromArgs(Clang->getInvocation(), Argv,
+                                          Clang->getDiagnostics(), "clang")) {
+    ADD_FAILURE() << "Failed to create invocation";
+    return;
+  }
+  assert(!Clang->getInvocation().getFrontendOpts().DisableFree);
+
+  // Set up a VFS with only the virtual file visible.
+  auto VFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+  VFS->addFile(Filename, /*ModificationTime=*/0,
+               llvm::MemoryBuffer::getMemBufferCopy(In.Code, Filename));
+  Clang->createFileManager(VFS);
+
+  // Running the FrontendAction creates the other components: SourceManager,
+  // Preprocessor, ASTContext, Sema. Preprocessor needs TargetInfo to be set.
+  EXPECT_TRUE(Clang->createTarget());
+  Action = std::make_unique<SyntaxOnlyAction>();
+  const FrontendInputFile &Main = Clang->getFrontendOpts().Inputs.front();
+  if (!Action->BeginSourceFile(*Clang, Main)) {
+    ADD_FAILURE() << "Failed to BeginSourceFile()";
+    Action.reset(); // Don't call EndSourceFile if BeginSourceFile failed.
+    return;
+  }
+  if (auto Err = Action->Execute())
+    ADD_FAILURE() << "Failed to Execute(): " << llvm::toString(std::move(Err));
+
+  // Action->EndSourceFile() would destroy the ASTContext, we want to keep it.
+  // But notify the preprocessor we're done now.
+  Clang->getPreprocessor().EndSourceFile();
+  // We're done gathering diagnostics, detach the consumer so we can destroy it.
+  Clang->getDiagnosticClient().EndSourceFile();
+  Clang->getDiagnostics().setClient(new DiagnosticConsumer(),
+                                    /*ShouldOwnClient=*/true);
+}
+
+void TestAST::clear() {
+  if (Action) {
+    // We notified the preprocessor of EOF already, so detach it first.
+    // Sema needs the PP alive until after EndSourceFile() though.
+    auto PP = Clang->getPreprocessorPtr(); // Keep PP alive for now.
+    Clang->setPreprocessor(nullptr);       // Detach so we don't send EOF twice.
+    Action->EndSourceFile();               // Destroy ASTContext and Sema.
+    // Now Sema is gone, PP can safely be destroyed.
+  }
+  Action.reset();
+  Clang.reset();
+  Diagnostics.clear();
+}
+
+TestAST &TestAST::operator=(TestAST &&M) {
+  clear();
+  Action = std::move(M.Action);
+  Clang = std::move(M.Clang);
+  Diagnostics = std::move(M.Diagnostics);
+  return *this;
+}
+
+TestAST::TestAST(TestAST &&M) { *this = std::move(M); }
+
+TestAST::~TestAST() { clear(); }
+
+} // end namespace clang

clang/unittests/Tooling/CMakeLists.txt

@@ -88,6 +88,7 @@ clang_target_link_libraries(ToolingTests
 target_link_libraries(ToolingTests
   PRIVATE
   LLVMTestingSupport
+  clangTesting
 )
 
 add_subdirectory(Syntax)

clang/unittests/Tooling/FixItTest.cpp

@@ -6,9 +6,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "TestVisitor.h"
-#include "clang/Basic/Diagnostic.h"
 #include "clang/Tooling/FixIt.h"
+#include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Testing/TestAST.h"
+#include "gtest/gtest.h"
 
 using namespace clang;
 
@@ -18,214 +20,169 @@ using tooling::fixit::createReplacement;
 
 namespace {
 
-struct CallsVisitor : TestVisitor<CallsVisitor> {
-  bool VisitCallExpr(CallExpr *Expr) {
-    OnCall(Expr, Context);
-    return true;
-  }
-
-  std::function<void(CallExpr *, ASTContext *Context)> OnCall;
-};
-
-std::string LocationToString(SourceLocation Loc, ASTContext *Context) {
-  return Loc.printToString(Context->getSourceManager());
+const CallExpr &onlyCall(ASTContext &Ctx) {
+  using namespace ast_matchers;
+  auto Calls = match(callExpr().bind(""), Ctx);
+  EXPECT_EQ(Calls.size(), 1u);
+  return *Calls.front().getNodeAs<CallExpr>("");
 }
 
 TEST(FixItTest, getText) {
-  CallsVisitor Visitor;
+  TestAST AST("void foo(int x, int y) { foo(x, y); }");
+  const CallExpr &CE = onlyCall(AST.context());
+  EXPECT_EQ("foo(x, y)", getText(CE, AST.context()));
+  EXPECT_EQ("foo(x, y)", getText(CE.getSourceRange(), AST.context()));
+  EXPECT_EQ("x", getText(*CE.getArg(0), AST.context()));
+  EXPECT_EQ("y", getText(*CE.getArg(1), AST.context()));
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    EXPECT_EQ("foo(x, y)", getText(*CE, *Context));
-    EXPECT_EQ("foo(x, y)", getText(CE->getSourceRange(), *Context));
-
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    EXPECT_EQ("x", getText(*P0, *Context));
-    EXPECT_EQ("y", getText(*P1, *Context));
-  };
-  Visitor.runOver("void foo(int x, int y) { foo(x, y); }");
-
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    EXPECT_EQ("APPLY(foo, x, y)", getText(*CE, *Context));
-  };
-  Visitor.runOver("#define APPLY(f, x, y) f(x, y)\n"
-                  "void foo(int x, int y) { APPLY(foo, x, y); }");
+  AST = TestAST("#define APPLY(f, x, y) f(x, y)\n"
+                "void foo(int x, int y) { APPLY(foo, x, y); }");
+  const CallExpr &CE2 = onlyCall(AST.context());
+  EXPECT_EQ("APPLY(foo, x, y)", getText(CE2, AST.context()));
 }
 
 TEST(FixItTest, getTextWithMacro) {
-  CallsVisitor Visitor;
+  TestAST AST("#define F foo(\n"
+              "#define OO x, y)\n"
+              "void foo(int x, int y) { F OO ; }");
+  const CallExpr &CE = onlyCall(AST.context());
+  EXPECT_EQ("F OO", getText(CE, AST.context()));
+  EXPECT_EQ("", getText(*CE.getArg(0), AST.context()));
+  EXPECT_EQ("", getText(*CE.getArg(1), AST.context()));
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    EXPECT_EQ("F OO", getText(*CE, *Context));
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    EXPECT_EQ("", getText(*P0, *Context));
-    EXPECT_EQ("", getText(*P1, *Context));
-  };
-  Visitor.runOver("#define F foo(\n"
-                  "#define OO x, y)\n"
-                  "void foo(int x, int y) { F OO ; }");
-
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    EXPECT_EQ("", getText(*CE, *Context));
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    EXPECT_EQ("x", getText(*P0, *Context));
-    EXPECT_EQ("y", getText(*P1, *Context));
-  };
-  Visitor.runOver("#define FOO(x, y) (void)x; (void)y; foo(x, y);\n"
-                  "void foo(int x, int y) { FOO(x,y) }");
+  AST = TestAST("#define FOO(x, y) (void)x; (void)y; foo(x, y);\n"
+                "void foo(int x, int y) { FOO(x,y) }");
+  const CallExpr &CE2 = onlyCall(AST.context());
+  EXPECT_EQ("", getText(CE2, AST.context()));
+  EXPECT_EQ("x", getText(*CE2.getArg(0), AST.context()));
+  EXPECT_EQ("y", getText(*CE2.getArg(1), AST.context()));
 }
 
 TEST(FixItTest, createRemoval) {
-  CallsVisitor Visitor;
+  TestAST AST("void foo(int x, int y) { foo(x, y); }");
+  const CallExpr &CE = onlyCall(AST.context());
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    FixItHint Hint = createRemoval(*CE);
-    EXPECT_EQ("foo(x, y)", getText(Hint.RemoveRange.getAsRange(), *Context));
-    EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint.CodeToInsert.empty());
+  FixItHint Hint = createRemoval(CE);
+  EXPECT_EQ("foo(x, y)", getText(Hint.RemoveRange.getAsRange(), AST.context()));
+  EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint.CodeToInsert.empty());
 
-    Expr *P0 = CE->getArg(0);
-    FixItHint Hint0 = createRemoval(*P0);
-    EXPECT_EQ("x", getText(Hint0.RemoveRange.getAsRange(), *Context));
-    EXPECT_TRUE(Hint0.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint0.CodeToInsert.empty());
+  FixItHint Hint0 = createRemoval(*CE.getArg(0));
+  EXPECT_EQ("x", getText(Hint0.RemoveRange.getAsRange(), AST.context()));
+  EXPECT_TRUE(Hint0.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint0.CodeToInsert.empty());
 
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint1 = createRemoval(*P1);
-    EXPECT_EQ("y", getText(Hint1.RemoveRange.getAsRange(), *Context));
-    EXPECT_TRUE(Hint1.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint1.CodeToInsert.empty());
-  };
-  Visitor.runOver("void foo(int x, int y) { foo(x, y); }");
+  FixItHint Hint1 = createRemoval(*CE.getArg(1));
+  EXPECT_EQ("y", getText(Hint1.RemoveRange.getAsRange(), AST.context()));
+  EXPECT_TRUE(Hint1.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint1.CodeToInsert.empty());
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    Expr *P0 = CE->getArg(0);
-    FixItHint Hint0 = createRemoval(*P0);
-    EXPECT_EQ("x + y", getText(Hint0.RemoveRange.getAsRange(), *Context));
+  AST = TestAST("void foo(int x, int y) { foo(x + y, y + x); }");
+  const CallExpr &CE2 = onlyCall(AST.context());
+  Hint0 = createRemoval(*CE2.getArg(0));
+  EXPECT_EQ("x + y", getText(Hint0.RemoveRange.getAsRange(), AST.context()));
 
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint1 = createRemoval(*P1);
-    EXPECT_EQ("y + x", getText(Hint1.RemoveRange.getAsRange(), *Context));
-  };
-  Visitor.runOver("void foo(int x, int y) { foo(x + y, y + x); }");
+  Hint1 = createRemoval(*CE2.getArg(1));
+  EXPECT_EQ("y + x", getText(Hint1.RemoveRange.getAsRange(), AST.context()));
 }
 
 TEST(FixItTest, createRemovalWithMacro) {
-  CallsVisitor Visitor;
+  TestAST AST("#define FOO foo(1, 1)\n"
+              "void foo(int x, int y) { FOO; }");
+  const CallExpr &CE = onlyCall(AST.context());
+  FixItHint Hint = createRemoval(CE);
+  EXPECT_EQ("FOO", getText(Hint.RemoveRange.getAsRange(), AST.context()));
+  EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint.CodeToInsert.empty());
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    FixItHint Hint = createRemoval(*CE);
-    EXPECT_EQ("FOO", getText(Hint.RemoveRange.getAsRange(), *Context));
-    EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint.CodeToInsert.empty());
+  FixItHint Hint0 = createRemoval(*CE.getArg(0));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:17>",
+            Hint0.RemoveRange.getBegin().printToString(AST.sourceManager()));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:17>",
+            Hint0.RemoveRange.getEnd().printToString(AST.sourceManager()));
+  EXPECT_TRUE(Hint0.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint0.CodeToInsert.empty());
 
-    Expr *P0 = CE->getArg(0);
-    FixItHint Hint0 = createRemoval(*P0);
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:17>",
-              LocationToString(Hint0.RemoveRange.getBegin(), Context));
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:17>",
-              LocationToString(Hint0.RemoveRange.getEnd(), Context));
-    EXPECT_TRUE(Hint0.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint0.CodeToInsert.empty());
+  FixItHint Hint1 = createRemoval(*CE.getArg(1));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:20>",
+            Hint1.RemoveRange.getBegin().printToString(AST.sourceManager()));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:20>",
+            Hint1.RemoveRange.getEnd().printToString(AST.sourceManager()));
+  EXPECT_TRUE(Hint1.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint1.CodeToInsert.empty());
 
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint1 = createRemoval(*P1);
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:20>",
-              LocationToString(Hint1.RemoveRange.getBegin(), Context));
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:20>",
-              LocationToString(Hint1.RemoveRange.getEnd(), Context));
-    EXPECT_TRUE(Hint1.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint1.CodeToInsert.empty());
-  };
-  Visitor.runOver("#define FOO foo(1, 1)\n"
-                  "void foo(int x, int y) { FOO; }");
-
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    FixItHint Hint = createRemoval(*CE);
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:37>",
-              LocationToString(Hint.RemoveRange.getBegin(), Context));
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:45>",
-              LocationToString(Hint.RemoveRange.getEnd(), Context));
-    EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint.CodeToInsert.empty());
-  };
-  Visitor.runOver("#define FOO(x, y) (void)x; (void)y; foo(x, y);\n"
-                  "void foo(int x, int y) { FOO(x,y) }");
+  AST = TestAST("#define FOO(x, y) (void)x; (void)y; foo(x, y);\n"
+                "void foo(int x, int y) { FOO(x,y) }");
+  const CallExpr &CE2 = onlyCall(AST.context());
+  Hint = createRemoval(CE2);
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:37>",
+            Hint.RemoveRange.getBegin().printToString(AST.sourceManager()));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:45>",
+            Hint.RemoveRange.getEnd().printToString(AST.sourceManager()));
+  EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint.CodeToInsert.empty());
 }
 
 TEST(FixItTest, createReplacement) {
-  CallsVisitor Visitor;
+  for (const char *Code : {
+           "void foo(int x, int y) { foo(x, y); }",
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint0 = createReplacement(*P0, *P1, *Context);
-    FixItHint Hint1 = createReplacement(*P1, *P0, *Context);
+           "#define APPLY(f, x, y) f(x, y)\n"
+           "void foo(int x, int y) { APPLY(foo, x, y); }",
+
+           "#define APPLY(f, P) f(P)\n"
+           "#define PAIR(x, y) x, y\n"
+           "void foo(int x, int y) { APPLY(foo, PAIR(x, y)); }\n",
+       }) {
+    TestAST AST(Code);
+    const CallExpr &CE = onlyCall(AST.context());
+    const Expr *P0 = CE.getArg(0);
+    const Expr *P1 = CE.getArg(1);
+    FixItHint Hint0 = createReplacement(*P0, *P1, AST.context());
+    FixItHint Hint1 = createReplacement(*P1, *P0, AST.context());
 
     // Validate Hint0 fields.
-    EXPECT_EQ("x", getText(Hint0.RemoveRange.getAsRange(), *Context));
+    EXPECT_EQ("x", getText(Hint0.RemoveRange.getAsRange(), AST.context()));
     EXPECT_TRUE(Hint0.InsertFromRange.isInvalid());
     EXPECT_EQ(Hint0.CodeToInsert, "y");
 
     // Validate Hint1 fields.
-    EXPECT_EQ("y", getText(Hint1.RemoveRange.getAsRange(), *Context));
+    EXPECT_EQ("y", getText(Hint1.RemoveRange.getAsRange(), AST.context()));
     EXPECT_TRUE(Hint1.InsertFromRange.isInvalid());
     EXPECT_EQ(Hint1.CodeToInsert, "x");
-  };
-
-  Visitor.runOver("void foo(int x, int y) { foo(x, y); }");
-
-  Visitor.runOver("#define APPLY(f, x, y) f(x, y)\n"
-                  "void foo(int x, int y) { APPLY(foo, x, y); }");
-
-  Visitor.runOver("#define APPLY(f, P) f(P)\n"
-                  "#define PAIR(x, y) x, y\n"
-                  "void foo(int x, int y) { APPLY(foo, PAIR(x, y)); }\n");
+  }
 }
 
 TEST(FixItTest, createReplacementWithMacro) {
-  CallsVisitor Visitor;
+  TestAST AST("#define FOO foo(1, 1)\n"
+              "void foo(int x, int y) { FOO; }");
+  const CallExpr &CE = onlyCall(AST.context());
+  FixItHint Hint =
+      createReplacement(*CE.getArg(0), *CE.getArg(1), AST.context());
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:17>",
+            Hint.RemoveRange.getBegin().printToString(AST.sourceManager()));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:1:17>",
+            Hint.RemoveRange.getEnd().printToString(AST.sourceManager()));
+  EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
+  EXPECT_TRUE(Hint.CodeToInsert.empty());
 
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint = createReplacement(*P0, *P1, *Context);
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:17>",
-              LocationToString(Hint.RemoveRange.getBegin(), Context));
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:1:17>",
-              LocationToString(Hint.RemoveRange.getEnd(), Context));
-    EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
-    EXPECT_TRUE(Hint.CodeToInsert.empty());
-  };
+  AST = TestAST("#define FOO(x, y) (void)x; (void)y; foo(x, y);\n"
+                "void foo(int x, int y) { FOO(x,y) }");
+  const CallExpr &CE2 = onlyCall(AST.context());
+  Hint = createReplacement(*CE2.getArg(0), *CE2.getArg(1), AST.context());
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:2:30>",
+            Hint.RemoveRange.getEnd().printToString(AST.sourceManager()));
+  EXPECT_EQ("input.mm:2:26 <Spelling=input.mm:2:30>",
+            Hint.RemoveRange.getBegin().printToString(AST.sourceManager()));
+  EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
+  EXPECT_EQ("y", Hint.CodeToInsert);
 
-  Visitor.runOver("#define FOO foo(1, 1)\n"
-                  "void foo(int x, int y) { FOO; }");
-
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint = createReplacement(*P0, *P1, *Context);
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:2:30>",
-              LocationToString(Hint.RemoveRange.getBegin(), Context));
-    EXPECT_EQ("input.cc:2:26 <Spelling=input.cc:2:30>",
-              LocationToString(Hint.RemoveRange.getEnd(), Context));
-    EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
-    EXPECT_EQ("y", Hint.CodeToInsert);
-  };
-  Visitor.runOver("#define FOO(x, y) (void)x; (void)y; foo(x, y);\n"
-                  "void foo(int x, int y) { FOO(x,y) }");
-
-  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
-    Expr *P0 = CE->getArg(0);
-    Expr *P1 = CE->getArg(1);
-    FixItHint Hint = createReplacement(*P0, *P1, *Context);
-    EXPECT_EQ("x + y", getText(Hint.RemoveRange.getAsRange(), *Context));
-    EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
-    EXPECT_EQ("y + x", Hint.CodeToInsert);
-  };
-  Visitor.runOver("void foo(int x, int y) { foo(x + y, y + x); }");
+  AST = TestAST("void foo(int x, int y) { foo(x + y, y + x); }");
+  const CallExpr &CE3 = onlyCall(AST.context());
+  Hint = createReplacement(*CE3.getArg(0), *CE3.getArg(1), AST.context());
+  EXPECT_EQ("x + y", getText(Hint.RemoveRange.getAsRange(), AST.context()));
+  EXPECT_TRUE(Hint.InsertFromRange.isInvalid());
+  EXPECT_EQ("y + x", Hint.CodeToInsert);
 }
 
 } // end anonymous namespace

clang/unittests/Tooling/StandardLibraryTest.cpp

@@ -7,10 +7,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
+#include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclarationName.h"
-#include "clang/Frontend/ASTUnit.h"
-#include "clang/Tooling/Tooling.h"
+#include "clang/Testing/TestAST.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ScopedPrinter.h"
@@ -24,10 +24,9 @@ namespace clang {
 namespace tooling {
 namespace {
 
-const NamedDecl &lookup(ASTUnit &AST, llvm::StringRef Name) {
-  auto &Ctx = AST.getASTContext();
-  TranslationUnitDecl *TU = Ctx.getTranslationUnitDecl();
-  auto Result = TU->lookup(DeclarationName(&Ctx.Idents.get(Name)));
+const NamedDecl &lookup(TestAST &AST, llvm::StringRef Name) {
+  TranslationUnitDecl *TU = AST.context().getTranslationUnitDecl();
+  auto Result = TU->lookup(DeclarationName(&AST.context().Idents.get(Name)));
   assert(!Result.empty() && "Lookup failed");
   assert(Result.isSingleResult() && "Lookup returned multiple results");
   return *Result.front();
@@ -50,7 +49,7 @@ TEST(StdlibTest, All) {
 }
 
 TEST(StdlibTest, Recognizer) {
-  std::unique_ptr<ASTUnit> AST = buildASTFromCode(R"cpp(
+  TestAST AST(R"cpp(
     namespace std {
     inline namespace inl {
 
@@ -83,17 +82,15 @@ TEST(StdlibTest, Recognizer) {
     div_t div;
   )cpp");
 
-  auto &VectorNonstd = lookup(*AST, "vector");
-  auto *Vec =
-      cast<VarDecl>(lookup(*AST, "vec")).getType()->getAsCXXRecordDecl();
+  auto &VectorNonstd = lookup(AST, "vector");
+  auto *Vec = cast<VarDecl>(lookup(AST, "vec")).getType()->getAsCXXRecordDecl();
   auto *Nest =
-      cast<VarDecl>(lookup(*AST, "nest")).getType()->getAsCXXRecordDecl();
+      cast<VarDecl>(lookup(AST, "nest")).getType()->getAsCXXRecordDecl();
   auto *Clock =
-      cast<VarDecl>(lookup(*AST, "clock")).getType()->getAsCXXRecordDecl();
-  auto *Sec =
-      cast<VarDecl>(lookup(*AST, "sec")).getType()->getAsCXXRecordDecl();
+      cast<VarDecl>(lookup(AST, "clock")).getType()->getAsCXXRecordDecl();
+  auto *Sec = cast<VarDecl>(lookup(AST, "sec")).getType()->getAsCXXRecordDecl();
   auto *CDivT =
-      cast<VarDecl>(lookup(*AST, "div")).getType()->getAsCXXRecordDecl();
+      cast<VarDecl>(lookup(AST, "div")).getType()->getAsCXXRecordDecl();
 
   stdlib::Recognizer Recognizer;