[libTooling] Move RewriteRule include edits to ASTEdit granularity.

[ymandel]

Summary

Currently, changes to includes are applied to an entire rule. However,
include changes may be specific to particular edits within a rule (for example,
they may apply to one file but not another). Also, include changes may need to
carry metadata, just like other changes. So, we make include changes first-class
edits.

[ymandel]

tweak

[ymandel]

Updated clang-tidy transformer interpreter correspondingly.

[ymandel]

fixed typo

[gribozavr2]

clang/include/clang/Tooling/Transformer/RewriteRule.h

39

[ymandel]

reword comment per suggestion.

clang/include/clang/Tooling/Transformer/RewriteRule.h

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Error.h"
 #include <functional>
 #include <string>
 #include <utility>
 
 namespace clang {
 namespace transformer {
+// Specifies how to interpret an edit.
+enum class EditKind {
+  // Edits a source range in the file.
+  Range,
+  // Inserts an include in the file. The `Replacement` field is the name of the
+  // file for which to add an include.
+  AddInclude,
+};
+
 /// A concrete description of a source edit, represented by a character range in
 /// the source to be replaced and a corresponding replacement string.
 struct Edit {
+  EditKind Kind = EditKind::Range;
   CharSourceRange Range;
   std::string Replacement;
   llvm::Any Metadata;
 };
 
+/// Format of the path in an include directive -- angle brackets or quotes.
+enum class IncludeFormat {
+  Quoted,
+  Angled,
+};
+
 /// Maps a match result to a list of concrete edits (with possible
 /// failure). This type is a building block of rewrite rules, but users will
 /// generally work in terms of `ASTEdit`s (below) rather than directly in terms
 /// of `EditGenerator`.
 using EditGenerator = MatchConsumer<llvm::SmallVector<Edit, 1>>;
 
 using TextGenerator = std::shared_ptr<MatchComputation<std::string>>;
 

 //   changeTo(callArgs(call), cat(x, ",", y))
 // \endcode
 // Or, if you are changing the node corresponding to the rule's matcher, you can
 // use the single-argument override of \c change:
 // \code
 //   changeTo(cat("different_expr"))
 // \endcode
 struct ASTEdit {
+  EditKind Kind = EditKind::Range;
   RangeSelector TargetRange;
   TextGenerator Replacement;
   TextGenerator Note;
   // Not all transformations will want or need to attach metadata and therefore
   // should not be requierd to do so.
   AnyGenerator Metadata = [](const ast_matchers::MatchFinder::MatchResult &)
       -> llvm::Expected<llvm::Any> {
     return llvm::Expected<llvm::Any>(llvm::Any());

 /// unchanged.
 inline ASTEdit insertAfter(RangeSelector S, TextGenerator Replacement) {
   return changeTo(after(std::move(S)), std::move(Replacement));
 }
 
 /// Removes the source selected by \p S.
 ASTEdit remove(RangeSelector S);
 
+/// Adds an include directive for the given header to the file of `Target`. The
+/// particular location specified by `Target` is ignored.
+ASTEdit addInclude(RangeSelector Target, StringRef Header,
+                   IncludeFormat Format = IncludeFormat::Quoted);
+
+/// Adds an include directive for the given header to the file associated with
+/// `RootID`.
+inline ASTEdit addInclude(StringRef Header,
+                          IncludeFormat Format = IncludeFormat::Quoted) {
+  return addInclude(before(node(RootID)), Header, Format);
+}
+
 // FIXME: If `Metadata` returns an `llvm::Expected<T>` the `AnyGenerator` will
 // construct an `llvm::Expected<llvm::Any>` where no error is present but the
 // `llvm::Any` holds the error. This is unlikely but potentially surprising.
 // Perhaps the `llvm::Expected` should be unwrapped, or perhaps this should be a
 // compile-time error. No solution here is perfect.
 //
 // Note: This function template accepts any type callable with a MatchResult
 // rather than a `std::function` because the return-type needs to be deduced. If

 /// Assuming that the inner range is enclosed by the outer range, creates
 /// precision edits to remove the parts of the outer range that are not included
 /// in the inner range.
 inline EditGenerator shrinkTo(RangeSelector outer, RangeSelector inner) {
   return editList({remove(enclose(before(outer), before(inner))),
                    remove(enclose(after(inner), after(outer)))});
 }
 
-/// 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;
     EditGenerator 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;
 
   /// DEPRECATED: use `::clang::transformer::RootID` instead.
   static const llvm::StringRef RootID;
 };
 

clang/lib/Tooling/Transformer/RewriteRule.cpp

if (!EditRange)
       return SmallVector<Edit, 0>();
     auto Replacement = E.Replacement->eval(Result);
     if (!Replacement)
       return Replacement.takeError();
     auto Metadata = E.Metadata(Result);
     if (!Metadata)
       return Metadata.takeError();
     transformer::Edit T;
+    T.Kind = E.Kind;
     T.Range = *EditRange;
     T.Replacement = std::move(*Replacement);
     T.Metadata = std::move(*Metadata);
     Edits.push_back(std::move(T));
   }
   return Edits;
 }
 
llvm::Error eval(const ast_matchers::MatchFinder::MatchResult &,
     return llvm::Error::success();
   }
   std::string toString() const override {
     return (llvm::Twine("text(\"") + S + "\")").str();
   }
 };
 } // namespace
 
+static TextGenerator makeText(std::string S) {
+  return std::make_shared<SimpleTextGenerator>(std::move(S));
+}
+
 ASTEdit transformer::remove(RangeSelector S) {
-  return change(std::move(S), std::make_shared<SimpleTextGenerator>(""));
+  return change(std::move(S), makeText(""));
+}
+
+static std::string formatHeaderPath(StringRef Header, IncludeFormat Format) {
+  switch (Format) {
+  case transformer::IncludeFormat::Quoted:
+    return Header.str();
+  case transformer::IncludeFormat::Angled:
+    return ("<" + Header + ">").str();
+  }
+}
+
+ASTEdit transformer::addInclude(RangeSelector Target, StringRef Header,
+                                IncludeFormat Format) {
+  ASTEdit E;
+  E.Kind = EditKind::AddInclude;
+  E.TargetRange = Target;
+  E.Replacement = makeText(formatHeaderPath(Header, Format));
+  return E;
 }
 
 RewriteRule transformer::makeRule(DynTypedMatcher M, EditGenerator 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)}}};
 }
 
 namespace {
 
 /// Unconditionally binds the given node set before trying `InnerMatcher` and
 /// keeps the bound nodes on a successful match.
 template <typename T>
 class BindingsMatcher : public ast_matchers::internal::MatcherInterface<T> {
return llvm::make_error<llvm::StringError>(
         "type unsupported for recursive rewriting, ID=\"" + NodeId +
             "\", Kind=" + It->second.getNodeKind().asStringRef());
   };
 }
 
 void transformer::addInclude(RewriteRule &Rule, StringRef Header,
                              IncludeFormat Format) {
   for (auto &Case : Rule.Cases)
-    Case.AddedIncludes.emplace_back(Header.str(), Format);
+    Case.Edits = flatten(std::move(Case.Edits), addInclude(Header, Format));
 }
 
 #ifndef NDEBUG
 // Filters for supported matcher kinds. FIXME: Explicitly list the allowed kinds
 // (all node matcher types except for `QualType` and `Type`), rather than just
 // banning `QualType` and `Type`.
 static bool hasValidKind(const DynTypedMatcher &M) {
   return !M.canConvertTo<QualType>();

clang/lib/Tooling/Transformer/Transformer.cpp

void Transformer::run(const MatchFinder::MatchResult &Result) {
   std::map<FileID, AtomicChange> ChangesByFileID;
   for (const auto &T : *Transformations) {
     auto ID = Result.SourceManager->getFileID(T.Range.getBegin());
     auto Iter = ChangesByFileID
                     .emplace(ID, AtomicChange(*Result.SourceManager,
                                               T.Range.getBegin(), T.Metadata))
                     .first;
     auto &AC = Iter->second;
-    if (auto Err = AC.replace(*Result.SourceManager, T.Range, T.Replacement)) {
+    switch (T.Kind) {
+    case transformer::EditKind::Range:
+      if (auto Err =
+              AC.replace(*Result.SourceManager, T.Range, T.Replacement)) {
         Consumer(std::move(Err));
         return;
       }
-  }
-
-  for (auto &IDChangePair : ChangesByFileID) {
-    auto &AC = IDChangePair.second;
-    // FIXME: this will add includes to *all* changed files, which may not be
-    // the intent. We should upgrade the representation to allow associating
-    // headers with specific edits.
-    for (const auto &I : Case.AddedIncludes) {
-      auto &Header = I.first;
-      switch (I.second) {
-      case transformer::IncludeFormat::Quoted:
-        AC.addHeader(Header);
-        break;
-      case transformer::IncludeFormat::Angled:
-        AC.addHeader((llvm::Twine("<") + Header + ">").str());
-        break;
-      }
-    }
+    case transformer::EditKind::AddInclude:
+      AC.addHeader(T.Replacement);
 
-    Consumer(std::move(AC));
-  }
+  for (auto &IDChangePair : ChangesByFileID)
+    Consumer(std::move(IDChangePair.second));
 }

clang/unittests/Tooling/TransformerTest.cpp

 #include "llvm/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using namespace clang;
 using namespace tooling;
 using namespace ast_matchers;
 namespace {
+using ::testing::ElementsAre;
 using ::testing::IsEmpty;
 using transformer::cat;
 using transformer::changeTo;
 using transformer::RewriteRule;
 
 constexpr char KHeaderContents[] = R"cc(
   struct string {
     string(const char*);
std::string Expected = R"cc(
     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")))),
+               {addInclude("clang/OtherLib.h"), changeTo(cat("other()"))});
+
+  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")))),
+      {addInclude("clang/OtherLib.h", transformer::IncludeFormat::Angled),
+       changeTo(cat("other()"))});
+
+  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, AddIncludeQuotedForRule) {
   RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))),
                               changeTo(cat("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) {
+TEST_F(TransformerTest, AddIncludeAngledForRule) {
   RewriteRule Rule = makeRule(callExpr(callee(functionDecl(hasName("f")))),
                               changeTo(cat("other()")));
   addInclude(Rule, "clang/OtherLib.h", transformer::IncludeFormat::Angled);
 
   std::string Input = R"cc(
     int f(int x);
     int h(int x) { return f(x); }
   )cc";
TEST_F(TransformerTest, MultipleFiles) {
   EXPECT_THAT(Changes[1].getRemovedHeaders(), IsEmpty());
   UpdatedCode = clang::tooling::applyAllReplacements(
       Source, Changes[1].getReplacements());
   ASSERT_TRUE(static_cast<bool>(UpdatedCode))
       << "Could not update code: " << llvm::toString(UpdatedCode.takeError());
   EXPECT_EQ(format(*UpdatedCode), format(R"cc(#include "input.h"
                         ;)cc"));
 }
+
+TEST_F(TransformerTest, AddIncludeMultipleFiles) {
+  std::string Header = R"cc(void RemoveThisFunction();)cc";
+  std::string Source = R"cc(#include "input.h"
+                            void Foo() {RemoveThisFunction();})cc";
+  Transformer T(
+      makeRule(callExpr(callee(
+                   functionDecl(hasName("RemoveThisFunction")).bind("fun"))),
+               addInclude(node("fun"), "header.h")),
+      consumer());
+  T.registerMatchers(&MatchFinder);
+  auto Factory = newFrontendActionFactory(&MatchFinder);
+  EXPECT_TRUE(runToolOnCodeWithArgs(
+      Factory->create(), Source, std::vector<std::string>(), "input.cc",
+      "clang-tool", std::make_shared<PCHContainerOperations>(),
+      {{"input.h", Header}}));
+
+  ASSERT_EQ(Changes.size(), 1U);
+  ASSERT_EQ(Changes[0].getFilePath(), "./input.h");
+  EXPECT_THAT(Changes[0].getInsertedHeaders(), ElementsAre("header.h"));
+  EXPECT_THAT(Changes[0].getRemovedHeaders(), IsEmpty());
+  llvm::Expected<std::string> UpdatedCode =
+      clang::tooling::applyAllReplacements(Header,
+                                           Changes[0].getReplacements());
+  ASSERT_TRUE(static_cast<bool>(UpdatedCode))
+      << "Could not update code: " << llvm::toString(UpdatedCode.takeError());
+  EXPECT_EQ(format(*UpdatedCode), format(Header));
+}
 } // namespace