[LibTooling] Extend `RewriteRule` with support for adding includes.

Summary:
This revision allows users to specify the insertion of an included directive (at
the top of the file being rewritten) as part of a rewrite rule.  These
directives are bundled with `RewriteRule` cases, so that different cases can
potentially result in different include actions.

Reviewers: ilya-biryukov, gribozavr

Subscribers: cfe-commits

Tags: #clang

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

llvm-svn: 364917

3 files changed, 81 insertions, 4 deletions

diff --git a/clang/include/clang/Tooling/Refactoring/Transformer.h b/clang/include/clang/Tooling/Refactoring/Transformer.h
index 1b71d1541d0..6d9c5a37cc1 100644
--- a/clang/include/clang/Tooling/Refactoring/Transformer.h
+++ b/clang/include/clang/Tooling/Refactoring/Transformer.h
@@ -86,6 +86,12 @@ struct ASTEdit {
   TextGenerator Note;
 };
 
+/// Format of the path in an include directive -- angle brackets or quotes.
+enum class IncludeFormat {
+  Quoted,
+  Angled,
+};
+
 /// Description of a source-code transformation.
 //
 // A *rewrite rule* describes a transformation of source code. A simple rule
@@ -114,6 +120,10 @@ struct RewriteRule {
     ast_matchers::internal::DynTypedMatcher Matcher;
     SmallVector<ASTEdit, 1> Edits;
     TextGenerator Explanation;
+    // Include paths to add to the file affected by this case.  These are
+    // bundled with the `Case`, rather than the `RewriteRule`, because each case
+    // might have different associated changes to the includes.
+    std::vector<std::pair<std::string, IncludeFormat>> AddedIncludes;
   };
   // We expect RewriteRules will most commonly include only one case.
   SmallVector<Case, 1> Cases;
@@ -137,6 +147,19 @@ inline RewriteRule makeRule(ast_matchers::internal::DynTypedMatcher M,
   return makeRule(std::move(M), std::move(Edits), std::move(Explanation));
 }
 
+/// For every case in Rule, adds an include directive for the given header. The
+/// common use is assumed to be a rule with only one case. For example, to
+/// replace a function call and add headers corresponding to the new code, one
+/// could write:
+/// \code
+///   auto R = makeRule(callExpr(callee(functionDecl(hasName("foo")))),
+///            change(text("bar()")));
+///   AddInclude(R, "path/to/bar_header.h");
+///   AddInclude(R, "vector", IncludeFormat::Angled);
+/// \endcode
+void addInclude(RewriteRule &Rule, llvm::StringRef Header,
+                IncludeFormat Format = IncludeFormat::Quoted);
+
 /// Applies the first rule whose pattern matches; other rules are ignored.
 ///
 /// N.B. All of the rules must use the same kind of matcher (that is, share a
diff --git a/clang/lib/Tooling/Refactoring/Transformer.cpp b/clang/lib/Tooling/Refactoring/Transformer.cpp
index 76573d6b59f..8e6fe6c7a94 100644
--- a/clang/lib/Tooling/Refactoring/Transformer.cpp
+++ b/clang/lib/Tooling/Refactoring/Transformer.cpp
@@ -98,8 +98,14 @@ ASTEdit tooling::change(RangeSelector S, TextGenerator Replacement) {
 
 RewriteRule tooling::makeRule(DynTypedMatcher M, SmallVector<ASTEdit, 1> Edits,
                               TextGenerator Explanation) {
-  return RewriteRule{{RewriteRule::Case{std::move(M), std::move(Edits),
-                                        std::move(Explanation)}}};
+  return RewriteRule{{RewriteRule::Case{
+      std::move(M), std::move(Edits), std::move(Explanation), {}}}};
+}
+
+void tooling::addInclude(RewriteRule &Rule, StringRef Header,
+                         IncludeFormat Format) {
+  for (auto &Case : Rule.Cases)
+    Case.AddedIncludes.emplace_back(Header.str(), Format);
 }
 
 // Determines whether A is a base type of B in the class hierarchy, including
@@ -217,8 +223,8 @@ void Transformer::run(const MatchResult &Result) {
       Root->second.getSourceRange().getBegin());
   assert(RootLoc.isValid() && "Invalid location for Root node of match.");
 
-  auto Transformations = tooling::detail::translateEdits(
-      Result, tooling::detail::findSelectedCase(Result, Rule).Edits);
+  RewriteRule::Case Case = tooling::detail::findSelectedCase(Result, Rule);
+  auto Transformations = tooling::detail::translateEdits(Result, Case.Edits);
   if (!Transformations) {
     Consumer(Transformations.takeError());
     return;
@@ -241,5 +247,17 @@ void Transformer::run(const MatchResult &Result) {
     }
   }
 
+  for (const auto &I : Case.AddedIncludes) {
+    auto &Header = I.first;
+    switch (I.second) {
+      case IncludeFormat::Quoted:
+        AC.addHeader(Header);
+        break;
+      case IncludeFormat::Angled:
+        AC.addHeader((llvm::Twine("<") + Header + ">").str());
+        break;
+    }
+  }
+
   Consumer(std::move(AC));
 }
diff --git a/clang/unittests/Tooling/TransformerTest.cpp b/clang/unittests/Tooling/TransformerTest.cpp
index e9de00d22a5..64f511b5fab 100644
--- a/clang/unittests/Tooling/TransformerTest.cpp
+++ b/clang/unittests/Tooling/TransformerTest.cpp
@@ -198,6 +198,42 @@ TEST_F(TransformerTest, Flag) {
   testRule(std::move(Rule), Input, Expected);
 }
 
+TEST_F(TransformerTest, AddIncludeQuoted) {
+  RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))),
+                              change(text("other()")));
+  addInclude(Rule, "clang/OtherLib.h");
+
+  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";
+
+  testRule(Rule, Input, Expected);
+}
+
+TEST_F(TransformerTest, AddIncludeAngled) {
+  RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))),
+                              change(text("other()")));
+  addInclude(Rule, "clang/OtherLib.h", IncludeFormat::Angled);
+
+  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";
+
+  testRule(Rule, Input, Expected);
+}
+
 TEST_F(TransformerTest, NodePartNameNamedDecl) {
   StringRef Fun = "fun";
   RewriteRule Rule = makeRule(functionDecl(hasName("bad")).bind(Fun),