1 files changed, 46 insertions, 22 deletions

diff --git a/clang/lib/Tooling/Transformer/RewriteRule.cpp b/clang/lib/Tooling/Transformer/RewriteRule.cpp
index 20d3a371950a..995bec03cd66 100644
--- a/clang/lib/Tooling/Transformer/RewriteRule.cpp
+++ b/clang/lib/Tooling/Transformer/RewriteRule.cpp
@@ -25,16 +25,14 @@ using namespace transformer;
 
 using ast_matchers::MatchFinder;
 using ast_matchers::internal::DynTypedMatcher;
-using ast_type_traits::ASTNodeKind;
 
 using MatchResult = MatchFinder::MatchResult;
 
-Expected<SmallVector<transformer::detail::Transformation, 1>>
-transformer::detail::translateEdits(const MatchResult &Result,
-                                llvm::ArrayRef<ASTEdit> Edits) {
-  SmallVector<transformer::detail::Transformation, 1> Transformations;
-  for (const auto &Edit : Edits) {
-    Expected<CharSourceRange> Range = Edit.TargetRange(Result);
+static Expected<SmallVector<transformer::Edit, 1>>
+translateEdits(const MatchResult &Result, ArrayRef<ASTEdit> ASTEdits) {
+  SmallVector<transformer::Edit, 1> Edits;
+  for (const auto &E : ASTEdits) {
+    Expected<CharSourceRange> Range = E.TargetRange(Result);
     if (!Range)
       return Range.takeError();
     llvm::Optional<CharSourceRange> EditRange =
@@ -42,21 +40,34 @@ transformer::detail::translateEdits(const MatchResult &Result,
     // FIXME: let user specify whether to treat this case as an error or ignore
     // it as is currently done.
     if (!EditRange)
-      return SmallVector<Transformation, 0>();
-    auto Replacement = Edit.Replacement->eval(Result);
+      return SmallVector<Edit, 0>();
+    auto Replacement = E.Replacement->eval(Result);
     if (!Replacement)
       return Replacement.takeError();
-    transformer::detail::Transformation T;
+    transformer::Edit T;
     T.Range = *EditRange;
     T.Replacement = std::move(*Replacement);
-    Transformations.push_back(std::move(T));
+    T.Metadata = E.Metadata;
+    Edits.push_back(std::move(T));
   }
-  return Transformations;
+  return Edits;
 }
 
-ASTEdit transformer::changeTo(RangeSelector S, TextGenerator Replacement) {
+EditGenerator transformer::editList(SmallVector<ASTEdit, 1> Edits) {
+  return [Edits = std::move(Edits)](const MatchResult &Result) {
+    return translateEdits(Result, Edits);
+  };
+}
+
+EditGenerator transformer::edit(ASTEdit Edit) {
+  return [Edit = std::move(Edit)](const MatchResult &Result) {
+    return translateEdits(Result, {Edit});
+  };
+}
+
+ASTEdit transformer::changeTo(RangeSelector Target, TextGenerator Replacement) {
   ASTEdit E;
-  E.TargetRange = std::move(S);
+  E.TargetRange = std::move(Target);
   E.Replacement = std::move(Replacement);
   return E;
 }
@@ -83,8 +94,9 @@ ASTEdit transformer::remove(RangeSelector S) {
   return change(std::move(S), std::make_shared<SimpleTextGenerator>(""));
 }
 
-RewriteRule transformer::makeRule(DynTypedMatcher M, SmallVector<ASTEdit, 1> Edits,
-                              TextGenerator Explanation) {
+RewriteRule transformer::makeRule(ast_matchers::internal::DynTypedMatcher M,
+                                  EditGenerator Edits,
+                                  TextGenerator Explanation) {
   return RewriteRule{{RewriteRule::Case{
       std::move(M), std::move(Edits), std::move(Explanation), {}}}};
 }
@@ -105,10 +117,13 @@ static bool hasValidKind(const DynTypedMatcher &M) {
 #endif
 
 // Binds each rule's matcher to a unique (and deterministic) tag based on
-// `TagBase` and the id paired with the case.
+// `TagBase` and the id paired with the case. All of the returned matchers have
+// their traversal kind explicitly set, either based on a pre-set kind or to the
+// provided `DefaultTraversalKind`.
 static std::vector<DynTypedMatcher> taggedMatchers(
     StringRef TagBase,
-    const SmallVectorImpl<std::pair<size_t, RewriteRule::Case>> &Cases) {
+    const SmallVectorImpl<std::pair<size_t, RewriteRule::Case>> &Cases,
+    ast_type_traits::TraversalKind DefaultTraversalKind) {
   std::vector<DynTypedMatcher> Matchers;
   Matchers.reserve(Cases.size());
   for (const auto &Case : Cases) {
@@ -116,8 +131,10 @@ static std::vector<DynTypedMatcher> taggedMatchers(
     // HACK: Many matchers are not bindable, so ensure that tryBind will work.
     DynTypedMatcher BoundMatcher(Case.second.Matcher);
     BoundMatcher.setAllowBind(true);
-    auto M = BoundMatcher.tryBind(Tag);
-    Matchers.push_back(*std::move(M));
+    auto M = *BoundMatcher.tryBind(Tag);
+    Matchers.push_back(!M.getTraversalKind()
+                           ? M.withTraversalKind(DefaultTraversalKind)
+                           : std::move(M));
   }
   return Matchers;
 }
@@ -147,14 +164,21 @@ transformer::detail::buildMatchers(const RewriteRule &Rule) {
     Buckets[Cases[I].Matcher.getSupportedKind()].emplace_back(I, Cases[I]);
   }
 
+  // Each anyOf explicitly controls the traversal kind. The anyOf itself is set
+  // to `TK_AsIs` to ensure no nodes are skipped, thereby deferring to the kind
+  // of the branches. Then, each branch is either left as is, if the kind is
+  // already set, or explicitly set to `TK_IgnoreUnlessSpelledInSource`. We
+  // choose this setting, because we think it is the one most friendly to
+  // beginners, who are (largely) the target audience of Transformer.
   std::vector<DynTypedMatcher> Matchers;
   for (const auto &Bucket : Buckets) {
     DynTypedMatcher M = DynTypedMatcher::constructVariadic(
         DynTypedMatcher::VO_AnyOf, Bucket.first,
-        taggedMatchers("Tag", Bucket.second));
+        taggedMatchers("Tag", Bucket.second, TK_IgnoreUnlessSpelledInSource));
     M.setAllowBind(true);
     // `tryBind` is guaranteed to succeed, because `AllowBind` was set to true.
-    Matchers.push_back(*M.tryBind(RewriteRule::RootID));
+    Matchers.push_back(
+        M.tryBind(RewriteRule::RootID)->withTraversalKind(TK_AsIs));
   }
   return Matchers;
 }