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

[ymandel]

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.

[gribozavr]

clang/include/clang/Tooling/Refactoring/Transformer.h

123

s/that//

152

The example looks incomplete to me because it does not call addInclude...

[ymandel]

Responded to comments.

[ymandel]

comments tweak.

[gribozavr]

LGTM, thanks!

clang/include/clang/Tooling/Refactoring/Transformer.h

 //   change(cat("different_expr"))
 // \endcode
 struct ASTEdit {
   RangeSelector TargetRange;
   TextGenerator Replacement;
   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
 // contains each of the following components:
 //
 // * Matcher: the pattern term, expressed as clang matchers (with Transformer
 //   extensions).
 //

 //
 // The \c Transformer class can be used to apply the rewrite rule and obtain the
 // corresponding replacements.
 struct RewriteRule {
   struct Case {
     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;
 
   // ID used as the default target of each match. The node described by the
   // matcher is should always be bound to this id.
   static constexpr llvm::StringLiteral RootID = "___root___";
 };
 
 /// Convenience function for constructing a simple \c RewriteRule.
 RewriteRule makeRule(ast_matchers::internal::DynTypedMatcher M,
                      SmallVector<ASTEdit, 1> Edits,
                      TextGenerator Explanation = nullptr);
 
 /// Convenience overload of \c makeRule for common case of only one edit.
 inline RewriteRule makeRule(ast_matchers::internal::DynTypedMatcher M,
                             ASTEdit Edit,
                             TextGenerator Explanation = nullptr) {
   SmallVector<ASTEdit, 1> Edits;
   Edits.emplace_back(std::move(Edit));
   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
 /// base class in the AST hierarchy).  However, this constraint is caused by an
 /// implementation detail and should be lifted in the future.
 //
 // `applyFirst` is like an `anyOf` matcher with an edit action attached to each
 // of its cases. Anywhere you'd use `anyOf(m1.bind("id1"), m2.bind("id2"))` and

clang/lib/Tooling/Refactoring/Transformer.cpp

ASTEdit tooling::change(RangeSelector S, TextGenerator Replacement) {
   ASTEdit E;
   E.TargetRange = std::move(S);
   E.Replacement = std::move(Replacement);
   return E;
 }
 
 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
 // the implicit relationship of Type and QualType.
 static bool isBaseOf(ASTNodeKind A, ASTNodeKind B) {
   static auto TypeKind = ASTNodeKind::getFromNodeKind<Type>();
   static auto QualKind = ASTNodeKind::getFromNodeKind<QualType>();
   /// Mimic the implicit conversions of Matcher<>.
void Transformer::run(const MatchResult &Result) {
   // Verify the existence and validity of the AST node that roots this rule.
   auto &NodesMap = Result.Nodes.getMap();
   auto Root = NodesMap.find(RewriteRule::RootID);
   assert(Root != NodesMap.end() && "Transformation failed: missing root node.");
   SourceLocation RootLoc = Result.SourceManager->getExpansionLoc(
       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;
   }
 
   if (Transformations->empty()) {
     // No rewrite applied (but no error encountered either).
     RootLoc.print(llvm::errs() << "note: skipping match at loc ",
                   *Result.SourceManager);
     llvm::errs() << "\n";
     return;
   }
 
   // Record the results in the AtomicChange.
   AtomicChange AC(*Result.SourceManager, RootLoc);
   for (const auto &T : *Transformations) {
     if (auto Err = AC.replace(*Result.SourceManager, T.Range, T.Replacement)) {
       Consumer(std::move(Err));
       return;
     }
   }
 
+  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));
 }

clang/unittests/Tooling/TransformerTest.cpp

std::string Expected = R"cc(
     proto::ProtoCommandLineFlag flag;
     int x = EXPR.foo();
     int y = flag.GetProto().foo();
   )cc";
 
   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),
                               change(name(Fun), text("good")));
 
   std::string Input = R"cc(
     int bad(int x);
     int bad(int x) { return x * x; }