[Syntax] Build mapping from AST to syntax tree nodes

[hlopko]

Summary

Copy of https://reviews.llvm.org/D72446, submitting with Ilya's permission.

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.

[hlopko]

Addressing comments.

[hlopko]

Reformat

[gribozavr2]

clang/lib/Tooling/Syntax/BuildTree.cpp

173

The comment is not needed anymore.

489

Could you add a private setter that performs this cast? (The cast is repeated at least 3 times in this patch.)

571

It might make sense to add a helper Builder.getRange(SourceRange) and simplify these calls to Builder.getRange(something.getBegin(), something.getEnd()) throughout the patch.

976

Add a Decl *From parameter and pass it through to Builder.foldNode() below?

[hlopko]

Address comments

clang/lib/Tooling/Syntax/BuildTree.cpp

571

Done and refactored all the calls that made sense to use the SourceRange method. I didn't modify call like Builder.getRange(L.getLBracketLoc(), L.getRBracketLoc()).

976

Done, but to get rid of all nullptr parents in BuildTree.cpp we'd have to implement support for Types in the AST mapping. Let's not do that in this patch.

[gribozavr2]

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

129

setRole() (in new code).

clang/lib/Tooling/Syntax/BuildTree.cpp

266

Don't repeat the function name in comments.

/// Finds the syntax tokens that correspond to the provided \c SourceRange.

266

Ditto.

976

For consistency with, for example, foldTemplateDeclaration, I think Decl *From should be the last parameter.

[hlopko]

Address comments

[hlopko]

PTAL.

clang/lib/Tooling/Syntax/BuildTree.cpp

173

Clang tidy complains when I remove it :/

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

private:
   friend class Tree;
   // TreeBuilder is allowed to set the Original and CanModify flags.
   friend class TreeBuilder;
   // MutationsImpl sets roles and CanModify flag.
   friend class MutationsImpl;
   // FactoryImpl sets CanModify flag.
   friend class FactoryImpl;
 
   Tree *Parent;
   Node *NextSibling;
   unsigned Kind : 16;
   unsigned Role : 8;
   unsigned Original : 1;
   unsigned CanModify : 1;
 };
 
 /// A leaf node points to a single token inside the expanded token stream.

 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() && {
     auto Tokens = Arena.tokenBuffer().expandedTokens();
     assert(!Tokens.empty());
     assert(Tokens.back().kind() == tok::eof);
 
     // Build the root of the tree, consuming all the children.
     Pending.foldChildren(Arena, Tokens.drop_back(),
                          new (Arena.allocator()) syntax::TranslationUnit);
 
     auto *TU = cast<syntax::TranslationUnit>(std::move(Pending).finalize());
     TU->assertInvariantsRecursive();
     return TU;
   }
 
   /// Finds a token starting at \p L. The token must exist if \p L is valid.
   const syntax::Token *findToken(SourceLocation L) const;
 
   /// getRange() finds the syntax tokens corresponding to the passed source
   /// locations.
   /// \p First is the start position of the first token and \p Last is the start
   /// position of the last token.
   llvm::ArrayRef<syntax::Token> getRange(SourceLocation First,
                                          SourceLocation Last) const {
     assert(First.isValid());
     assert(Last.isValid());
     assert(First == Last ||
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 BeginFolds = 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 += std::string(llvm::formatv(
-            "- '{0}' covers '{1}'+{2} tokens\n", It->second.Node->kind(),
+            "- '{0}' covers '{1}'+{2} tokens\n", 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());
     const 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