[Syntax] Build mapping from AST to syntax tree nodes

[ilya-biryukov]

Summary

Only used to assign roles to child nodes for now. This is more efficient
than doing range-based queries.

In the future, will be exposed in the public API of syntax trees.

[ilya-biryukov]

Remove the (now redundant) NodeAndRole class

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[ilya-biryukov]

• Remove debug logs
• Cosmetics

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[gribozavr2]

clang/lib/Tooling/Syntax/BuildTree.cpp

147

Also assert that From is not null either?

160

I think it is possible to use a PointerUnion as a DenseMap key (there's a DenseMapInfo for it).

491

I don't understand the first fixme (maybe it is stale).

The second fixme is not applicable after this patch.

[hlopko]

Taking over the patch in https://reviews.llvm.org/D76355.

clang/lib/Tooling/Syntax/BuildTree.cpp

147

Done.

160

Done.

491

Removed both FIXMEs now, asking Ilya offline for clarification.

[gribozavr2]

Superseded by https://reviews.llvm.org/D76355.

clang/include/clang/Tooling/Syntax/Tree.h

 protected:
   /// Find the first node with a corresponding role.
   syntax::Node *findChild(NodeRole R);
 
 private:
   /// Prepend \p Child to the list of children and and sets the parent pointer.
   /// A very low-level operation that does not check any invariants, only used
   /// by TreeBuilder and FactoryImpl.
-  /// EXPECTS: Role != NodeRoleDetached.
+  /// EXPECTS: Role != Detached.
   void prependChildLowLevel(Node *Child, NodeRole Role);
+  /// Like the previous overload, but does not set role for \p Child.
+  /// EXPECTS: Child->Role != Detached
+  void prependChildLowLevel(Node *Child);
   friend class TreeBuilder;
   friend class FactoryImpl;
 
   /// Replace a range of children [BeforeBegin->NextSibling, End) with a list of
   /// new nodes starting at \p New.
   /// Only used by MutationsImpl to implement higher-level mutation operations.
   /// (!) \p New can be null to model removal of the child range.
   void replaceChildRangeLowLevel(Node *BeforeBegin, Node *End, Node *New);

clang/lib/Tooling/Syntax/BuildTree.cpp

 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Tooling/Syntax/Nodes.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "clang/Tooling/Syntax/Tree.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/raw_ostream.h"
+#include <cstddef>
 #include <map>
 
 using namespace clang;
 
 LLVM_ATTRIBUTE_UNUSED
 static bool isImplicitExpr(clang::Expr *E) { return E->IgnoreImplicit() != E; }
 
 static SourceLocation getQualifiedNameStart(DeclaratorDecl *D) {
static SourceRange getDeclaratorRange(const SourceManager &SM, TypeLoc T,
   }
   if (Initializer.isValid()) {
     assert(SM.isBeforeInTranslationUnit(End, Initializer.getEnd()));
     End = Initializer.getEnd();
   }
   return SourceRange(Start, End);
 }
 
+namespace {
+/// All AST hierarchy roots that can be represented as pointers.
+using ASTPtr = llvm::PointerUnion<Stmt *, Decl *>;
+/// Maintains a mapping from AST to syntax tree nodes. This class will get more
+/// complicated as we support more kinds of AST nodes, e.g. TypeLocs.
+/// FIXME: expose this as public API.
+class ASTToSyntaxMapping {
+public:
+  void add(ASTPtr From, syntax::Tree *To) {
+    assert(To != nullptr);
+
+    bool Added = Nodes.insert({From.getOpaqueValue(), To}).second;
+    (void)Added;
+    assert(Added && "mapping added twice");
+  }
+
+  syntax::Tree *find(ASTPtr P) const {
+    return Nodes.lookup(P.getOpaqueValue());
+  }
+
+private:
+  // Keys are either Stmt* or Decl*.
+  llvm::DenseMap<void *, syntax::Tree *> Nodes;
+};
+} // namespace
+
 /// A helper class for constructing the syntax tree while traversing a clang
 /// AST.
 ///
 /// At each point of the traversal we maintain a list of pending nodes.
 /// Initially all tokens are added as pending nodes. When processing a clang AST
 /// node, the clients need to:
 ///   - create a corresponding syntax node,
 ///   - assign roles to all pending child nodes with 'markChild' and
for (const auto &T : Arena.tokenBuffer().expandedTokens())
       LocationToToken.insert({T.location().getRawEncoding(), &T});
   }
 
   llvm::BumpPtrAllocator &allocator() { return Arena.allocator(); }
   const SourceManager &sourceManager() const { return Arena.sourceManager(); }
 
   /// Populate children for \p New node, assuming it covers tokens from \p
   /// Range.
-  void foldNode(llvm::ArrayRef<syntax::Token> Range, syntax::Tree *New);
+  void foldNode(llvm::ArrayRef<syntax::Token> Range, syntax::Tree *New,
+                ASTPtr From) {
+    assert(New);
+    Pending.foldChildren(Arena, Range, New);
+    if (From)
+      Mapping.add(From, New);
+  }
+  void foldNode(llvm::ArrayRef<syntax::Token> Range, syntax::Tree *New,
+                TypeLoc L) {
+    // FIXME: add mapping for TypeLocs
+    foldNode(Range, New, nullptr);
+  }
 
   /// Must be called with the range of each `DeclaratorDecl`. Ensures the
   /// corresponding declarator nodes are covered by `SimpleDeclaration`.
   void noticeDeclRange(llvm::ArrayRef<syntax::Token> Range);
 
   /// Notifies that we should not consume trailing semicolon when computing
   /// token range of \p D.
   void noticeDeclWithoutSemicolon(Decl *D);
 
   /// Mark the \p Child node with a corresponding \p Role. All marked children
   /// should be consumed by foldNode.
   /// (!) when called on expressions (clang::Expr is derived from clang::Stmt),
   ///     wraps expressions into expression statement.
   void markStmtChild(Stmt *Child, NodeRole Role);
   /// Should be called for expressions in non-statement position to avoid
   /// wrapping into expression statement.
   void markExprChild(Expr *Child, NodeRole Role);
   /// Set role for a token starting at \p Loc.
   void markChildToken(SourceLocation Loc, NodeRole R);
   /// Set role for \p T.
   void markChildToken(const syntax::Token *T, NodeRole R);
 
-  /// Set role for the node that spans exactly \p Range.
-  void markChild(llvm::ArrayRef<syntax::Token> Range, NodeRole R);
+  /// Set role for \p N.
+  void markChild(syntax::Node *N, NodeRole R);
+  /// Set role for the syntax node matching \p N.
+  void markChild(ASTPtr N, NodeRole R);
+
   /// Set role for the delayed node that spans exactly \p Range.
   void markDelayedChild(llvm::ArrayRef<syntax::Token> Range, NodeRole R);
   // Set role for the node that may or may not be delayed. Node must span
   // exactly \p Range.
   void markMaybeDelayedChild(llvm::ArrayRef<syntax::Token> Range, NodeRole R);
 
   /// Finish building the tree and consume the root node.
   syntax::TranslationUnit *finalize() && {
withTrailingSemicolon(llvm::ArrayRef<syntax::Token> Tokens) const {
     assert(!Tokens.empty());
     assert(Tokens.back().kind() != tok::eof);
     // (!) we never consume 'eof', so looking at the next token is ok.
     if (Tokens.back().kind() != tok::semi && Tokens.end()->kind() == tok::semi)
       return llvm::makeArrayRef(Tokens.begin(), Tokens.end() + 1);
     return Tokens;
   }
 
+  void setRole(syntax::Node *N, NodeRole R) {
+    assert(N->role() == NodeRole::Detached);
+    N->Role = static_cast<unsigned>(R);
+  }
+
   /// A collection of trees covering the input tokens.
   /// When created, each tree corresponds to a single token in the file.
   /// Clients call 'foldChildren' to attach one or more subtrees to a parent
   /// node and update the list of trees accordingly.
   ///
   /// Ensures that added nodes properly nest and cover the whole token stream.
   struct Forest {
     Forest(syntax::Arena &A) {
       assert(!A.tokenBuffer().expandedTokens().empty());
       assert(A.tokenBuffer().expandedTokens().back().kind() == tok::eof);
       // Create all leaf nodes.
       // Note that we do not have 'eof' in the tree.
       for (auto &T : A.tokenBuffer().expandedTokens().drop_back()) {
         auto *L = new (A.allocator()) syntax::Leaf(&T);
         L->Original = true;
         L->CanModify = A.tokenBuffer().spelledForExpanded(T).hasValue();
-        Trees.insert(Trees.end(), {&T, NodeAndRole{L}});
+        Trees.insert(Trees.end(), {&T, L});
       }
     }
 
     ~Forest() { assert(DelayedFolds.empty()); }
 
     void assignRoleDelayed(llvm::ArrayRef<syntax::Token> Range,
                            syntax::NodeRole Role) {
       auto It = DelayedFolds.find(Range.begin());
void assignRole(llvm::ArrayRef<syntax::Token> Range,
                     syntax::NodeRole Role) {
       assert(!Range.empty());
       auto It = Trees.lower_bound(Range.begin());
       assert(It != Trees.end() && "no node found");
       assert(It->first == Range.begin() && "no child with the specified range");
       assert((std::next(It) == Trees.end() ||
               std::next(It)->first == Range.end()) &&
              "no child with the specified range");
-      It->second.Role = Role;
+      assert(It->second->role() == NodeRole::Detached &&
+             "re-assigning role for a child");
+      It->second->Role = static_cast<unsigned>(Role);
     }
 
     /// Add \p Node to the forest and attach child nodes based on \p Tokens.
     void foldChildren(const syntax::Arena &A,
                       llvm::ArrayRef<syntax::Token> Tokens,
                       syntax::Tree *Node) {
       // Execute delayed folds inside `Tokens`.
       auto BeginExecuted = DelayedFolds.lower_bound(Tokens.begin());
bool extendDelayedFold(llvm::ArrayRef<syntax::Token> ExtendedRange) {
       assert(It->second.End <= ExtendedRange.end());
       It->second.End = ExtendedRange.end();
       return true;
     }
 
     // EXPECTS: all tokens were consumed and are owned by a single root node.
     syntax::Node *finalize() && {
       assert(Trees.size() == 1);
-      auto *Root = Trees.begin()->second.Node;
+      auto *Root = Trees.begin()->second;
       Trees = {};
       return Root;
     }
 
     std::string str(const syntax::Arena &A) const {
       std::string R;
       for (auto It = Trees.begin(); It != Trees.end(); ++It) {
         unsigned CoveredTokens =
             It != Trees.end()
                 ? (std::next(It)->first - It->first)
                 : A.tokenBuffer().expandedTokens().end() - It->first;
 
         R += llvm::formatv("- '{0}' covers '{1}'+{2} tokens\n",
-                           It->second.Node->kind(),
+                           It->second->kind(),
                            It->first->text(A.sourceManager()), CoveredTokens);
-        R += It->second.Node->dump(A);
+        R += It->second->dump(A);
       }
       return R;
     }
 
   private:
     /// Implementation detail of `foldChildren`, does acutal folding ignoring
     /// delayed folds.
     void foldChildrenEager(const syntax::Arena &A,
                            llvm::ArrayRef<syntax::Token> Tokens,
                            syntax::Tree *Node) {
       assert(Node->firstChild() == nullptr && "node already has children");
 
       auto *FirstToken = Tokens.begin();
       auto BeginChildren = Trees.lower_bound(FirstToken);
       assert((BeginChildren == Trees.end() ||
               BeginChildren->first == FirstToken) &&
              "fold crosses boundaries of existing subtrees");
       auto EndChildren = Trees.lower_bound(Tokens.end());
       assert(
           (EndChildren == Trees.end() || EndChildren->first == Tokens.end()) &&
           "fold crosses boundaries of existing subtrees");
 
       // (!) we need to go in reverse order, because we can only prepend.
-      for (auto It = EndChildren; It != BeginChildren; --It)
-        Node->prependChildLowLevel(std::prev(It)->second.Node,
-                                   std::prev(It)->second.Role);
+      for (auto It = EndChildren; It != BeginChildren; --It) {
+        auto *C = std::prev(It)->second;
+        if (C->role() == NodeRole::Detached)
+          C->Role = static_cast<unsigned>(NodeRole::Unknown);
+        Node->prependChildLowLevel(C);
+      }
 
       // Mark that this node came from the AST and is backed by the source code.
       Node->Original = true;
       Node->CanModify = A.tokenBuffer().spelledForExpanded(Tokens).hasValue();
 
       Trees.erase(BeginChildren, EndChildren);
-      Trees.insert({FirstToken, NodeAndRole(Node)});
+      Trees.insert({FirstToken, Node});
     }
-    /// A with a role that should be assigned to it when adding to a parent.
-    struct NodeAndRole {
-      explicit NodeAndRole(syntax::Node *Node)
-          : Node(Node), Role(NodeRole::Unknown) {}
-
-      syntax::Node *Node;
-      NodeRole Role;
-    };
 
     /// Maps from the start token to a subtree starting at that token.
     /// Keys in the map are pointers into the array of expanded tokens, so
     /// pointer order corresponds to the order of preprocessor tokens.
     /// FIXME: storing the end tokens is redundant.
     /// FIXME: the key of a map is redundant, it is also stored in NodeForRange.
-    std::map<const syntax::Token *, NodeAndRole> Trees;
+    std::map<const syntax::Token *, syntax::Node *> Trees;
 
     /// See documentation of `foldChildrenDelayed` for details.
     struct DelayedFold {
       const syntax::Token *End = nullptr;
       syntax::Tree *Node = nullptr;
       NodeRole Role = NodeRole::Unknown;
     };
     std::map<const syntax::Token *, DelayedFold> DelayedFolds;
   };
 
   /// For debugging purposes.
   std::string str() { return Pending.str(Arena); }
 
   syntax::Arena &Arena;
   /// To quickly find tokens by their start location.
   llvm::DenseMap</*SourceLocation*/ unsigned, const syntax::Token *>
       LocationToToken;
   Forest Pending;
   llvm::DenseSet<Decl *> DeclsWithoutSemicolons;
+  ASTToSyntaxMapping Mapping;
 };
 
 namespace {
 class BuildTreeVisitor : public RecursiveASTVisitor<BuildTreeVisitor> {
 public:
   explicit BuildTreeVisitor(ASTContext &Ctx, syntax::TreeBuilder &Builder)
       : Builder(Builder), LangOpts(Ctx.getLangOpts()) {}
 
if (auto *V = llvm::dyn_cast<VarDecl>(DD)) {
       // (!) initializers in range-based-for are not part of the declarator
       if (I && !V->isCXXForRangeDecl())
         Initializer = I->getSourceRange();
     }
     auto Declarator = getDeclaratorRange(
         Builder.sourceManager(), DD->getTypeSourceInfo()->getTypeLoc(),
         getQualifiedNameStart(DD), Initializer);
     if (Declarator.isValid()) {
-      auto Tokens =
-          Builder.getRange(Declarator.getBegin(), Declarator.getEnd());
-      Builder.foldNode(Tokens, new (allocator()) syntax::SimpleDeclarator);
-      Builder.markChild(Tokens, syntax::NodeRole::SimpleDeclaration_declarator);
+      auto *N = new (allocator()) syntax::SimpleDeclarator;
+      Builder.foldNode(
+          Builder.getRange(Declarator.getBegin(), Declarator.getEnd()), N, DD);
+      Builder.markChild(N, syntax::NodeRole::SimpleDeclaration_declarator);
     }
 
     return true;
   }
 
   bool WalkUpFromTypedefNameDecl(TypedefNameDecl *D) {
     // Ensure declarators are covered by SimpleDeclaration.
     Builder.noticeDeclRange(Builder.getDeclRange(D));
 
     auto R = getDeclaratorRange(
         Builder.sourceManager(), D->getTypeSourceInfo()->getTypeLoc(),
         /*Name=*/D->getLocation(), /*Initializer=*/SourceRange());
     if (R.isValid()) {
-      auto Tokens = Builder.getRange(R.getBegin(), R.getEnd());
-      Builder.foldNode(Tokens, new (allocator()) syntax::SimpleDeclarator);
-      Builder.markChild(Tokens, syntax::NodeRole::SimpleDeclaration_declarator);
+      auto *N = new (allocator()) syntax::SimpleDeclarator;
+      Builder.foldNode(Builder.getRange(R.getBegin(), R.getEnd()), N, D);
+      Builder.markChild(N, syntax::NodeRole::SimpleDeclaration_declarator);
     }
     return true;
   }
 
   bool VisitDecl(Decl *D) {
     assert(!D->isImplicit());
     Builder.foldNode(Builder.getDeclRange(D),
-                     new (allocator()) syntax::UnknownDeclaration());
+                     new (allocator()) syntax::UnknownDeclaration(), D);
     return true;
   }
 
   // RAV does not call WalkUpFrom* on explicit instantiations, so we have to
   // override Traverse.
   // FIXME: make RAV call WalkUpFrom* instead.
   bool
   TraverseClassTemplateSpecializationDecl(ClassTemplateSpecializationDecl *C) {
     if (!RecursiveASTVisitor::TraverseClassTemplateSpecializationDecl(C))
       return false;
     if (C->isExplicitSpecialization())
       return true; // we are only interested in explicit instantiations.
-    if (!WalkUpFromClassTemplateSpecializationDecl(C))
-      return false;
+    auto *Declaration =
+        cast<syntax::SimpleDeclaration>(handleFreeStandingTagDecl(C));
     foldExplicitTemplateInstantiation(
         Builder.getTemplateRange(C), Builder.findToken(C->getExternLoc()),
-        Builder.findToken(C->getTemplateKeywordLoc()), Builder.getDeclRange(C));
+        Builder.findToken(C->getTemplateKeywordLoc()), Declaration);
     return true;
   }
 
   bool WalkUpFromTemplateDecl(TemplateDecl *S) {
     foldTemplateDeclaration(
         Builder.getDeclRange(S),
         Builder.findToken(S->getTemplateParameters()->getTemplateLoc()),
-        Builder.getDeclRange(S->getTemplatedDecl()));
+        Builder.getDeclRange(S->getTemplatedDecl()), S);
     return true;
   }
 
   bool WalkUpFromTagDecl(TagDecl *C) {
     // FIXME: build the ClassSpecifier node.
     if (!C->isFreeStanding()) {
       assert(C->getNumTemplateParameterLists() == 0);
       return true;
     }
+    handleFreeStandingTagDecl(C);
+    return true;
+  }
+
+  syntax::Declaration *handleFreeStandingTagDecl(TagDecl *C) {
+    assert(C->isFreeStanding());
     // Class is a declaration specifier and needs a spanning declaration node.
     auto DeclarationRange = Builder.getDeclRange(C);
-    Builder.foldNode(DeclarationRange,
-                     new (allocator()) syntax::SimpleDeclaration);
+    syntax::Declaration *Result = new (allocator()) syntax::SimpleDeclaration;
+    Builder.foldNode(DeclarationRange, Result, nullptr);
 
     // Build TemplateDeclaration nodes if we had template parameters.
     auto ConsumeTemplateParameters = [&](const TemplateParameterList &L) {
       auto *TemplateKW = Builder.findToken(L.getTemplateLoc());
       auto R = llvm::makeArrayRef(TemplateKW, DeclarationRange.end());
-      foldTemplateDeclaration(R, TemplateKW, DeclarationRange);
-
+      Result =
+          foldTemplateDeclaration(R, TemplateKW, DeclarationRange, nullptr);
       DeclarationRange = R;
     };
     if (auto *S = llvm::dyn_cast<ClassTemplatePartialSpecializationDecl>(C))
       ConsumeTemplateParameters(*S->getTemplateParameters());
     for (unsigned I = C->getNumTemplateParameterLists(); 0 < I; --I)
       ConsumeTemplateParameters(*C->getTemplateParameterList(I - 1));
-    return true;
+    return Result;
   }
 
   bool WalkUpFromTranslationUnitDecl(TranslationUnitDecl *TU) {
     // (!) we do not want to call VisitDecl(), the declaration for translation
     // unit is built by finalize().
     return true;
   }
 
   bool WalkUpFromCompoundStmt(CompoundStmt *S) {
     using NodeRole = syntax::NodeRole;
 
     Builder.markChildToken(S->getLBracLoc(), NodeRole::OpenParen);
     for (auto *Child : S->body())
       Builder.markStmtChild(Child, NodeRole::CompoundStatement_statement);
     Builder.markChildToken(S->getRBracLoc(), NodeRole::CloseParen);
 
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::CompoundStatement);
+                     new (allocator()) syntax::CompoundStatement, S);
     return true;
   }
 
   // Some statements are not yet handled by syntax trees.
   bool WalkUpFromStmt(Stmt *S) {
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::UnknownStatement);
+                     new (allocator()) syntax::UnknownStatement, S);
     return true;
   }
 
   bool TraverseCXXForRangeStmt(CXXForRangeStmt *S) {
     // We override to traverse range initializer as VarDecl.
     // RAV traverses it as a statement, we produce invalid node kinds in that
     // case.
     // FIXME: should do this in RAV instead?
bool TraverseStmt(Stmt *S) {
     }
     return RecursiveASTVisitor::TraverseStmt(S);
   }
 
   // Some expressions are not yet handled by syntax trees.
   bool WalkUpFromExpr(Expr *E) {
     assert(!isImplicitExpr(E) && "should be handled by TraverseStmt");
     Builder.foldNode(Builder.getExprRange(E),
-                     new (allocator()) syntax::UnknownExpression);
+                     new (allocator()) syntax::UnknownExpression, E);
     return true;
   }
 
   bool WalkUpFromNamespaceDecl(NamespaceDecl *S) {
     auto Tokens = Builder.getDeclRange(S);
     if (Tokens.front().kind() == tok::coloncolon) {
       // Handle nested namespace definitions. Those start at '::' token, e.g.
       // namespace a^::b {}
       // FIXME: build corresponding nodes for the name of this namespace.
       return true;
     }
-    Builder.foldNode(Tokens, new (allocator()) syntax::NamespaceDefinition);
+    Builder.foldNode(Tokens, new (allocator()) syntax::NamespaceDefinition, S);
     return true;
   }
 
   bool TraverseParenTypeLoc(ParenTypeLoc L) {
     // (!) we reverse order of traversal to get the proper syntax structure.
     if (!WalkUpFromParenTypeLoc(L))
       return false;
     return TraverseTypeLoc(L.getInnerLoc());
   }
 
   bool WalkUpFromParenTypeLoc(ParenTypeLoc L) {
     Builder.markChildToken(L.getLParenLoc(), syntax::NodeRole::OpenParen);
     Builder.markChildToken(L.getRParenLoc(), syntax::NodeRole::CloseParen);
     Builder.foldNode(Builder.getRange(L.getLParenLoc(), L.getRParenLoc()),
-                     new (allocator()) syntax::ParenDeclarator);
+                     new (allocator()) syntax::ParenDeclarator, L);
     return true;
   }
 
   // Declarator chunks, they are produced by type locs and some clang::Decls.
   bool WalkUpFromArrayTypeLoc(ArrayTypeLoc L) {
     Builder.markChildToken(L.getLBracketLoc(), syntax::NodeRole::OpenParen);
     Builder.markExprChild(L.getSizeExpr(),
                           syntax::NodeRole::ArraySubscript_sizeExpression);
     Builder.markChildToken(L.getRBracketLoc(), syntax::NodeRole::CloseParen);
     Builder.foldNode(Builder.getRange(L.getLBracketLoc(), L.getRBracketLoc()),
-                     new (allocator()) syntax::ArraySubscript);
+                     new (allocator()) syntax::ArraySubscript, L);
     return true;
   }
 
   bool WalkUpFromFunctionTypeLoc(FunctionTypeLoc L) {
     Builder.markChildToken(L.getLParenLoc(), syntax::NodeRole::OpenParen);
     for (auto *P : L.getParams())
       Builder.markDelayedChild(
           Builder.getDeclRange(P),
           syntax::NodeRole::ParametersAndQualifiers_parameter);
     Builder.markChildToken(L.getRParenLoc(), syntax::NodeRole::CloseParen);
     Builder.foldNode(Builder.getRange(L.getLParenLoc(), L.getEndLoc()),
-                     new (allocator()) syntax::ParametersAndQualifiers);
+                     new (allocator()) syntax::ParametersAndQualifiers, L);
     return true;
   }
 
   bool WalkUpFromFunctionProtoTypeLoc(FunctionProtoTypeLoc L) {
     if (!L.getTypePtr()->hasTrailingReturn())
       return WalkUpFromFunctionTypeLoc(L);
 
     auto TrailingReturnTokens = BuildTrailingReturn(L);
     // Finish building the node for parameters.
     Builder.markChild(TrailingReturnTokens,
                       syntax::NodeRole::ParametersAndQualifiers_trailingReturn);
     return WalkUpFromFunctionTypeLoc(L);
   }
 
   bool WalkUpFromMemberPointerTypeLoc(MemberPointerTypeLoc L) {
     auto SR = L.getLocalSourceRange();
     Builder.foldNode(Builder.getRange(SR.getBegin(), SR.getEnd()),
-                     new (allocator()) syntax::MemberPointer);
+                     new (allocator()) syntax::MemberPointer, L);
     return true;
   }
 
   // The code below is very regular, it could even be generated with some
   // preprocessor magic. We merely assign roles to the corresponding children
   // and fold resulting nodes.
   bool WalkUpFromDeclStmt(DeclStmt *S) {
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::DeclarationStatement);
+                     new (allocator()) syntax::DeclarationStatement, S);
     return true;
   }
 
   bool WalkUpFromNullStmt(NullStmt *S) {
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::EmptyStatement);
+                     new (allocator()) syntax::EmptyStatement, S);
     return true;
   }
 
   bool WalkUpFromSwitchStmt(SwitchStmt *S) {
     Builder.markChildToken(S->getSwitchLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.markStmtChild(S->getBody(), syntax::NodeRole::BodyStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::SwitchStatement);
+                     new (allocator()) syntax::SwitchStatement, S);
     return true;
   }
 
   bool WalkUpFromCaseStmt(CaseStmt *S) {
     Builder.markChildToken(S->getKeywordLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.markExprChild(S->getLHS(), syntax::NodeRole::CaseStatement_value);
     Builder.markStmtChild(S->getSubStmt(), syntax::NodeRole::BodyStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::CaseStatement);
+                     new (allocator()) syntax::CaseStatement, S);
     return true;
   }
 
   bool WalkUpFromDefaultStmt(DefaultStmt *S) {
     Builder.markChildToken(S->getKeywordLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.markStmtChild(S->getSubStmt(), syntax::NodeRole::BodyStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::DefaultStatement);
+                     new (allocator()) syntax::DefaultStatement, S);
     return true;
   }
 
   bool WalkUpFromIfStmt(IfStmt *S) {
     Builder.markChildToken(S->getIfLoc(), syntax::NodeRole::IntroducerKeyword);
     Builder.markStmtChild(S->getThen(),
                           syntax::NodeRole::IfStatement_thenStatement);
     Builder.markChildToken(S->getElseLoc(),
                            syntax::NodeRole::IfStatement_elseKeyword);
     Builder.markStmtChild(S->getElse(),
                           syntax::NodeRole::IfStatement_elseStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::IfStatement);
+                     new (allocator()) syntax::IfStatement, S);
     return true;
   }
 
   bool WalkUpFromForStmt(ForStmt *S) {
     Builder.markChildToken(S->getForLoc(), syntax::NodeRole::IntroducerKeyword);
     Builder.markStmtChild(S->getBody(), syntax::NodeRole::BodyStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::ForStatement);
+                     new (allocator()) syntax::ForStatement, S);
     return true;
   }
 
   bool WalkUpFromWhileStmt(WhileStmt *S) {
     Builder.markChildToken(S->getWhileLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.markStmtChild(S->getBody(), syntax::NodeRole::BodyStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::WhileStatement);
+                     new (allocator()) syntax::WhileStatement, S);
     return true;
   }
 
   bool WalkUpFromContinueStmt(ContinueStmt *S) {
     Builder.markChildToken(S->getContinueLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::ContinueStatement);
+                     new (allocator()) syntax::ContinueStatement, S);
     return true;
   }
 
   bool WalkUpFromBreakStmt(BreakStmt *S) {
     Builder.markChildToken(S->getBreakLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::BreakStatement);
+                     new (allocator()) syntax::BreakStatement, S);
     return true;
   }
 
   bool WalkUpFromReturnStmt(ReturnStmt *S) {
     Builder.markChildToken(S->getReturnLoc(),
                            syntax::NodeRole::IntroducerKeyword);
     Builder.markExprChild(S->getRetValue(),
                           syntax::NodeRole::ReturnStatement_value);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::ReturnStatement);
+                     new (allocator()) syntax::ReturnStatement, S);
     return true;
   }
 
   bool WalkUpFromCXXForRangeStmt(CXXForRangeStmt *S) {
     Builder.markChildToken(S->getForLoc(), syntax::NodeRole::IntroducerKeyword);
     Builder.markStmtChild(S->getBody(), syntax::NodeRole::BodyStatement);
     Builder.foldNode(Builder.getStmtRange(S),
-                     new (allocator()) syntax::RangeBasedForStatement);
+                     new (allocator()) syntax::RangeBasedForStatement, S);
     return true;
   }
 
   bool WalkUpFromEmptyDecl(EmptyDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::EmptyDeclaration);
+                     new (allocator()) syntax::EmptyDeclaration, S);
     return true;
   }
 
   bool WalkUpFromStaticAssertDecl(StaticAssertDecl *S) {
     Builder.markExprChild(S->getAssertExpr(),
                           syntax::NodeRole::StaticAssertDeclaration_condition);
     Builder.markExprChild(S->getMessage(),
                           syntax::NodeRole::StaticAssertDeclaration_message);
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::StaticAssertDeclaration);
+                     new (allocator()) syntax::StaticAssertDeclaration, S);
     return true;
   }
 
   bool WalkUpFromLinkageSpecDecl(LinkageSpecDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::LinkageSpecificationDeclaration);
+                     new (allocator()) syntax::LinkageSpecificationDeclaration,
+                     S);
     return true;
   }
 
   bool WalkUpFromNamespaceAliasDecl(NamespaceAliasDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::NamespaceAliasDefinition);
+                     new (allocator()) syntax::NamespaceAliasDefinition, S);
     return true;
   }
 
   bool WalkUpFromUsingDirectiveDecl(UsingDirectiveDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::UsingNamespaceDirective);
+                     new (allocator()) syntax::UsingNamespaceDirective, S);
     return true;
   }
 
   bool WalkUpFromUsingDecl(UsingDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::UsingDeclaration);
+                     new (allocator()) syntax::UsingDeclaration, S);
     return true;
   }
 
   bool WalkUpFromUnresolvedUsingValueDecl(UnresolvedUsingValueDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::UsingDeclaration);
+                     new (allocator()) syntax::UsingDeclaration, S);
     return true;
   }
 
   bool WalkUpFromUnresolvedUsingTypenameDecl(UnresolvedUsingTypenameDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::UsingDeclaration);
+                     new (allocator()) syntax::UsingDeclaration, S);
     return true;
   }
 
   bool WalkUpFromTypeAliasDecl(TypeAliasDecl *S) {
     Builder.foldNode(Builder.getDeclRange(S),
-                     new (allocator()) syntax::TypeAliasDeclaration);
+                     new (allocator()) syntax::TypeAliasDeclaration, S);
     return true;
   }
 
 private:
   /// Returns the range of the built node.
-  llvm::ArrayRef<syntax::Token> BuildTrailingReturn(FunctionProtoTypeLoc L) {
+  syntax::TrailingReturnType *BuildTrailingReturn(FunctionProtoTypeLoc L) {
     assert(L.getTypePtr()->hasTrailingReturn());
 
     auto ReturnedType = L.getReturnLoc();
     // Build node for the declarator, if any.
     auto ReturnDeclaratorRange =
         getDeclaratorRange(this->Builder.sourceManager(), ReturnedType,
                            /*Name=*/SourceLocation(),
                            /*Initializer=*/SourceLocation());
-    llvm::ArrayRef<syntax::Token> ReturnDeclaratorTokens;
+    syntax::SimpleDeclarator *ReturnDeclarator = nullptr;
     if (ReturnDeclaratorRange.isValid()) {
-      ReturnDeclaratorTokens = Builder.getRange(
-          ReturnDeclaratorRange.getBegin(), ReturnDeclaratorRange.getEnd());
-      Builder.foldNode(ReturnDeclaratorTokens,
-                       new (allocator()) syntax::SimpleDeclarator);
+      ReturnDeclarator = new (allocator()) syntax::SimpleDeclarator;
+      Builder.foldNode(Builder.getRange(ReturnDeclaratorRange.getBegin(),
+                                        ReturnDeclaratorRange.getEnd()),
+                       ReturnDeclarator, nullptr);
     }
 
     // Build node for trailing return type.
     auto Return =
         Builder.getRange(ReturnedType.getBeginLoc(), ReturnedType.getEndLoc());
     auto *Arrow = Return.begin() - 1;
     assert(Arrow->kind() == tok::arrow);
     auto Tokens = llvm::makeArrayRef(Arrow, Return.end());
     Builder.markChildToken(Arrow, syntax::NodeRole::TrailingReturnType_arrow);
-    if (!ReturnDeclaratorTokens.empty())
-      Builder.markChild(ReturnDeclaratorTokens,
+    if (ReturnDeclarator)
+      Builder.markChild(ReturnDeclarator,
                         syntax::NodeRole::TrailingReturnType_declarator);
-    Builder.foldNode(Tokens, new (allocator()) syntax::TrailingReturnType);
-    return Tokens;
+    auto *R = new (allocator()) syntax::TrailingReturnType;
+    Builder.foldNode(Tokens, R, nullptr);
+    return R;
   }
 
-  void
-  foldExplicitTemplateInstantiation(ArrayRef<syntax::Token> Range,
-                                    const syntax::Token *ExternKW,
-                                    const syntax::Token *TemplateKW,
-                                    ArrayRef<syntax::Token> InnerDeclaration) {
+  void foldExplicitTemplateInstantiation(
+      ArrayRef<syntax::Token> Range, const syntax::Token *ExternKW,
+      syntax::SimpleDeclaration *InnerDeclaration) {
     assert(!ExternKW || ExternKW->kind() == tok::kw_extern);
     assert(TemplateKW && TemplateKW->kind() == tok::kw_template);
     Builder.markChildToken(
         ExternKW,
         syntax::NodeRole::ExplicitTemplateInstantiation_externKeyword);
     Builder.markChildToken(TemplateKW, syntax::NodeRole::IntroducerKeyword);
     Builder.markChild(
         InnerDeclaration,
         syntax::NodeRole::ExplicitTemplateInstantiation_declaration);
     Builder.foldNode(Range,
-                     new (allocator()) syntax::ExplicitTemplateInstantiation);
+                     new (allocator()) syntax::ExplicitTemplateInstantiation,
+                     nullptr);
   }
 
-  void foldTemplateDeclaration(ArrayRef<syntax::Token> Range,
-                               const syntax::Token *TemplateKW,
-                               ArrayRef<syntax::Token> TemplatedDeclaration) {
+  syntax::TemplateDeclaration *foldTemplateDeclaration(
+      ArrayRef<syntax::Token> Range, const syntax::Token *TemplateKW,
+      ArrayRef<syntax::Token> TemplatedDeclaration, Decl *From) {
     assert(TemplateKW && TemplateKW->kind() == tok::kw_template);
     Builder.markChildToken(TemplateKW, syntax::NodeRole::IntroducerKeyword);
     Builder.markMaybeDelayedChild(
         TemplatedDeclaration,
         syntax::NodeRole::TemplateDeclaration_declaration);
-    Builder.foldNode(Range, new (allocator()) syntax::TemplateDeclaration);
+
+    auto *N = new (allocator()) syntax::TemplateDeclaration;
+    Builder.foldNode(Range, N, From);
+    return N;
   }
 
   /// A small helper to save some typing.
   llvm::BumpPtrAllocator &allocator() { return Builder.allocator(); }
 
   syntax::TreeBuilder &Builder;
   const LangOptions &LangOpts;
 };
 } // namespace
 
-void syntax::TreeBuilder::foldNode(llvm::ArrayRef<syntax::Token> Range,
-                                   syntax::Tree *New) {
-  Pending.foldChildren(Arena, Range, New);
-}
-
 void syntax::TreeBuilder::noticeDeclRange(llvm::ArrayRef<syntax::Token> Range) {
   if (Pending.extendDelayedFold(Range))
     return;
   Pending.foldChildrenDelayed(Range,
                               new (allocator()) syntax::SimpleDeclaration);
 }
 
 void syntax::TreeBuilder::noticeDeclWithoutSemicolon(Decl *D) {
   DeclsWithoutSemicolons.insert(D);
 }
 
 void syntax::TreeBuilder::markChildToken(SourceLocation Loc, NodeRole Role) {
   if (Loc.isInvalid())
     return;
   Pending.assignRole(*findToken(Loc), Role);
 }
 
 void syntax::TreeBuilder::markChildToken(const syntax::Token *T, NodeRole R) {
   if (!T)
     return;
   Pending.assignRole(*T, R);
 }
 
-void syntax::TreeBuilder::markChild(llvm::ArrayRef<syntax::Token> Range,
-                                    NodeRole R) {
-  Pending.assignRole(Range, R);
+void syntax::TreeBuilder::markChild(syntax::Node *N, NodeRole R) {
+  assert(N);
+  setRole(N, R);
+}
+
+void syntax::TreeBuilder::markChild(ASTPtr N, NodeRole R) {
+  auto *SN = Mapping.find(N);
+  assert(SN != nullptr);
+  setRole(SN, R);
 }
 
 void syntax::TreeBuilder::markDelayedChild(llvm::ArrayRef<syntax::Token> Range,
                                            NodeRole R) {
   Pending.assignRoleDelayed(Range, R);
 }
 
 void syntax::TreeBuilder::markMaybeDelayedChild(
     llvm::ArrayRef<syntax::Token> Range, NodeRole R) {
   Pending.assignRoleMaybeDelayed(Range, R);
 }
 
 void syntax::TreeBuilder::markStmtChild(Stmt *Child, NodeRole Role) {
   if (!Child)
     return;
 
-  auto Range = getStmtRange(Child);
+  syntax::Tree *ChildNode = Mapping.find(Child);
+  assert(ChildNode != nullptr);
+
   // This is an expression in a statement position, consume the trailing
   // semicolon and form an 'ExpressionStatement' node.
   if (auto *E = dyn_cast<Expr>(Child)) {
-    Pending.assignRole(getExprRange(E),
-                       NodeRole::ExpressionStatement_expression);
-    // (!) 'getRange(Stmt)' ensures this already covers a trailing semicolon.
-    Pending.foldChildren(Arena, Range,
-                         new (allocator()) syntax::ExpressionStatement);
+    setRole(ChildNode, NodeRole::ExpressionStatement_expression);
+    ChildNode = new (allocator()) syntax::ExpressionStatement;
+    // (!) 'getStmtRange()' ensures this covers a trailing semicolon.
+    Pending.foldChildren(Arena, getStmtRange(Child), ChildNode);
   }
-  Pending.assignRole(Range, Role);
+  setRole(ChildNode, Role);
 }
 
 void syntax::TreeBuilder::markExprChild(Expr *Child, NodeRole Role) {
   if (!Child)
     return;
+  Child = Child->IgnoreImplicit();
 
-  Pending.assignRole(getExprRange(Child), Role);
+  syntax::Tree *ChildNode = Mapping.find(Child);
+  assert(ChildNode != nullptr);
+  setRole(ChildNode, Role);
 }
 
 const syntax::Token *syntax::TreeBuilder::findToken(SourceLocation L) const {
   if (L.isInvalid())
     return nullptr;
   auto It = LocationToToken.find(L.getRawEncoding());
   assert(It != LocationToToken.end());
   return It->second;
 }
 
 syntax::TranslationUnit *
 syntax::buildSyntaxTree(Arena &A, const TranslationUnitDecl &TU) {
   TreeBuilder Builder(A);
   BuildTreeVisitor(TU.getASTContext(), Builder).TraverseAST(TU.getASTContext());
   return std::move(Builder).finalize();
 }

clang/lib/Tooling/Syntax/Tree.cpp

: Parent(nullptr), NextSibling(nullptr), Kind(static_cast<unsigned>(Kind)),
       Role(static_cast<unsigned>(NodeRole::Detached)), Original(false),
       CanModify(false) {}
 
 bool syntax::Node::isDetached() const { return role() == NodeRole::Detached; }
 
 bool syntax::Tree::classof(const Node *N) { return N->kind() > NodeKind::Leaf; }
 
 void syntax::Tree::prependChildLowLevel(Node *Child, NodeRole Role) {
-  assert(Child->Parent == nullptr);
-  assert(Child->NextSibling == nullptr);
   assert(Child->role() == NodeRole::Detached);
   assert(Role != NodeRole::Detached);
 
+  Child->Role = static_cast<unsigned>(Role);
+  prependChildLowLevel(Child);
+}
+
+void syntax::Tree::prependChildLowLevel(Node *Child) {
+  assert(Child->Parent == nullptr);
+  assert(Child->NextSibling == nullptr);
+  assert(Child->role() != NodeRole::Detached);
+
   Child->Parent = this;
   Child->NextSibling = this->FirstChild;
-  Child->Role = static_cast<unsigned>(Role);
   this->FirstChild = Child;
 }
 
 void syntax::Tree::replaceChildRangeLowLevel(Node *BeforeBegin, Node *End,
                                              Node *New) {
   assert(!BeforeBegin || BeforeBegin->Parent == this);
 
 #ifndef NDEBUG