[clang-tidy] Extend TransformerClangTidyCheck to support adding includes.

Summary:
This revision implements support for the `AddedIncludes` field in
RewriteRule cases; that is, it supports specifying the addition of include
directives in files modified by the clang tidy check.

Reviewers: ilya-biryukov, gribozavr

Subscribers: xazax.hun, cfe-commits

Tags: #clang

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

llvm-svn: 364922

clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.cpp

@@ -46,6 +46,19 @@ TransformerClangTidyCheck::TransformerClangTidyCheck(RewriteRule R,
          " explicitly provide an empty explanation if none is desired");
 }
 
+void TransformerClangTidyCheck::registerPPCallbacks(
+    const SourceManager &SM, Preprocessor *PP, Preprocessor *ModuleExpanderPP) {
+  // Only allocate and register the IncludeInsert when some `Case` will add
+  // includes.
+  if (Rule && llvm::any_of(Rule->Cases, [](const RewriteRule::Case &C) {
+        return !C.AddedIncludes.empty();
+      })) {
+    Inserter = llvm::make_unique<IncludeInserter>(
+        SM, getLangOpts(), utils::IncludeSorter::IS_LLVM);
+    PP->addPPCallbacks(Inserter->CreatePPCallbacks());
+  }
+}
+
 void TransformerClangTidyCheck::registerMatchers(
     ast_matchers::MatchFinder *Finder) {
   if (Rule)
@@ -89,6 +102,15 @@ void TransformerClangTidyCheck::check(
   for (const auto &T : *Transformations) {
     Diag << FixItHint::CreateReplacement(T.Range, T.Replacement);
   }
+
+  for (const auto &I : Case.AddedIncludes) {
+    auto &Header = I.first;
+    if (Optional<FixItHint> Fix = Inserter->CreateIncludeInsertion(
+            Result.SourceManager->getMainFileID(), Header,
+            /*IsAngled=*/I.second == tooling::IncludeFormat::Angled)) {
+      Diag << *Fix;
+    }
+  }
 }
 
 } // namespace utils

clang-tools-extra/clang-tidy/utils/TransformerClangTidyCheck.h

@@ -10,7 +10,9 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANG_TIDY_TRANSFORMER_CLANG_TIDY_CHECK_H
 
 #include "../ClangTidy.h"
+#include "../utils/IncludeInserter.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/Frontend/CompilerInstance.h"
 #include "clang/Tooling/Refactoring/Transformer.h"
 #include <deque>
 #include <vector>
@@ -52,11 +54,14 @@ public:
   TransformerClangTidyCheck(tooling::RewriteRule R, StringRef Name,
                             ClangTidyContext *Context);
 
+  void registerPPCallbacks(const SourceManager &SM, Preprocessor *PP,
+                           Preprocessor *ModuleExpanderPP) override;
   void registerMatchers(ast_matchers::MatchFinder *Finder) final;
   void check(const ast_matchers::MatchFinder::MatchResult &Result) final;
 
 private:
   Optional<tooling::RewriteRule> Rule;
+  std::unique_ptr<clang::tidy::utils::IncludeInserter> Inserter;
 };
 
 } // namespace utils

clang-tools-extra/unittests/clang-tidy/TransformerClangTidyCheckTest.cpp

@@ -19,6 +19,7 @@ namespace tidy {
 namespace utils {
 namespace {
 using tooling::change;
+using tooling::IncludeFormat;
 using tooling::RewriteRule;
 using tooling::text;
 using tooling::stencil::cat;
@@ -121,6 +122,54 @@ TEST(TransformerClangTidyCheckTest, DisableByConfig) {
                           Input, nullptr, "input.cc", None, Options));
 }
 
+RewriteRule replaceCall(IncludeFormat Format) {
+  using namespace ::clang::ast_matchers;
+  RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))),
+                              change(text("other()")), text("no message"));
+  addInclude(Rule, "clang/OtherLib.h", Format);
+  return Rule;
+}
+
+template <IncludeFormat Format>
+class IncludeCheck : public TransformerClangTidyCheck {
+public:
+  IncludeCheck(StringRef Name, ClangTidyContext *Context)
+      : TransformerClangTidyCheck(replaceCall(Format), Name, Context) {}
+};
+
+TEST(TransformerClangTidyCheckTest, AddIncludeQuoted) {
+
+  std::string Input = R"cc(
+    int f(int x);
+    int h(int x) { return f(x); }
+  )cc";
+  std::string Expected = R"cc(#include "clang/OtherLib.h"
+
+
+    int f(int x);
+    int h(int x) { return other(); }
+  )cc";
+
+  EXPECT_EQ(Expected,
+            test::runCheckOnCode<IncludeCheck<IncludeFormat::Quoted>>(Input));
+}
+
+TEST(TransformerClangTidyCheckTest, AddIncludeAngled) {
+  std::string Input = R"cc(
+    int f(int x);
+    int h(int x) { return f(x); }
+  )cc";
+  std::string Expected = R"cc(#include <clang/OtherLib.h>
+
+
+    int f(int x);
+    int h(int x) { return other(); }
+  )cc";
+
+  EXPECT_EQ(Expected,
+            test::runCheckOnCode<IncludeCheck<IncludeFormat::Angled>>(Input));
+}
+
 } // namespace
 } // namespace utils
 } // namespace tidy