[syntax] Introduce a TokenManager interface.

[hokein]

Summary

The TokenManager defines Token interfaces for the clang syntax-tree. This is level of
abstraction that the syntax-tree should use to operate on Tokens.

It decouples the syntax-tree from a particular token implementation (TokenBuffer
previously). This enables us to use a different underlying token implementation
for the syntax Leaf node -- in clang pseudoparser, we want to produce a
syntax-tree with its own pseudo::Token rather than syntax::Token.

[sammccall]

As discussed offline this has some problems:

• putting virtual methods on BaseToken gives it a vtable, which makes it (and syntax::Token) large
• being able to use ArrayRef<syntax::Token> but not ArrayRef<BaseToken> is a bit weird
• unusual uses of inheritance can be hard to reason about

We suggested rather having Leaf store an opaque ID.
Callers who know what kind of tokens are in use can use this to associate with the original token.
For generic use (e.g dump()), we can have a TokenManager interface which provides the common operations (like getText()). This generalizes where SourceManager is needed today.

[hokein]

Revised to the TokenManager approach:

• Inroduce a Base Token class (TokenManager) for syntax-tree, the motivation is to allow using different underlying token implementation in syntax-tree
• Decouple the syntax-tree from the TokenBuffer:
  • syntax-tree structure (Tree.h) doesn't depend on the TokenBuffer, SourceManager Source location etc, it communicates with TokenManager interfaces;
  • syntax-tree Arena is simpler, the token-managing responsiblity is transferred to TokenManager;
  • in SyntaxTree directory, we implement a TokenBuffer-based SyntaxTokenManager, which mangues all token-related stuff
  • For the mutation/replacement computation APIs, currently they only work on a TokenBuffer-based token manager. Asssertion will be raised if it is not satisfied. It is an NFC change

[hokein]

I'm quite happy about the interfaces now, it should be in a good shape for the API review (would be nice to get some initial feedback before I do further cleanup).

[ilya-biryukov]

clang/include/clang/Tooling/Syntax/TokenManager.h

24

NIT: maybe explain the TokenManager concept here, the comment seems to be a leftover from the previous revision.

clang/lib/Tooling/Syntax/Tree.cpp

271

Maybe add TokenManager::getKind(Key) right away and remove this FIXME.
This should as simple as cast<syntax::Token>(T)->Kind, right? Or am I missing some complications?

[ilya-biryukov]

clang/include/clang/Tooling/Syntax/TokenManager.h

34

I have just realized that we were discussing having opaque index as a key, but there may also be tokens in the syntax tree that are not backed by the TokenBuffer.
Those can be synthesized, e.g. imagine someone wants to change !(A && B) to !A || !B. They will need to synthesize at least the || token as it's not in the source code. There is a way to do this now and it prohibits the use of an index to the TokenBuffer.

So having the opaque pointer is probably ok for now, it should enable the pseudo-parser to build syntax trees.
We might want to add an operation to synthesize tokens into the TokenManager at some point, but that's a discussion for another day.

[hokein]

clang/include/clang/Tooling/Syntax/TokenManager.h

34

Those can be synthesized, e.g. imagine someone wants to change !(A && B) to !A || !B. They will need to synthesize at least the || token as it's not in the source code. There is a way to do this now and it prohibits the use of an index to the TokenBuffer.

Yes, this is the exact reason why the Key is an opaque pointer, my first attempt was to use an index integer, but failed -- we already have some APIs doing this stuff (see createLeaf in BuildTree.h), the token can be a synthesized token backed up by the SourceManager...

Personally, I don't like the Key to be an opaque pointer as well, but considering the effort, it seems to be the best approach so far -- it enables the pseudoparser to build syntax trees with a different Token implementation while keeping the rest syntax stuff unchanged.

We might want to add an operation to synthesize tokens into the TokenManager at some point, but that's a discussion for another day.

Agree, we will encounter this in the future, but we're still far away from there (the layering mutation/syntax-tree is not perfect at the moment, mutation still depends on the TokenBuffer). And our initial application of syntax-tree in pseudoparser focuses on the read use-case, we should be fine now.

clang/lib/Tooling/Syntax/Tree.cpp

271

Yeah, the main problem is that we don't have the TokenManager object in the syntax::Node, we somehow need to pass it (e.g. a function parameter), which seems a heavy change. I'm not sure it is worth for this small assert.

[ilya-biryukov]

clang/lib/Tooling/Syntax/BuildTree.cpp

571–573

Could we accept a TokenBuffer here directly?
If TokenManager is needed, we can use it directly in the arena.

627–628

Same suggestion here: accept TokenBuffer instead of TokenManager

clang/lib/Tooling/Syntax/Tree.cpp

271

That makes sense. WDYT about the alternative fix: to pass ̀TokenManager` to assertInvariants?
Not necessary to do it now, just thinking about changing the FIXME

clang/unittests/Tooling/Syntax/TreeTestBase.cpp

117

NIT: it´s not breaking anything now, but I suggest putting SyntaxTokenManager after TokenBuffer.
The reason is that it´s the right destruction order, TokenManager has references to TokenBuffer, so it could potentially access it in destructor some time in the future (e.g. imagine asserting something on tokens).

Not that it actually breaks today, but seems like a potential surprising bug in the future if we happen to refactor code in a certain way.

[sammccall]

clang/include/clang/Tooling/Syntax/TokenManager.h

10

It's important that we have comments explaining what the concept is *now*, rather than how it changes the code structure from the previous state (decouples tokenbuffer, enables pseudoparser).

(I think this comment is actually fine, but be careful when writing the class comment for TokenManager)

34

Consider uintptr_t instead, which more naturally supports both approaches. (Stashing an integer in a void* is incredibly weird)

38

I wouldn't want to prejudge this: this is a very basic attribute similar to kind/role, and we may want to store it in Leaf and avoid the virtual stuff.

There's certainly enough space, e.g. of the current 16-bit kind use the top bit to denote leaf-or-not and the bottom 15 bits to store kind-or-tokenkind.

clang/include/clang/Tooling/Syntax/Tokens.h

463

conceptually this is just "TokenBuffer implements TokenManager"

The main reason I can see not to actually write that is to avoid the dependency from TokenBuffer (tokens library) to TokenManager (syntax library). But here you've added that dependency anyway.

So I think we'd be better either with TokenBuffer : TokenManager or moving this class to its own header.

465

no need to take SourceManager, TokenBuffer already includes it.
LangOpts is unused here.

475

Empty string seems like the correct return value here to me.
If you want a special case for dump, I think that belongs in dump().

If this is because we currently provide no way to get the token kind, then this should be a FIXME

487

just tokenBuffer(), consistent with TokenBuffer::sourceManager()

490

manager

504

aren't we trying to store these on the syntax arena?
We never do lookups into this map, maybe lexbuffer just allocates storage on the arena('s allocator) instead of using this map?

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

156

why is this comment removed?

clang/lib/Tooling/Syntax/Tree.cpp

271

per my comment above: Leaf can store the tok::Kind directly and I think it's appropriate to do so.
But maybe fiddly enough that it's worth deferring for one patch

[ilya-biryukov]

clang/include/clang/Tooling/Syntax/Tokens.h

463

I would argue they should be separate concepts.

• TokenBuffer is about storing tokens and mapping between expanded and spelled token streams.
• SyntaxTokenManager is about implementing relevant certain operations on syntax::Token and hiding the actual token type.

Conceptually those things are different and decoupling them allows for more flexibility and allows reasoning about them independently. In particular, one could imagine having a SyntaxTokenManager without a TokenBuffer if we do not need to deal with two streams of tokens.

I would suggest keeping SyntaxTokenManager as a separate class.

[hokein]

address review comments.

[hokein]

clang/include/clang/Tooling/Syntax/TokenManager.h

38

This sounds good. Adjust the FIXME.

clang/include/clang/Tooling/Syntax/Tokens.h

463

I would suggest keeping SyntaxTokenManager as a separate class.

+1. Moved to a separate file.

465

no need to take SourceManager, TokenBuffer already includes it.

The SourceMgr can be mutated by the class (it stores the underlying tokens for ExtraTokens), while the SourceManager in TokenBuffer is immutable.

LangOpts is unused here.

It is used in SyntaxTokenManager::lexBuffer.

475

Yeah, this special case is for the Leaf node dump. Added a FIXME.

(I even double whether we need this special case at all, do we really want to build a Leaf node for eof token?)

504

aren't we trying to store these on the syntax arena?

We could do that, but I'd try to avoid that. Now the allocator of syntax::Arena is a storage for syntax-tree nodes only. Allocation for token-related stuff is on the SyntaxTokenManager.

We never do lookups into this map, maybe lexbuffer just allocates storage on the arena('s allocator) instead of using this map?

This map is moved from the syntax::Arena.

You're right, there is no usage of the key at the moment, and the only use-case is to create a syntax leaf node that not backed by the source code (for refactoring usecase), it is unclear whether we will use it in the future. I will keep it as-is (this is not the scope of this patch).

clang/unittests/Tooling/Syntax/TreeTestBase.cpp

117

good point!

[sammccall]

clang/include/clang/Tooling/Syntax/Nodes.h

23–24

I think Token and TokenKinds are also no longer needed?

clang/include/clang/Tooling/Syntax/SyntaxTokenManager.h

20 ↗

(On Diff #443173)

I don't think "syntax" in "syntax token manager" is particularly disambiguating here, both TokenBuffer and TokenManager are part of syntax, so it's not clear what it refers to (and it doens't have any obvious plain-english meaning).

Maybe some combination like TokenBufferTokenManager or BufferTokenManager.
In fact I think best is TokenBuffer::TokenManager, still defined in a separate header, though I'm not sure if you think that's too weird.

clang/include/clang/Tooling/Syntax/TokenManager.h

15

unclear: different from what? it's not clear what "enables" means if there's no default. Maybe replace the sentence with:

"For example, a TokenBuffer captured from a clang parse may track macro expansions and associate tokens with clang's SourceManager, while a pseudo-parser would use a flat array of raw-lexed tokens in memory."

39

This is not a useful comment, either remove it or add more content to make it useful.

In particular the guarantees (or lack thereof) of exactly what this text is would be helpful. (This is some source code that would produce this token, though it may differ from exactly what was spelled in the file when preprocessing is involved)?

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

9

this is an implementation detail.
At a high level, leaf nodes correspond to tokens. I'd just delete "expanded" from the original comment

41–42

this is not a helpful comment, just remove it?

clang/lib/Tooling/Syntax/BuildTree.cpp

370

need changes to the public API to make this cast valid

clang/lib/Tooling/Syntax/ComputeReplacements.cpp

95–96

Need a change to the public interface to guarantee this cast will succeed.
The cleanest would be just to take the SyntaxTokenManager as a param (moving the cast to the call site - I don't think Arena is otherwise used for anything).

Failing that we at least need to update the contract

[hokein]

address comments.

[hokein]

more update

clang/include/clang/Tooling/Syntax/SyntaxTokenManager.h

20 ↗

(On Diff #443173)

Renamed to TokenBufferTokenManager.

BufferTokenManager name is short, but it has BufferToken prefix, which seems confusing with TokenBuffer. TokenBuffer::TokenManager is weird to me, and doesn't reflect the layering IMO

clang/lib/Tooling/Syntax/BuildTree.cpp

370

Are you suggesting to change all public APIs where there is such a cast usage?

If yes, this seems not a trivial change, I think at least we will change all APIs (buildSyntaxTree, createLeaf, createTree, deepCopyExpandingMacros) in BuildTree.h (by adding a new TokenBufferTokenManager parameter).
And the Arena probably doesn't need to have a TokenManager field (it could be simplified as a single BumpPtrAllocator), as the TokenManager is passed in parallel with the Arena.

I'm happy to do the change, but IMO, the current version doesn't seem too bad for me (if we pass an non-SyntaxTokenManager, it will trigger an assertion in debug mode).

clang/lib/Tooling/Syntax/ComputeReplacements.cpp

95–96

Done, this is a trivial change.

[sammccall]

clang/lib/Tooling/Syntax/BuildTree.cpp

370

Are you suggesting to change all public APIs where there is such a cast usage?

Yes, at the very least they should document the requirement ("this arena must use a TokenBuffer").
An unchecked downcast with no indication on the public API that a specific subclass is required just looks like a bug.

[hokein]

remove all TokenBufferTokenManager cast usage, make it as a contract in the APIs.

[hokein]

update the API changes in clangd part.

[DavidSpickett]

FYI, after this change I get:

/home/david.spickett/llvm-project/clang/include/clang/Tooling/Syntax/TokenBufferTokenManager.h:20:7: warning: 'clang::syntax::TokenBufferTokenManager' has virtual functions but non-virtual destructor [-Wnon-virtual-dtor]
class TokenBufferTokenManager : public TokenManager {
      ^

[hokein]

In D128411#3654452, @DavidSpickett wrote:

FYI, after this change I get:

/home/david.spickett/llvm-project/clang/include/clang/Tooling/Syntax/TokenBufferTokenManager.h:20:7: warning: 'clang::syntax::TokenBufferTokenManager' has virtual functions but non-virtual destructor [-Wnon-virtual-dtor]
class TokenBufferTokenManager : public TokenManager {
      ^

Sorry, it is fixed in 30c2406e270cc5dab8da813ce5c54e4bb8c40e49.

clang-tools-extra/clangd/SemanticSelection.cpp

 #include "Protocol.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Tooling/Syntax/BuildTree.h"
 #include "clang/Tooling/Syntax/Nodes.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 #include "clang/Tooling/Syntax/Tree.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
 #include <queue>
 #include <vector>
 
 namespace clang {
llvm::Optional<FoldingRange> toFoldingRange(SourceRange SR,
   FoldingRange Range;
   Range.startCharacter = SM.getColumnNumber(Begin.first, Begin.second) - 1;
   Range.startLine = SM.getLineNumber(Begin.first, Begin.second) - 1;
   Range.endCharacter = SM.getColumnNumber(End.first, End.second) - 1;
   Range.endLine = SM.getLineNumber(End.first, End.second) - 1;
   return Range;
 }
 
-llvm::Optional<FoldingRange> extractFoldingRange(const syntax::Node *Node,
-                                                 const SourceManager &SM) {
+llvm::Optional<FoldingRange>
+extractFoldingRange(const syntax::Node *Node,
+                    const syntax::TokenBufferTokenManager &TM) {
   if (const auto *Stmt = dyn_cast<syntax::CompoundStatement>(Node)) {
     const auto *LBrace = cast_or_null<syntax::Leaf>(
         Stmt->findChild(syntax::NodeRole::OpenParen));
     // FIXME(kirillbobyrev): This should find the last child. Compound
     // statements have only one pair of braces so this is valid but for other
     // node kinds it might not be correct.
     const auto *RBrace = cast_or_null<syntax::Leaf>(
         Stmt->findChild(syntax::NodeRole::CloseParen));
     if (!LBrace || !RBrace)
       return llvm::None;
     // Fold the entire range within braces, including whitespace.
-    const SourceLocation LBraceLocInfo = LBrace->getToken()->endLocation(),
-                         RBraceLocInfo = RBrace->getToken()->location();
-    auto Range = toFoldingRange(SourceRange(LBraceLocInfo, RBraceLocInfo), SM);
+    const SourceLocation LBraceLocInfo =
+                             TM.getToken(LBrace->getTokenKey())->endLocation(),
+                         RBraceLocInfo =
+                             TM.getToken(RBrace->getTokenKey())->location();
+    auto Range = toFoldingRange(SourceRange(LBraceLocInfo, RBraceLocInfo),
+                                TM.sourceManager());
     // Do not generate folding range for compound statements without any
     // nodes and newlines.
     if (Range && Range->startLine != Range->endLine)
       return Range;
   }
   return llvm::None;
 }
 
 // Traverse the tree and collect folding ranges along the way.
-std::vector<FoldingRange> collectFoldingRanges(const syntax::Node *Root,
-                                               const SourceManager &SM) {
+std::vector<FoldingRange>
+collectFoldingRanges(const syntax::Node *Root,
+                     const syntax::TokenBufferTokenManager &TM) {
   std::queue<const syntax::Node *> Nodes;
   Nodes.push(Root);
   std::vector<FoldingRange> Result;
   while (!Nodes.empty()) {
     const syntax::Node *Node = Nodes.front();
     Nodes.pop();
-    const auto Range = extractFoldingRange(Node, SM);
+    const auto Range = extractFoldingRange(Node, TM);
     if (Range)
       Result.push_back(*Range);
     if (const auto *T = dyn_cast<syntax::Tree>(Node))
       for (const auto *NextNode = T->getFirstChild(); NextNode;
            NextNode = NextNode->getNextSibling())
         Nodes.push(NextNode);
   }
   return Result;
llvm::Expected<SelectionRange> getSemanticRanges(ParsedAST &AST, Position Pos) {
   return std::move(Head);
 }
 
 // FIXME(kirillbobyrev): Collect comments, PP conditional regions, includes and
 // other code regions (e.g. public/private/protected sections of classes,
 // control flow statement bodies).
 // Related issue: https://github.com/clangd/clangd/issues/310
 llvm::Expected<std::vector<FoldingRange>> getFoldingRanges(ParsedAST &AST) {
-  syntax::Arena A(AST.getSourceManager(), AST.getLangOpts(), AST.getTokens());
-  const auto *SyntaxTree = syntax::buildSyntaxTree(A, AST.getASTContext());
-  return collectFoldingRanges(SyntaxTree, AST.getSourceManager());
+  syntax::Arena A;
+  syntax::TokenBufferTokenManager TM(AST.getTokens(), AST.getLangOpts(),
+                                     AST.getSourceManager());
+  const auto *SyntaxTree = syntax::buildSyntaxTree(A, TM, AST.getASTContext());
+  return collectFoldingRanges(SyntaxTree, TM);
 }
 
 } // namespace clangd
 } // namespace clang

clang/include/clang/Tooling/Syntax/BuildTree.h

 //===- BuildTree.h - build syntax trees -----------------------*- C++ -*-=====//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 // Functions to construct a syntax tree from an AST.
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLING_SYNTAX_BUILDTREE_H
 #define LLVM_CLANG_TOOLING_SYNTAX_BUILDTREE_H
 
 #include "clang/AST/Decl.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Tooling/Syntax/Nodes.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 #include "clang/Tooling/Syntax/Tree.h"
 
 namespace clang {
 namespace syntax {
 
 /// Build a syntax tree for the main file.
 /// This usually covers the whole TranslationUnitDecl, but can be restricted by
 /// the ASTContext's traversal scope.
-syntax::TranslationUnit *buildSyntaxTree(Arena &A, ASTContext &Context);
+syntax::TranslationUnit *
+buildSyntaxTree(Arena &A, TokenBufferTokenManager &TBTM, ASTContext &Context);
 
 // Create syntax trees from subtrees not backed by the source code.
 
 // Synthesis of Leafs
 /// Create `Leaf` from token with `Spelling` and assert it has the desired
 /// `TokenKind`.
-syntax::Leaf *createLeaf(syntax::Arena &A, tok::TokenKind K,
-                         StringRef Spelling);
+syntax::Leaf *createLeaf(syntax::Arena &A, TokenBufferTokenManager &TBTM,
+                         tok::TokenKind K, StringRef Spelling);
 
 /// Infer the token spelling from its `TokenKind`, then create `Leaf` from
 /// this token
-syntax::Leaf *createLeaf(syntax::Arena &A, tok::TokenKind K);
+syntax::Leaf *createLeaf(syntax::Arena &A, TokenBufferTokenManager &TBTM,
+                         tok::TokenKind K);
 
 // Synthesis of Trees
 /// Creates the concrete syntax node according to the specified `NodeKind` `K`.
 /// Returns it as a pointer to the base class `Tree`.
 syntax::Tree *
 createTree(syntax::Arena &A,
            ArrayRef<std::pair<syntax::Node *, syntax::NodeRole>> Children,
            syntax::NodeKind K);
 
 // Synthesis of Syntax Nodes
-syntax::EmptyStatement *createEmptyStatement(syntax::Arena &A);
+syntax::EmptyStatement *createEmptyStatement(syntax::Arena &A,
+                                             TokenBufferTokenManager &TBTM);
 
 /// Creates a completely independent copy of `N` with its macros expanded.
 ///
 /// The copy is:
 /// * Detached, i.e. `Parent == NextSibling == nullptr` and
 /// `Role == Detached`.
 /// * Synthesized, i.e. `Original == false`.
-syntax::Node *deepCopyExpandingMacros(syntax::Arena &A, const syntax::Node *N);
+syntax::Node *deepCopyExpandingMacros(syntax::Arena &A,
+                                      TokenBufferTokenManager &TBTM,
+                                      const syntax::Node *N);
 } // namespace syntax
 } // namespace clang
 #endif

clang/include/clang/Tooling/Syntax/Mutations.h

 //===- Mutations.h - mutate syntax trees --------------------*- C++ ---*-=====//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 // Defines high-level APIs for transforming syntax trees and producing the
 // corresponding textual replacements.
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLING_SYNTAX_MUTATIONS_H
 #define LLVM_CLANG_TOOLING_SYNTAX_MUTATIONS_H
 
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Syntax/Nodes.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 #include "clang/Tooling/Syntax/Tree.h"
 
 namespace clang {
 namespace syntax {
 
 /// Computes textual replacements required to mimic the tree modifications made
 /// to the syntax tree.
-tooling::Replacements computeReplacements(const Arena &A,
+tooling::Replacements computeReplacements(const TokenBufferTokenManager &TBTM,
                                           const syntax::TranslationUnit &TU);
 
 /// Removes a statement or replaces it with an empty statement where one is
 /// required syntactically. E.g., in the following example:
 ///     if (cond) { foo(); } else bar();
 /// One can remove `foo();` completely and to remove `bar();` we would need to
 /// replace it with an empty statement.
 /// EXPECTS: S->canModify() == true
-void removeStatement(syntax::Arena &A, syntax::Statement *S);
+void removeStatement(syntax::Arena &A, TokenBufferTokenManager &TBTM,
+                     syntax::Statement *S);
 
 } // namespace syntax
 } // namespace clang
 
 #endif

clang/include/clang/Tooling/Syntax/Nodes.h

 //   - the corresponding subnode is optional in the C++ grammar, e.g. an else
 //     branch of an if statement,
 //   - syntactic errors occurred while parsing the corresponding subnode.
 // One notable exception is "introducer" keywords, e.g. the accessor for the
 // 'if' keyword of an if statement will never return null.
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLING_SYNTAX_NODES_H
 #define LLVM_CLANG_TOOLING_SYNTAX_NODES_H
 
-#include "clang/Basic/TokenKinds.h"
-#include "clang/Lex/Token.h"
-#include "clang/Tooling/Syntax/Tokens.h"
+#include "clang/Basic/LLVM.h"
 #include "clang/Tooling/Syntax/Tree.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 namespace clang {
 namespace syntax {
 
 /// A kind of a syntax node, used for implementing casts. The ordering and

clang/include/clang/Tooling/Syntax/TokenBufferTokenManager.h

+//===- TokenBufferTokenManager.h  -----------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLING_SYNTAX_TOKEN_BUFFER_TOKEN_MANAGER_H
+#define LLVM_CLANG_TOOLING_SYNTAX_TOKEN_BUFFER_TOKEN_MANAGER_H
+
+#include "clang/Tooling/Syntax/TokenManager.h"
+#include "clang/Tooling/Syntax/Tokens.h"
+
+namespace clang {
+namespace syntax {
+
+/// A TokenBuffer-powered token manager.
+/// It tracks the underlying token buffers, source manager, etc.
+class TokenBufferTokenManager : public TokenManager {
+public:
+  TokenBufferTokenManager(const TokenBuffer &Tokens,
+                          const LangOptions &LangOpts, SourceManager &SourceMgr)
+      : Tokens(Tokens), LangOpts(LangOpts), SM(SourceMgr) {}
+
+  static bool classof(const TokenManager *N) { return N->kind() == Kind; }
+  llvm::StringLiteral kind() const override { return Kind; }
+
+  llvm::StringRef getText(Key I) const override {
+    const auto *Token = getToken(I);
+    assert(Token);
+    // Handle 'eof' separately, calling text() on it produces an empty string.
+    // FIXME: this special logic is for syntax::Leaf dump, move it when we
+    // have a direct way to retrive token kind in the syntax::Leaf.
+    if (Token->kind() == tok::eof)
+      return "<eof>";
+    return Token->text(SM);
+  }
+
+  const syntax::Token *getToken(Key I) const {
+    return reinterpret_cast<const syntax::Token *>(I);
+  }
+  SourceManager &sourceManager() { return SM; }
+  const SourceManager &sourceManager() const { return SM; }
+  const TokenBuffer &tokenBuffer() const { return Tokens; }
+
+private:
+  // This manager is powered by the TokenBuffer.
+  static constexpr llvm::StringLiteral Kind = "TokenBuffer";
+
+  /// Add \p Buffer to the underlying source manager, tokenize it and store the
+  /// resulting tokens. Used exclusively in `FactoryImpl` to materialize tokens
+  /// that were not written in user code.
+  std::pair<FileID, ArrayRef<Token>>
+  lexBuffer(std::unique_ptr<llvm::MemoryBuffer> Buffer);
+  friend class FactoryImpl;
+
+  const TokenBuffer &Tokens;
+  const LangOptions &LangOpts;
+
+  /// The underlying source manager for the ExtraTokens.
+  SourceManager &SM;
+  /// IDs and storage for additional tokenized files.
+  llvm::DenseMap<FileID, std::vector<Token>> ExtraTokens;
+};
+
+} // namespace syntax
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLING_SYNTAX_TOKEN_BUFFER_TOKEN_MANAGER_H

clang/include/clang/Tooling/Syntax/TokenManager.h

+//===- TokenManager.h - Manage Tokens for syntax-tree ------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines Token interfaces for the clang syntax-tree. This is the level of
+// abstraction that the syntax-tree uses to operate on Token.
+//
+// TokenManager decouples the syntax-tree from a particular token
+// implementation. For example, a TokenBuffer captured from a clang parser may
+// track macro expansions and associate tokens with clang's SourceManager, while
+// a clang pseudoparser would use a flat array of raw-lexed tokens in memory.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLING_SYNTAX_TOKEN_MANAGER_H
+#define LLVM_CLANG_TOOLING_SYNTAX_TOKEN_MANAGER_H
+
+#include "llvm/ADT/StringRef.h"
+#include <cstdint>
+
+namespace clang {
+namespace syntax {
+
+/// Defines interfaces for operating "Token" in the clang syntax-tree.
+class TokenManager {
+public:
+  /// Describes what the exact class kind of the TokenManager is.
+  virtual llvm::StringLiteral kind() const = 0;
+
+  /// A key to identify a specific token. The token concept depends on the
+  /// underlying implementation -- it can be a spelled token from the original
+  /// source file or an expanded token.
+  /// The syntax-tree Leaf node holds a Key.
+  using Key = uintptr_t;
+  virtual llvm::StringRef getText(Key K) const = 0;
+};
+
+} // namespace syntax
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLING_SYNTAX_TOKEN_MANAGER_H

clang/include/clang/Tooling/Syntax/Tokens.h

 #define LLVM_CLANG_TOOLING_SYNTAX_TOKENS_H
 
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Lex/Token.h"
+#include "clang/Tooling/Syntax/TokenManager.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 #include <tuple>
private:
   // FIXME: we only store macro expansions, also add directives(#pragma, etc.)
   PPExpansions Expansions;
   Preprocessor &PP;
   CollectPPExpansions *Collector;
 };
 
 } // namespace syntax
 } // namespace clang
 
 #endif

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

 //===- Tree.h - structure of the syntax tree ------------------*- C++ -*-=====//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 // Defines the basic structure of the syntax tree. There are two kinds of nodes:
-//   - leaf nodes correspond to a token in the expanded token stream,
+//   - leaf nodes correspond to tokens,
 //   - tree nodes correspond to language grammar constructs.
 //
 // The tree is initially built from an AST. Each node of a newly built tree
 // covers a continous subrange of expanded tokens (i.e. tokens after
 // preprocessing), the specific tokens coverered are stored in the leaf nodes of
 // a tree. A post-order traversal of a tree will visit leaf nodes in an order
 // corresponding the original order of expanded tokens.
 //
 // This is still work in progress and highly experimental, we leave room for
 // ourselves to completely change the design and/or implementation.
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLING_SYNTAX_TREE_H
 #define LLVM_CLANG_TOOLING_SYNTAX_TREE_H
 
-#include "clang/Basic/LangOptions.h"
-#include "clang/Basic/SourceLocation.h"
-#include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
-#include "clang/Tooling/Syntax/Tokens.h"
+#include "clang/Tooling/Syntax/TokenManager.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Support/Allocator.h"
 #include <cstdint>
 #include <iterator>
 
 namespace clang {
 namespace syntax {
 
-/// A memory arena for syntax trees. Also tracks the underlying token buffers,
-/// source manager, etc.
+/// A memory arena for syntax trees.
+// FIXME: use BumpPtrAllocator directly.
 class Arena {
 public:
-  Arena(SourceManager &SourceMgr, const LangOptions &LangOpts,
-        const TokenBuffer &Tokens);
-
-  const SourceManager &getSourceManager() const { return SourceMgr; }
-  const LangOptions &getLangOptions() const { return LangOpts; }
-
-  const TokenBuffer &getTokenBuffer() const;
   llvm::BumpPtrAllocator &getAllocator() { return Allocator; }
-
-private:
-  /// Add \p Buffer to the underlying source manager, tokenize it and store the
-  /// resulting tokens. Used exclusively in `FactoryImpl` to materialize tokens
-  /// that were not written in user code.
-  std::pair<FileID, ArrayRef<Token>>
-  lexBuffer(std::unique_ptr<llvm::MemoryBuffer> Buffer);
-  friend class FactoryImpl;
-
 private:
-  SourceManager &SourceMgr;
-  const LangOptions &LangOpts;
-  const TokenBuffer &Tokens;
-  /// IDs and storage for additional tokenized files.
-  llvm::DenseMap<FileID, std::vector<Token>> ExtraTokens;
   /// Keeps all the allocated nodes and their intermediate data structures.
   llvm::BumpPtrAllocator Allocator;
 };
 
 class Tree;
 class TreeBuilder;
 class FactoryImpl;
 class MutationsImpl;
 
public:
   Tree *getParent() { return Parent; }
 
   const Node *getNextSibling() const { return NextSibling; }
   Node *getNextSibling() { return NextSibling; }
   const Node *getPreviousSibling() const { return PreviousSibling; }
   Node *getPreviousSibling() { return PreviousSibling; }
 
   /// Dumps the structure of a subtree. For debugging and testing purposes.
-  std::string dump(const SourceManager &SM) const;
+  std::string dump(const TokenManager &SM) const;
   /// Dumps the tokens forming this subtree.
-  std::string dumpTokens(const SourceManager &SM) const;
+  std::string dumpTokens(const TokenManager &SM) const;
 
   /// Asserts invariants on this node of the tree and its immediate children.
   /// Will not recurse into the subtree. No-op if NDEBUG is set.
   void assertInvariants() const;
   /// Runs checkInvariants on all nodes in the subtree. No-op if NDEBUG is set.
   void assertInvariantsRecursive() const;
 
 private:
private:
   Node *NextSibling;
   Node *PreviousSibling;
   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.
+/// A leaf node points to a single token.
+// FIXME: add TokenKind field (borrow some bits from the Node::kind).
 class Leaf final : public Node {
 public:
-  Leaf(const Token *T);
+  Leaf(TokenManager::Key K);
   static bool classof(const Node *N);
 
-  const Token *getToken() const { return Tok; }
+  TokenManager::Key getTokenKey() const { return K; }
 
 private:
-  const Token *Tok;
+  TokenManager::Key K;
 };
 
 /// A node that has children and represents a syntactic language construct.
 class Tree : public Node {
   /// Iterator over children (common base for const/non-const).
   /// Not invalidated by tree mutations (holds a stable node pointer).
   template <typename DerivedT, typename NodeT>
   class ChildIteratorBase

clang/lib/Tooling/Syntax/BuildTree.cpp

 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/LiteralSupport.h"
 #include "clang/Tooling/Syntax/Nodes.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.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"

 ///   - replace the child nodes with the new syntax node in the pending list
 ///     with 'foldNode'.
 ///
 /// Note that all children are expected to be processed when building a node.
 ///
 /// Call finalize() to finish building the tree and consume the root node.
 class syntax::TreeBuilder {
 public:
-  TreeBuilder(syntax::Arena &Arena) : Arena(Arena), Pending(Arena) {
-    for (const auto &T : Arena.getTokenBuffer().expandedTokens())
+  TreeBuilder(syntax::Arena &Arena, TokenBufferTokenManager& TBTM)
+      : Arena(Arena),
+        TBTM(TBTM),
+        Pending(Arena, TBTM.tokenBuffer()) {
+    for (const auto &T : TBTM.tokenBuffer().expandedTokens())
       LocationToToken.insert({T.location(), &T});
   }
 
   llvm::BumpPtrAllocator &allocator() { return Arena.getAllocator(); }
   const SourceManager &sourceManager() const {
-    return Arena.getSourceManager();
+    return TBTM.sourceManager();
   }
 
   /// Populate children for \p New node, assuming it covers tokens from \p
   /// Range.
   void foldNode(ArrayRef<syntax::Token> Range, syntax::Tree *New, ASTPtr From) {
     assert(New);
-    Pending.foldChildren(Arena, Range, New);
+    Pending.foldChildren(TBTM.tokenBuffer(), Range, New);
     if (From)
       Mapping.add(From, New);
   }
 
   void foldNode(ArrayRef<syntax::Token> Range, syntax::Tree *New, TypeLoc L) {
     // FIXME: add mapping for TypeLocs
     foldNode(Range, New, nullptr);
   }
 
   void foldNode(llvm::ArrayRef<syntax::Token> Range, syntax::Tree *New,
                 NestedNameSpecifierLoc From) {
     assert(New);
-    Pending.foldChildren(Arena, Range, New);
+    Pending.foldChildren(TBTM.tokenBuffer(), Range, New);
     if (From)
       Mapping.add(From, New);
   }
 
   /// Populate children for \p New list, assuming it covers tokens from a
   /// subrange of \p SuperRange.
   void foldList(ArrayRef<syntax::Token> SuperRange, syntax::List *New,
                 ASTPtr From) {
     assert(New);
     auto ListRange = Pending.shrinkToFitList(SuperRange);
-    Pending.foldChildren(Arena, ListRange, New);
+    Pending.foldChildren(TBTM.tokenBuffer(), ListRange, New);
     if (From)
       Mapping.add(From, New);
   }
 
   /// Notifies that we should not consume trailing semicolon when computing
   /// token range of \p D.
   void noticeDeclWithoutSemicolon(Decl *D);
 
public:
   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 syntax node matching \p N.
   void markChild(NestedNameSpecifierLoc N, NodeRole R);
 
   /// Finish building the tree and consume the root node.
   syntax::TranslationUnit *finalize() && {
-    auto Tokens = Arena.getTokenBuffer().expandedTokens();
+    auto Tokens = TBTM.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(),
+    Pending.foldChildren(TBTM.tokenBuffer(), Tokens.drop_back(),
                          new (Arena.getAllocator()) 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;
 
   /// Finds the syntax tokens corresponding to the \p SourceRange.
   ArrayRef<syntax::Token> getRange(SourceRange Range) const {
     assert(Range.isValid());
     return getRange(Range.getBegin(), Range.getEnd());
   }
 
   /// 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.
   ArrayRef<syntax::Token> getRange(SourceLocation First,
                                    SourceLocation Last) const {
     assert(First.isValid());
     assert(Last.isValid());
     assert(First == Last ||
-           Arena.getSourceManager().isBeforeInTranslationUnit(First, Last));
+           TBTM.sourceManager().isBeforeInTranslationUnit(First, Last));
     return llvm::makeArrayRef(findToken(First), std::next(findToken(Last)));
   }
 
   ArrayRef<syntax::Token>
   getTemplateRange(const ClassTemplateSpecializationDecl *D) const {
     auto Tokens = getRange(D->getSourceRange());
     return maybeAppendSemicolon(Tokens, D);
   }
private:
 
   /// 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.getTokenBuffer().expandedTokens().empty());
-      assert(A.getTokenBuffer().expandedTokens().back().kind() == tok::eof);
+    Forest(syntax::Arena &A, const syntax::TokenBuffer &TB) {
+      assert(!TB.expandedTokens().empty());
+      assert(TB.expandedTokens().back().kind() == tok::eof);
       // Create all leaf nodes.
       // Note that we do not have 'eof' in the tree.
-      for (const auto &T : A.getTokenBuffer().expandedTokens().drop_back()) {
-        auto *L = new (A.getAllocator()) syntax::Leaf(&T);
+      for (const auto &T : TB.expandedTokens().drop_back()) {
+        auto *L = new (A.getAllocator())
+            syntax::Leaf(reinterpret_cast<TokenManager::Key>(&T));
         L->Original = true;
-        L->CanModify = A.getTokenBuffer().spelledForExpanded(T).has_value();
+        L->CanModify = TB.spelledForExpanded(T).has_value();
         Trees.insert(Trees.end(), {&T, L});
       }
     }
 
     void assignRole(ArrayRef<syntax::Token> Range, syntax::NodeRole Role) {
       assert(!Range.empty());
       auto It = Trees.lower_bound(Range.begin());
       assert(It != Trees.end() && "no node found");
ArrayRef<syntax::Token> shrinkToFitList(ArrayRef<syntax::Token> Range) {
 
       auto EndListChildren =
           std::find_if_not(BeginListChildren, EndChildren, BelongsToList);
 
       return ArrayRef<syntax::Token>(BeginListChildren->first,
                                      EndListChildren->first);
     }
 
     /// Add \p Node to the forest and attach child nodes based on \p Tokens.
-    void foldChildren(const syntax::Arena &A, ArrayRef<syntax::Token> Tokens,
-                      syntax::Tree *Node) {
+    void foldChildren(const syntax::TokenBuffer &TB,
+                      ArrayRef<syntax::Token> Tokens, syntax::Tree *Node) {
       // Attach children to `Node`.
       assert(Node->getFirstChild() == 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");
 
       for (auto It = BeginChildren; It != EndChildren; ++It) {
         auto *C = It->second;
         if (C->getRole() == NodeRole::Detached)
           C->setRole(NodeRole::Unknown);
         Node->appendChildLowLevel(C);
       }
 
       // Mark that this node came from the AST and is backed by the source code.
       Node->Original = true;
       Node->CanModify =
-          A.getTokenBuffer().spelledForExpanded(Tokens).has_value();
+          TB.spelledForExpanded(Tokens).has_value();
 
       Trees.erase(BeginChildren, EndChildren);
       Trees.insert({FirstToken, Node});
     }
 
     // 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;
       Trees = {};
       return Root;
     }
 
-    std::string str(const syntax::Arena &A) const {
+    std::string str(const syntax::TokenBufferTokenManager &STM) 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.getTokenBuffer().expandedTokens().end() - It->first;
+                : STM.tokenBuffer().expandedTokens().end() - It->first;
 
         R += std::string(
             formatv("- '{0}' covers '{1}'+{2} tokens\n", It->second->getKind(),
-                    It->first->text(A.getSourceManager()), CoveredTokens));
-        R += It->second->dump(A.getSourceManager());
+                    It->first->text(STM.sourceManager()), CoveredTokens));
+        R += It->second->dump(STM);
       }
       return R;
     }
 
   private:
     /// 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.
     std::map<const syntax::Token *, syntax::Node *> Trees;
   };
 
   /// For debugging purposes.
-  std::string str() { return Pending.str(Arena); }
+  std::string str() { return Pending.str(TBTM); }
 
   syntax::Arena &Arena;
+  TokenBufferTokenManager& TBTM;
   /// To quickly find tokens by their start location.
   llvm::DenseMap<SourceLocation, const syntax::Token *> LocationToToken;
   Forest Pending;
   llvm::DenseSet<Decl *> DeclsWithoutSemicolons;
   ASTToSyntaxMapping Mapping;
 };
 
 namespace {
void syntax::TreeBuilder::markStmtChild(Stmt *Child, NodeRole Role) {
 
   syntax::Tree *ChildNode;
   if (Expr *ChildExpr = dyn_cast<Expr>(Child)) {
     // This is an expression in a statement position, consume the trailing
     // semicolon and form an 'ExpressionStatement' node.
     markExprChild(ChildExpr, NodeRole::Expression);
     ChildNode = new (allocator()) syntax::ExpressionStatement;
     // (!) 'getStmtRange()' ensures this covers a trailing semicolon.
-    Pending.foldChildren(Arena, getStmtRange(Child), ChildNode);
+    Pending.foldChildren(TBTM.tokenBuffer(), getStmtRange(Child), ChildNode);
   } else {
     ChildNode = Mapping.find(Child);
   }
   assert(ChildNode != nullptr);
   setRole(ChildNode, Role);
 }
 
 void syntax::TreeBuilder::markExprChild(Expr *Child, NodeRole Role) {
const syntax::Token *syntax::TreeBuilder::findToken(SourceLocation L) const {
   if (L.isInvalid())
     return nullptr;
   auto It = LocationToToken.find(L);
   assert(It != LocationToToken.end());
   return It->second;
 }
 
 syntax::TranslationUnit *syntax::buildSyntaxTree(Arena &A,
+                                                 TokenBufferTokenManager& TBTM,
                                                  ASTContext &Context) {
-  TreeBuilder Builder(A);
+  TreeBuilder Builder(A, TBTM);
   BuildTreeVisitor(Context, Builder).TraverseAST(Context);
   return std::move(Builder).finalize();
 }

clang/lib/Tooling/Syntax/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS Support)
 
 add_clang_library(clangToolingSyntax
   BuildTree.cpp
   ComputeReplacements.cpp
   Nodes.cpp
   Mutations.cpp
+  TokenBufferTokenManager.cpp
   Synthesis.cpp
   Tokens.cpp
   Tree.cpp
 
   LINK_LIBS
   clangAST
   clangBasic
   clangFrontend
   clangLex
   clangToolingCore
 
   DEPENDS
   omp_gen
   )

clang/lib/Tooling/Syntax/ComputeReplacements.cpp

 //===- ComputeReplacements.cpp --------------------------------*- C++ -*-=====//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Syntax/Mutations.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 #include "clang/Tooling/Syntax/Tokens.h"
+#include "clang/Tooling/Syntax/Tree.h"
 #include "llvm/Support/Error.h"
 
 using namespace clang;
 
 namespace {
 using ProcessTokensFn = llvm::function_ref<void(llvm::ArrayRef<syntax::Token>,
                                                 bool /*IsOriginal*/)>;
 /// Enumerates spans of tokens from the tree consecutively laid out in memory.
-void enumerateTokenSpans(const syntax::Tree *Root, ProcessTokensFn Callback) {
+void enumerateTokenSpans(const syntax::Tree *Root,
+                         const syntax::TokenBufferTokenManager &STM,
+                         ProcessTokensFn Callback) {
   struct Enumerator {
-    Enumerator(ProcessTokensFn Callback)
-        : SpanBegin(nullptr), SpanEnd(nullptr), SpanIsOriginal(false),
+    Enumerator(const syntax::TokenBufferTokenManager &STM,
+               ProcessTokensFn Callback)
+        : STM(STM), SpanBegin(nullptr), SpanEnd(nullptr), SpanIsOriginal(false),
           Callback(Callback) {}
 
     void run(const syntax::Tree *Root) {
       process(Root);
       // Report the last span to the user.
       if (SpanBegin)
         Callback(llvm::makeArrayRef(SpanBegin, SpanEnd), SpanIsOriginal);
     }
 
   private:
     void process(const syntax::Node *N) {
       if (auto *T = dyn_cast<syntax::Tree>(N)) {
         for (const auto *C = T->getFirstChild(); C != nullptr;
              C = C->getNextSibling())
           process(C);
         return;
       }
 
       auto *L = cast<syntax::Leaf>(N);
-      if (SpanEnd == L->getToken() && SpanIsOriginal == L->isOriginal()) {
+      if (SpanEnd == STM.getToken(L->getTokenKey()) &&
+          SpanIsOriginal == L->isOriginal()) {
         // Extend the current span.
         ++SpanEnd;
         return;
       }
       // Report the current span to the user.
       if (SpanBegin)
         Callback(llvm::makeArrayRef(SpanBegin, SpanEnd), SpanIsOriginal);
       // Start recording a new span.
-      SpanBegin = L->getToken();
+      SpanBegin = STM.getToken(L->getTokenKey());
       SpanEnd = SpanBegin + 1;
       SpanIsOriginal = L->isOriginal();
     }
 
+    const syntax::TokenBufferTokenManager &STM;
     const syntax::Token *SpanBegin;
     const syntax::Token *SpanEnd;
     bool SpanIsOriginal;
     ProcessTokensFn Callback;
   };
 
-  return Enumerator(Callback).run(Root);
+  return Enumerator(STM, Callback).run(Root);
 }
 
-syntax::FileRange rangeOfExpanded(const syntax::Arena &A,
+syntax::FileRange rangeOfExpanded(const syntax::TokenBufferTokenManager &STM,
                                   llvm::ArrayRef<syntax::Token> Expanded) {
-  const auto &Buffer = A.getTokenBuffer();
-  const auto &SM = A.getSourceManager();
+  const auto &Buffer = STM.tokenBuffer();
+  const auto &SM = STM.sourceManager();
 
   // Check that \p Expanded actually points into expanded tokens.
   assert(Buffer.expandedTokens().begin() <= Expanded.begin());
   assert(Expanded.end() < Buffer.expandedTokens().end());
 
   if (Expanded.empty())
     // (!) empty tokens must always point before end().
     return syntax::FileRange(
         SM, SM.getExpansionLoc(Expanded.begin()->location()), /*Length=*/0);
 
   auto Spelled = Buffer.spelledForExpanded(Expanded);
   assert(Spelled && "could not find spelled tokens for expanded");
   return syntax::Token::range(SM, Spelled->front(), Spelled->back());
 }
 } // namespace
 
 tooling::Replacements
-syntax::computeReplacements(const syntax::Arena &A,
+syntax::computeReplacements(const TokenBufferTokenManager &TBTM,
                             const syntax::TranslationUnit &TU) {
-  const auto &Buffer = A.getTokenBuffer();
-  const auto &SM = A.getSourceManager();
+  const auto &Buffer = TBTM.tokenBuffer();
+  const auto &SM = TBTM.sourceManager();
 
   tooling::Replacements Replacements;
   // Text inserted by the replacement we are building now.
   std::string Replacement;
   auto emitReplacement = [&](llvm::ArrayRef<syntax::Token> ReplacedRange) {
     if (ReplacedRange.empty() && Replacement.empty())
       return;
     llvm::cantFail(Replacements.add(tooling::Replacement(
-        SM, rangeOfExpanded(A, ReplacedRange).toCharRange(SM), Replacement)));
+        SM, rangeOfExpanded(TBTM, ReplacedRange).toCharRange(SM),
+        Replacement)));
     Replacement = "";
   };
-
   const syntax::Token *NextOriginal = Buffer.expandedTokens().begin();
   enumerateTokenSpans(
-      &TU, [&](llvm::ArrayRef<syntax::Token> Tokens, bool IsOriginal) {
+      &TU, TBTM, [&](llvm::ArrayRef<syntax::Token> Tokens, bool IsOriginal) {
         if (!IsOriginal) {
           Replacement +=
               syntax::Token::range(SM, Tokens.front(), Tokens.back()).text(SM);
           return;
         }
         assert(NextOriginal <= Tokens.begin());
         // We are looking at a span of original tokens.
         if (NextOriginal != Tokens.begin()) {

clang/lib/Tooling/Syntax/Mutations.cpp

static void remove(syntax::Node *N) {
     P->replaceChildRangeLowLevel(N, N->getNextSibling(),
                                  /*New=*/nullptr);
 
     P->assertInvariants();
     N->assertInvariants();
   }
 };
 
-void syntax::removeStatement(syntax::Arena &A, syntax::Statement *S) {
+void syntax::removeStatement(syntax::Arena &A, TokenBufferTokenManager &TBTM,
+                             syntax::Statement *S) {
   assert(S);
   assert(S->canModify());
 
   if (isa<CompoundStatement>(S->getParent())) {
     // A child of CompoundStatement can just be safely removed.
     MutationsImpl::remove(S);
     return;
   }
   // For the rest, we have to replace with an empty statement.
   if (isa<EmptyStatement>(S))
     return; // already an empty statement, nothing to do.
 
-  MutationsImpl::replace(S, createEmptyStatement(A));
+  MutationsImpl::replace(S, createEmptyStatement(A, TBTM));
 }

clang/lib/Tooling/Syntax/Synthesis.cpp

 //===- Synthesis.cpp ------------------------------------------*- C++ -*-=====//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Tooling/Syntax/BuildTree.h"
 #include "clang/Tooling/Syntax/Tree.h"
+#include "clang/Tooling/Syntax/Tokens.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 
 using namespace clang;
 
 /// Exposes private syntax tree APIs required to implement node synthesis.
 /// Should not be used for anything else.
 class clang::syntax::FactoryImpl {
 public:
   static void setCanModify(syntax::Node *N) { N->CanModify = true; }
 
   static void prependChildLowLevel(syntax::Tree *T, syntax::Node *Child,
                                    syntax::NodeRole R) {
     T->prependChildLowLevel(Child, R);
   }
   static void appendChildLowLevel(syntax::Tree *T, syntax::Node *Child,
                                   syntax::NodeRole R) {
     T->appendChildLowLevel(Child, R);
   }
 
   static std::pair<FileID, ArrayRef<Token>>
-  lexBuffer(syntax::Arena &A, std::unique_ptr<llvm::MemoryBuffer> Buffer) {
-    return A.lexBuffer(std::move(Buffer));
+  lexBuffer(TokenBufferTokenManager &TBTM,
+            std::unique_ptr<llvm::MemoryBuffer> Buffer) {
+    return TBTM.lexBuffer(std::move(Buffer));
   }
 };
 
 // FIXME: `createLeaf` is based on `syntax::tokenize` internally, as such it
 // doesn't support digraphs or line continuations.
-syntax::Leaf *clang::syntax::createLeaf(syntax::Arena &A, tok::TokenKind K,
-                                        StringRef Spelling) {
+syntax::Leaf *clang::syntax::createLeaf(syntax::Arena &A,
+                                        TokenBufferTokenManager &TBTM,
+                                        tok::TokenKind K, StringRef Spelling) {
   auto Tokens =
-      FactoryImpl::lexBuffer(A, llvm::MemoryBuffer::getMemBufferCopy(Spelling))
+      FactoryImpl::lexBuffer(TBTM, llvm::MemoryBuffer::getMemBufferCopy(Spelling))
           .second;
   assert(Tokens.size() == 1);
   assert(Tokens.front().kind() == K &&
          "spelling is not lexed into the expected kind of token");
 
-  auto *Leaf = new (A.getAllocator()) syntax::Leaf(Tokens.begin());
+  auto *Leaf = new (A.getAllocator()) syntax::Leaf(
+    reinterpret_cast<TokenManager::Key>(Tokens.begin()));
   syntax::FactoryImpl::setCanModify(Leaf);
   Leaf->assertInvariants();
   return Leaf;
 }
 
-syntax::Leaf *clang::syntax::createLeaf(syntax::Arena &A, tok::TokenKind K) {
+syntax::Leaf *clang::syntax::createLeaf(syntax::Arena &A,
+                                        TokenBufferTokenManager &TBTM,
+                                        tok::TokenKind K) {
   const auto *Spelling = tok::getPunctuatorSpelling(K);
   if (!Spelling)
     Spelling = tok::getKeywordSpelling(K);
   assert(Spelling &&
          "Cannot infer the spelling of the token from its token kind.");
-  return createLeaf(A, K, Spelling);
+  return createLeaf(A, TBTM, K, Spelling);
 }
 
 namespace {
 // Allocates the concrete syntax `Tree` according to its `NodeKind`.
 syntax::Tree *allocateTree(syntax::Arena &A, syntax::NodeKind Kind) {
   switch (Kind) {
   case syntax::NodeKind::Leaf:
     assert(false);
syntax::Tree *clang::syntax::createTree(
   for (const auto &Child : Children)
     FactoryImpl::appendChildLowLevel(T, Child.first, Child.second);
 
   T->assertInvariants();
   return T;
 }
 
 syntax::Node *clang::syntax::deepCopyExpandingMacros(syntax::Arena &A,
+                                                    TokenBufferTokenManager &TBTM,
                                                      const syntax::Node *N) {
   if (const auto *L = dyn_cast<syntax::Leaf>(N))
     // `L->getToken()` gives us the expanded token, thus we implicitly expand
     // any macros here.
-    return createLeaf(A, L->getToken()->kind(),
-                      L->getToken()->text(A.getSourceManager()));
+    return createLeaf(A, TBTM, TBTM.getToken(L->getTokenKey())->kind(),
+                       TBTM.getText(L->getTokenKey()));
 
   const auto *T = cast<syntax::Tree>(N);
   std::vector<std::pair<syntax::Node *, syntax::NodeRole>> Children;
   for (const auto *Child = T->getFirstChild(); Child;
        Child = Child->getNextSibling())
-    Children.push_back({deepCopyExpandingMacros(A, Child), Child->getRole()});
+    Children.push_back({deepCopyExpandingMacros(A, TBTM, Child), Child->getRole()});
 
   return createTree(A, Children, N->getKind());
 }
 
-syntax::EmptyStatement *clang::syntax::createEmptyStatement(syntax::Arena &A) {
+syntax::EmptyStatement *clang::syntax::createEmptyStatement(syntax::Arena &A,  TokenBufferTokenManager &TBTM) {
   return cast<EmptyStatement>(
-      createTree(A, {{createLeaf(A, tok::semi), NodeRole::Unknown}},
+      createTree(A, {{createLeaf(A, TBTM, tok::semi), NodeRole::Unknown}},
                  NodeKind::EmptyStatement));
 }

clang/lib/Tooling/Syntax/TokenBufferTokenManager.cpp

+//===- TokenBufferTokenManager.cpp ----------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
+
+namespace clang {
+namespace syntax {
+constexpr llvm::StringLiteral syntax::TokenBufferTokenManager::Kind;
+
+std::pair<FileID, ArrayRef<syntax::Token>>
+syntax::TokenBufferTokenManager::lexBuffer(
+    std::unique_ptr<llvm::MemoryBuffer> Input) {
+  auto FID = SM.createFileID(std::move(Input));
+  auto It = ExtraTokens.try_emplace(FID, tokenize(FID, SM, LangOpts));
+  assert(It.second && "duplicate FileID");
+  return {FID, It.first->second};
+}
+
+} // namespace syntax
+} // namespace clang

clang/lib/Tooling/Syntax/Tree.cpp

 static void traverse(syntax::Node *N,
                      llvm::function_ref<void(syntax::Node *)> Visit) {
   traverse(static_cast<const syntax::Node *>(N), [&](const syntax::Node *N) {
     Visit(const_cast<syntax::Node *>(N));
   });
 }
 } // namespace
 
-syntax::Arena::Arena(SourceManager &SourceMgr, const LangOptions &LangOpts,
-                     const TokenBuffer &Tokens)
-    : SourceMgr(SourceMgr), LangOpts(LangOpts), Tokens(Tokens) {}
-
-const syntax::TokenBuffer &syntax::Arena::getTokenBuffer() const {
-  return Tokens;
-}
-
-std::pair<FileID, ArrayRef<syntax::Token>>
-syntax::Arena::lexBuffer(std::unique_ptr<llvm::MemoryBuffer> Input) {
-  auto FID = SourceMgr.createFileID(std::move(Input));
-  auto It = ExtraTokens.try_emplace(FID, tokenize(FID, SourceMgr, LangOpts));
-  assert(It.second && "duplicate FileID");
-  return {FID, It.first->second};
-}
-
-syntax::Leaf::Leaf(const syntax::Token *Tok) : Node(NodeKind::Leaf), Tok(Tok) {
-  assert(Tok != nullptr);
-}
+syntax::Leaf::Leaf(syntax::TokenManager::Key K) : Node(NodeKind::Leaf), K(K) {}
 
 syntax::Node::Node(NodeKind Kind)
     : Parent(nullptr), NextSibling(nullptr), PreviousSibling(nullptr),
       Kind(static_cast<unsigned>(Kind)), Role(0), Original(false),
       CanModify(false) {
   this->setRole(NodeRole::Detached);
 }
 
for (auto *N = New; N != nullptr; N = N->NextSibling) {
     LastInNew = N;
     N->Parent = this;
   }
   LastInNew->NextSibling = End;
   NewLast = LastInNew;
 }
 
 namespace {
-static void dumpLeaf(raw_ostream &OS, const syntax::Leaf *L,
-                     const SourceManager &SM) {
-  assert(L);
-  const auto *Token = L->getToken();
-  assert(Token);
-  // Handle 'eof' separately, calling text() on it produces an empty string.
-  if (Token->kind() == tok::eof)
-    OS << "<eof>";
-  else
-    OS << Token->text(SM);
-}
-
 static void dumpNode(raw_ostream &OS, const syntax::Node *N,
-                     const SourceManager &SM, llvm::BitVector IndentMask) {
+                     const syntax::TokenManager &TM, llvm::BitVector IndentMask) {
   auto DumpExtraInfo = [&OS](const syntax::Node *N) {
     if (N->getRole() != syntax::NodeRole::Unknown)
       OS << " " << N->getRole();
     if (!N->isOriginal())
       OS << " synthesized";
     if (!N->canModify())
       OS << " unmodifiable";
   };
 
   assert(N);
   if (const auto *L = dyn_cast<syntax::Leaf>(N)) {
     OS << "'";
-    dumpLeaf(OS, L, SM);
+    OS << TM.getText(L->getTokenKey());
     OS << "'";
     DumpExtraInfo(N);
     OS << "\n";
     return;
   }
 
   const auto *T = cast<syntax::Tree>(N);
   OS << T->getKind();
for (const syntax::Node &It : T->getChildren()) {
     }
     if (!It.getNextSibling()) {
       OS << "`-";
       IndentMask.push_back(false);
     } else {
       OS << "|-";
       IndentMask.push_back(true);
     }
-    dumpNode(OS, &It, SM, IndentMask);
+    dumpNode(OS, &It, TM, IndentMask);
     IndentMask.pop_back();
   }
 }
 } // namespace
 
-std::string syntax::Node::dump(const SourceManager &SM) const {
+std::string syntax::Node::dump(const TokenManager &TM) const {
   std::string Str;
   llvm::raw_string_ostream OS(Str);
-  dumpNode(OS, this, SM, /*IndentMask=*/{});
+  dumpNode(OS, this, TM, /*IndentMask=*/{});
   return std::move(OS.str());
 }
 
-std::string syntax::Node::dumpTokens(const SourceManager &SM) const {
+std::string syntax::Node::dumpTokens(const TokenManager &TM) const {
   std::string Storage;
   llvm::raw_string_ostream OS(Storage);
   traverse(this, [&](const syntax::Node *N) {
     if (const auto *L = dyn_cast<syntax::Leaf>(N)) {
-      dumpLeaf(OS, L, SM);
+      OS << TM.getText(L->getTokenKey());
       OS << " ";
     }
   });
   return Storage;
 }
 
 void syntax::Node::assertInvariants() const {
 #ifndef NDEBUG
#ifndef NDEBUG
   const auto *L = dyn_cast<List>(T);
   if (!L)
     return;
   for (const Node &C : T->getChildren()) {
     assert(C.getRole() == NodeRole::ListElement ||
            C.getRole() == NodeRole::ListDelimiter);
     if (C.getRole() == NodeRole::ListDelimiter) {
       assert(isa<Leaf>(C));
-      assert(cast<Leaf>(C).getToken()->kind() == L->getDelimiterTokenKind());
+      // FIXME: re-enable it when there is way to retrieve token kind in Leaf.
+      // assert(cast<Leaf>(C).getToken()->kind() == L->getDelimiterTokenKind());
     }
   }
 
 #endif
 }
 
 void syntax::Node::assertInvariantsRecursive() const {
 #ifndef NDEBUG

clang/tools/clang-check/ClangCheck.cpp

 #include "clang/Driver/Options.h"
 #include "clang/Frontend/ASTConsumers.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Rewrite/Frontend/FixItRewriter.h"
 #include "clang/Rewrite/Frontend/FrontendActions.h"
 #include "clang/StaticAnalyzer/Frontend/FrontendActions.h"
 #include "clang/Tooling/CommonOptionsParser.h"
 #include "clang/Tooling/Syntax/BuildTree.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "clang/Tooling/Syntax/Tree.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Option/OptTable.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/TargetSelect.h"
CreateASTConsumer(clang::CompilerInstance &CI, StringRef InFile) override {
     class Consumer : public clang::ASTConsumer {
     public:
       Consumer(clang::CompilerInstance &CI) : Collector(CI.getPreprocessor()) {}
 
       void HandleTranslationUnit(clang::ASTContext &AST) override {
         clang::syntax::TokenBuffer TB = std::move(Collector).consume();
         if (TokensDump)
           llvm::outs() << TB.dumpForTests();
-        clang::syntax::Arena A(AST.getSourceManager(), AST.getLangOpts(), TB);
-        llvm::outs() << clang::syntax::buildSyntaxTree(A, AST)->dump(
+        clang::syntax::TokenBufferTokenManager TBTM(TB, AST.getLangOpts(),
                                                     AST.getSourceManager());
+        clang::syntax::Arena A;
+        llvm::outs()
+            << clang::syntax::buildSyntaxTree(A, TBTM, AST)->dump(TBTM);
       }
 
     private:
       clang::syntax::TokenCollector Collector;
     };
     return std::make_unique<Consumer>(CI);
   }
 };

clang/unittests/Tooling/Syntax/BuildTreeTest.cpp

 protected:
   ::testing::AssertionResult treeDumpEqual(StringRef Code, StringRef Tree) {
     SCOPED_TRACE(llvm::join(GetParam().getCommandLineArgs(), " "));
 
     auto *Root = buildTree(Code, GetParam());
     auto ErrorOK = errorOK(Code);
     if (!ErrorOK)
       return ErrorOK;
-    auto Actual = StringRef(Root->dump(Arena->getSourceManager())).trim().str();
+    auto Actual = StringRef(Root->dump(*TM)).trim().str();
     // EXPECT_EQ shows the diff between the two strings if they are different.
     EXPECT_EQ(Tree.trim().str(), Actual);
     if (Actual != Tree.trim().str()) {
       return ::testing::AssertionFailure();
     }
     return ::testing::AssertionSuccess();
   }
 
if (AnnotatedRanges.size() != TreeDumps.size()) {
                 "different "
                 "to the number of their corresponding tree dumps.";
     }
     bool Failed = false;
     for (unsigned i = 0; i < AnnotatedRanges.size(); i++) {
       auto *AnnotatedNode = nodeByRange(AnnotatedRanges[i], Root);
       assert(AnnotatedNode);
       auto AnnotatedNodeDump =
-          StringRef(AnnotatedNode->dump(Arena->getSourceManager()))
+          StringRef(AnnotatedNode->dump(*TM))
               .trim()
               .str();
       // EXPECT_EQ shows the diff between the two strings if they are different.
       EXPECT_EQ(TreeDumps[i].trim().str(), AnnotatedNodeDump)
           << "Dumps diverged for the code:\n"
           << AnnotatedCode.code().slice(AnnotatedRanges[i].Begin,
                                         AnnotatedRanges[i].End);
       if (AnnotatedNodeDump != TreeDumps[i].trim().str())

clang/unittests/Tooling/Syntax/MutationsTest.cpp

using Transformation = std::function<void(const llvm::Annotations & /*Input*/,
                                             TranslationUnit * /*Root*/)>;
   void CheckTransformation(Transformation Transform, std::string Input,
                            std::string Expected) {
     llvm::Annotations Source(Input);
     auto *Root = buildTree(Source.code(), GetParam());
 
     Transform(Source, Root);
 
-    auto Replacements = syntax::computeReplacements(*Arena, *Root);
+    auto Replacements = syntax::computeReplacements(*TM, *Root);
     auto Output = tooling::applyAllReplacements(Source.code(), Replacements);
     if (!Output) {
       ADD_FAILURE() << "could not apply replacements: "
                     << llvm::toString(Output.takeError());
       return;
     }
 
     EXPECT_EQ(Expected, *Output) << "input is:\n" << Input;
   };
 
   // Removes the selected statement. Input should have exactly one selected
   // range and it should correspond to a single statement.
   Transformation RemoveStatement = [this](const llvm::Annotations &Input,
                                           TranslationUnit *Root) {
     auto *S = cast<syntax::Statement>(nodeByRange(Input.range(), Root));
     ASSERT_TRUE(S->canModify()) << "cannot remove a statement";
-    syntax::removeStatement(*Arena, S);
+    syntax::removeStatement(*Arena, *TM, S);
     EXPECT_TRUE(S->isDetached());
     EXPECT_FALSE(S->isOriginal())
         << "node removed from tree cannot be marked as original";
   };
 };
 
 INSTANTIATE_TEST_SUITE_P(SyntaxTreeTests, MutationTest,
                         ::testing::ValuesIn(allTestClangConfigs()) );

clang/unittests/Tooling/Syntax/SynthesisTest.cpp

 
 class SynthesisTest : public SyntaxTreeTest {
 protected:
   ::testing::AssertionResult treeDumpEqual(syntax::Node *Root, StringRef Dump) {
     if (!Root)
       return ::testing::AssertionFailure()
              << "Root was not built successfully.";
 
-    auto Actual = StringRef(Root->dump(Arena->getSourceManager())).trim().str();
+    auto Actual = StringRef(Root->dump(*TM)).trim().str();
     auto Expected = Dump.trim().str();
     // EXPECT_EQ shows the diff between the two strings if they are different.
     EXPECT_EQ(Expected, Actual);
     if (Actual != Expected) {
       return ::testing::AssertionFailure();
     }
     return ::testing::AssertionSuccess();
   }
 };
 
 INSTANTIATE_TEST_SUITE_P(SynthesisTests, SynthesisTest,
                         ::testing::ValuesIn(allTestClangConfigs()) );
 
 TEST_P(SynthesisTest, Leaf_Punctuation) {
   buildTree("", GetParam());
 
-  auto *Leaf = createLeaf(*Arena, tok::comma);
+  auto *Leaf = createLeaf(*Arena, *TM, tok::comma);
 
   EXPECT_TRUE(treeDumpEqual(Leaf, R"txt(
 ',' Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Leaf_Punctuation_CXX) {
   if (!GetParam().isCXX())
     return;
 
   buildTree("", GetParam());
 
-  auto *Leaf = createLeaf(*Arena, tok::coloncolon);
+  auto *Leaf = createLeaf(*Arena, *TM, tok::coloncolon);
 
   EXPECT_TRUE(treeDumpEqual(Leaf, R"txt(
 '::' Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Leaf_Keyword) {
   buildTree("", GetParam());
 
-  auto *Leaf = createLeaf(*Arena, tok::kw_if);
+  auto *Leaf = createLeaf(*Arena, *TM, tok::kw_if);
 
   EXPECT_TRUE(treeDumpEqual(Leaf, R"txt(
 'if' Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Leaf_Keyword_CXX11) {
   if (!GetParam().isCXX11OrLater())
     return;
 
   buildTree("", GetParam());
 
-  auto *Leaf = createLeaf(*Arena, tok::kw_nullptr);
+  auto *Leaf = createLeaf(*Arena, *TM, tok::kw_nullptr);
 
   EXPECT_TRUE(treeDumpEqual(Leaf, R"txt(
 'nullptr' Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Leaf_Identifier) {
   buildTree("", GetParam());
 
-  auto *Leaf = createLeaf(*Arena, tok::identifier, "a");
+  auto *Leaf = createLeaf(*Arena, *TM, tok::identifier, "a");
 
   EXPECT_TRUE(treeDumpEqual(Leaf, R"txt(
 'a' Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Leaf_Number) {
   buildTree("", GetParam());
 
-  auto *Leaf = createLeaf(*Arena, tok::numeric_constant, "1");
+  auto *Leaf = createLeaf(*Arena, *TM, tok::numeric_constant, "1");
 
   EXPECT_TRUE(treeDumpEqual(Leaf, R"txt(
 '1' Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Tree_Empty) {
   buildTree("", GetParam());
 
   auto *Tree = createTree(*Arena, {}, NodeKind::UnknownExpression);
 
   EXPECT_TRUE(treeDumpEqual(Tree, R"txt(
 UnknownExpression Detached synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Tree_Flat) {
   buildTree("", GetParam());
 
-  auto *LeafLParen = createLeaf(*Arena, tok::l_paren);
-  auto *LeafRParen = createLeaf(*Arena, tok::r_paren);
+  auto *LeafLParen = createLeaf(*Arena, *TM, tok::l_paren);
+  auto *LeafRParen = createLeaf(*Arena, *TM, tok::r_paren);
   auto *TreeParen = createTree(*Arena,
                                {{LeafLParen, NodeRole::LeftHandSide},
                                 {LeafRParen, NodeRole::RightHandSide}},
                                NodeKind::ParenExpression);
 
   EXPECT_TRUE(treeDumpEqual(TreeParen, R"txt(
 ParenExpression Detached synthesized
 |-'(' LeftHandSide synthesized
 `-')' RightHandSide synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Tree_OfTree) {
   buildTree("", GetParam());
 
-  auto *Leaf1 = createLeaf(*Arena, tok::numeric_constant, "1");
+  auto *Leaf1 = createLeaf(*Arena, *TM, tok::numeric_constant, "1");
   auto *Int1 = createTree(*Arena, {{Leaf1, NodeRole::LiteralToken}},
                           NodeKind::IntegerLiteralExpression);
 
-  auto *LeafPlus = createLeaf(*Arena, tok::plus);
+  auto *LeafPlus = createLeaf(*Arena, *TM, tok::plus);
 
-  auto *Leaf2 = createLeaf(*Arena, tok::numeric_constant, "2");
+  auto *Leaf2 = createLeaf(*Arena, *TM, tok::numeric_constant, "2");
   auto *Int2 = createTree(*Arena, {{Leaf2, NodeRole::LiteralToken}},
                           NodeKind::IntegerLiteralExpression);
 
   auto *TreeBinaryOperator = createTree(*Arena,
                                         {{Int1, NodeRole::LeftHandSide},
                                          {LeafPlus, NodeRole::OperatorToken},
                                          {Int2, NodeRole::RightHandSide}},
                                         NodeKind::BinaryOperatorExpression);
 
   EXPECT_TRUE(treeDumpEqual(TreeBinaryOperator, R"txt(
 BinaryOperatorExpression Detached synthesized
 |-IntegerLiteralExpression LeftHandSide synthesized
 | `-'1' LiteralToken synthesized
 |-'+' OperatorToken synthesized
 `-IntegerLiteralExpression RightHandSide synthesized
   `-'2' LiteralToken synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, DeepCopy_Synthesized) {
   buildTree("", GetParam());
 
-  auto *LeafContinue = createLeaf(*Arena, tok::kw_continue);
-  auto *LeafSemiColon = createLeaf(*Arena, tok::semi);
+  auto *LeafContinue = createLeaf(*Arena, *TM, tok::kw_continue);
+  auto *LeafSemiColon = createLeaf(*Arena, *TM, tok::semi);
   auto *StatementContinue = createTree(*Arena,
                                        {{LeafContinue, NodeRole::LiteralToken},
                                         {LeafSemiColon, NodeRole::Unknown}},
                                        NodeKind::ContinueStatement);
 
-  auto *Copy = deepCopyExpandingMacros(*Arena, StatementContinue);
-  EXPECT_TRUE(
-      treeDumpEqual(Copy, StatementContinue->dump(Arena->getSourceManager())));
+  auto *Copy = deepCopyExpandingMacros(*Arena, *TM, StatementContinue);
+  EXPECT_TRUE(treeDumpEqual(Copy, StatementContinue->dump(*TM)));
   // FIXME: Test that copy is independent of original, once the Mutations API is
   // more developed.
 }
 
 TEST_P(SynthesisTest, DeepCopy_Original) {
   auto *OriginalTree = buildTree("int a;", GetParam());
 
-  auto *Copy = deepCopyExpandingMacros(*Arena, OriginalTree);
+  auto *Copy = deepCopyExpandingMacros(*Arena, *TM, OriginalTree);
   EXPECT_TRUE(treeDumpEqual(Copy, R"txt(
 TranslationUnit Detached synthesized
 `-SimpleDeclaration synthesized
   |-'int' synthesized
   |-DeclaratorList Declarators synthesized
   | `-SimpleDeclarator ListElement synthesized
   |   `-'a' synthesized
   `-';' synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, DeepCopy_Child) {
   auto *OriginalTree = buildTree("int a;", GetParam());
 
-  auto *Copy = deepCopyExpandingMacros(*Arena, OriginalTree->getFirstChild());
+  auto *Copy =
+      deepCopyExpandingMacros(*Arena, *TM, OriginalTree->getFirstChild());
   EXPECT_TRUE(treeDumpEqual(Copy, R"txt(
 SimpleDeclaration Detached synthesized
 |-'int' synthesized
 |-DeclaratorList Declarators synthesized
 | `-SimpleDeclarator ListElement synthesized
 |   `-'a' synthesized
 `-';' synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, DeepCopy_Macro) {
   auto *OriginalTree = buildTree(R"cpp(
 #define HALF_IF if (1+
 #define HALF_IF_2 1) {}
 void test() {
   HALF_IF HALF_IF_2 else {}
 })cpp",
                                  GetParam());
 
-  auto *Copy = deepCopyExpandingMacros(*Arena, OriginalTree);
+  auto *Copy = deepCopyExpandingMacros(*Arena, *TM, OriginalTree);
 
   // The syntax tree stores already expanded Tokens, we can only see whether the
   // macro was expanded when computing replacements. The dump does show that
   // nodes in the copy are `modifiable`.
   EXPECT_TRUE(treeDumpEqual(Copy, R"txt(
 TranslationUnit Detached synthesized
 `-SimpleDeclaration synthesized
   |-'void' synthesized
`-CompoundStatement synthesized
     |   `-'}' CloseParen synthesized
     `-'}' CloseParen synthesized
   )txt"));
 }
 
 TEST_P(SynthesisTest, Statement_EmptyStatement) {
   buildTree("", GetParam());
 
-  auto *S = createEmptyStatement(*Arena);
+  auto *S = createEmptyStatement(*Arena, *TM);
   EXPECT_TRUE(treeDumpEqual(S, R"txt(
 EmptyStatement Detached synthesized
 `-';' synthesized
   )txt"));
 }
 } // namespace

clang/unittests/Tooling/Syntax/TreeTest.cpp

 
 class TreeTest : public SyntaxTreeTest {
 private:
   Tree *createTree(ArrayRef<const Node *> Children) {
     std::vector<std::pair<Node *, NodeRole>> ChildrenWithRoles;
     ChildrenWithRoles.reserve(Children.size());
     for (const auto *Child : Children) {
       ChildrenWithRoles.push_back(std::make_pair(
-          deepCopyExpandingMacros(*Arena, Child), NodeRole::Unknown));
+          deepCopyExpandingMacros(*Arena, *TM, Child), NodeRole::Unknown));
     }
     return clang::syntax::createTree(*Arena, ChildrenWithRoles,
                                      NodeKind::UnknownExpression);
   }
 
   // Generate Forests by combining `Children` into `ParentCount` Trees.
   //
   // We do this recursively.
protected:
   }
 };
 
 INSTANTIATE_TEST_SUITE_P(TreeTests, TreeTest,
                         ::testing::ValuesIn(allTestClangConfigs()) );
 
 TEST_P(TreeTest, FirstLeaf) {
   buildTree("", GetParam());
-  std::vector<const Node *> Leafs = {createLeaf(*Arena, tok::l_paren),
-                                     createLeaf(*Arena, tok::r_paren)};
+  std::vector<const Node *> Leafs = {createLeaf(*Arena, *TM, tok::l_paren),
+                                     createLeaf(*Arena, *TM, tok::r_paren)};
   for (const auto *Tree : generateAllTreesWithShape(Leafs, {3u})) {
     ASSERT_TRUE(Tree->findFirstLeaf() != nullptr);
-    EXPECT_EQ(Tree->findFirstLeaf()->getToken()->kind(), tok::l_paren);
+    EXPECT_EQ(TM->getToken(Tree->findFirstLeaf()->getTokenKey())->kind(), tok::l_paren);
   }
 }
 
 TEST_P(TreeTest, LastLeaf) {
   buildTree("", GetParam());
-  std::vector<const Node *> Leafs = {createLeaf(*Arena, tok::l_paren),
-                                     createLeaf(*Arena, tok::r_paren)};
+  std::vector<const Node *> Leafs = {createLeaf(*Arena, *TM, tok::l_paren),
+                                     createLeaf(*Arena, *TM, tok::r_paren)};
   for (const auto *Tree : generateAllTreesWithShape(Leafs, {3u})) {
     ASSERT_TRUE(Tree->findLastLeaf() != nullptr);
-    EXPECT_EQ(Tree->findLastLeaf()->getToken()->kind(), tok::r_paren);
+    EXPECT_EQ(TM->getToken(Tree->findLastLeaf()->getTokenKey())->kind(), tok::r_paren);
   }
 }
 
 TEST_F(TreeTest, Iterators) {
   buildTree("", allTestClangConfigs().front());
-  std::vector<Node *> Children = {createLeaf(*Arena, tok::identifier, "a"),
-                                  createLeaf(*Arena, tok::identifier, "b"),
-                                  createLeaf(*Arena, tok::identifier, "c")};
+  std::vector<Node *> Children = {createLeaf(*Arena, *TM, tok::identifier, "a"),
+                                  createLeaf(*Arena, *TM, tok::identifier, "b"),
+                                  createLeaf(*Arena, *TM, tok::identifier, "c")};
   auto *Tree = syntax::createTree(*Arena,
                                   {{Children[0], NodeRole::LeftHandSide},
                                    {Children[1], NodeRole::OperatorToken},
                                    {Children[2], NodeRole::RightHandSide}},
                                   NodeKind::TranslationUnit);
   const auto *ConstTree = Tree;
 
   auto Range = Tree->getChildren();
TEST_F(TreeTest, Iterators) {
   EXPECT_EQ(nullptr, It.asPointer());
   EXPECT_EQ(nullptr, CIt.asPointer());
 }
 
 class ListTest : public SyntaxTreeTest {
 private:
   std::string dumpQuotedTokensOrNull(const Node *N) {
     return N ? "'" +
-                   StringRef(N->dumpTokens(Arena->getSourceManager()))
+                   StringRef(N->dumpTokens(*TM))
                        .trim()
                        .str() +
                    "'"
              : "null";
   }
 
 protected:
   std::string

 /// "a, b, c"  <=> [("a", ","), ("b", ","), ("c", null)]
 TEST_P(ListTest, List_Separated_WellFormed) {
   buildTree("", GetParam());
 
   // "a, b, c"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "b"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "b"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
       },
       NodeKind::CallArguments));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', ','), ('b', ','), ('c', null)]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', 'b', 'c']");
 }
 
 /// "a,  , c"  <=> [("a", ","), (null, ","), ("c", null)]
 TEST_P(ListTest, List_Separated_MissingElement) {
   buildTree("", GetParam());
 
   // "a,  , c"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
       },
       NodeKind::CallArguments));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', ','), (null, ','), ('c', null)]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', null, 'c']");
 }
 
 /// "a, b  c"  <=> [("a", ","), ("b", null), ("c", null)]
 TEST_P(ListTest, List_Separated_MissingDelimiter) {
   buildTree("", GetParam());
 
   // "a, b  c"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "b"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "b"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
       },
       NodeKind::CallArguments));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', ','), ('b', null), ('c', null)]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', 'b', 'c']");
 }
 
 /// "a, b,"    <=> [("a", ","), ("b", ","), (null, null)]
 TEST_P(ListTest, List_Separated_MissingLastElement) {
   buildTree("", GetParam());
 
   // "a, b, c"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "b"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "b"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::comma), NodeRole::ListDelimiter},
       },
       NodeKind::CallArguments));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', ','), ('b', ','), (null, null)]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', 'b', null]");
 }
 
 /// "a:: b:: c::" <=> [("a", "::"), ("b", "::"), ("c", "::")]
 TEST_P(ListTest, List_Terminated_WellFormed) {
   if (!GetParam().isCXX()) {
     return;
   }
   buildTree("", GetParam());
 
   // "a:: b:: c::"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "b"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "b"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
       },
       NodeKind::NestedNameSpecifier));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', '::'), ('b', '::'), ('c', '::')]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', 'b', 'c']");
 }
 
 /// "a::  :: c::" <=> [("a", "::"), (null, "::"), ("c", "::")]
 TEST_P(ListTest, List_Terminated_MissingElement) {
   if (!GetParam().isCXX()) {
     return;
   }
   buildTree("", GetParam());
 
   // "a:: b:: c::"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
       },
       NodeKind::NestedNameSpecifier));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', '::'), (null, '::'), ('c', '::')]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', null, 'c']");
 }
 
 /// "a:: b  c::" <=> [("a", "::"), ("b", null), ("c", "::")]
 TEST_P(ListTest, List_Terminated_MissingDelimiter) {
   if (!GetParam().isCXX()) {
     return;
   }
   buildTree("", GetParam());
 
   // "a:: b  c::"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "b"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "b"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
       },
       NodeKind::NestedNameSpecifier));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', '::'), ('b', null), ('c', '::')]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', 'b', 'c']");
 }
 
 /// "a:: b:: c"  <=> [("a", "::"), ("b", "::"), ("c", null)]
 TEST_P(ListTest, List_Terminated_MissingLastDelimiter) {
   if (!GetParam().isCXX()) {
     return;
   }
   buildTree("", GetParam());
 
   // "a:: b:: c"
   auto *List = dyn_cast<syntax::List>(syntax::createTree(
       *Arena,
       {
-          {createLeaf(*Arena, tok::identifier, "a"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "b"), NodeRole::ListElement},
-          {createLeaf(*Arena, tok::coloncolon), NodeRole::ListDelimiter},
-          {createLeaf(*Arena, tok::identifier, "c"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::identifier, "a"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "b"), NodeRole::ListElement},
+          {createLeaf(*Arena, *TM, tok::coloncolon), NodeRole::ListDelimiter},
+          {createLeaf(*Arena, *TM, tok::identifier, "c"), NodeRole::ListElement},
       },
       NodeKind::NestedNameSpecifier));
 
   EXPECT_EQ(dumpElementsAndDelimiters(List->getElementsAsNodesAndDelimiters()),
             "[('a', '::'), ('b', '::'), ('c', null)]");
   EXPECT_EQ(dumpNodes(List->getElementsAsNodes()), "['a', 'b', 'c']");
 }
 
 } // namespace

clang/unittests/Tooling/Syntax/TreeTestBase.h

 
 #ifndef LLVM_CLANG_UNITTESTS_TOOLING_SYNTAX_TREETESTBASE_H
 #define LLVM_CLANG_UNITTESTS_TOOLING_SYNTAX_TREETESTBASE_H
 
 #include "clang/Basic/LLVM.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Testing/TestClangConfig.h"
 #include "clang/Tooling/Syntax/Nodes.h"
+#include "clang/Tooling/Syntax/TokenBufferTokenManager.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "clang/Tooling/Syntax/Tree.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Testing/Support/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> FS =
       new llvm::vfs::InMemoryFileSystem;
   IntrusiveRefCntPtr<FileManager> FileMgr =
       new FileManager(FileSystemOptions(), FS);
   IntrusiveRefCntPtr<SourceManager> SourceMgr =
       new SourceManager(*Diags, *FileMgr);
   std::shared_ptr<CompilerInvocation> Invocation;
   // Set after calling buildTree().
   std::unique_ptr<syntax::TokenBuffer> TB;
+  std::unique_ptr<syntax::TokenBufferTokenManager> TM;
   std::unique_ptr<syntax::Arena> Arena;
 };
 
 std::vector<TestClangConfig> allTestClangConfigs();
 
 MATCHER_P(role, R, "") {
   if (arg.getRole() == R)
     return true;
   *result_listener << "role is " << llvm::to_string(arg.getRole());
   return false;
 }
 
 } // namespace syntax
 } // namespace clang
 #endif // LLVM_CLANG_UNITTESTS_TOOLING_SYNTAX_TREETESTBASE_H

clang/unittests/Tooling/Syntax/TreeTestBase.cpp

 #include "llvm/Support/Error.h"
 #include "llvm/Testing/Support/Annotations.h"
 #include "gtest/gtest.h"
 
 using namespace clang;
 using namespace clang::syntax;
 
 namespace {
-ArrayRef<syntax::Token> tokens(syntax::Node *N) {
+ArrayRef<syntax::Token> tokens(syntax::Node *N,
+                               const TokenBufferTokenManager &STM) {
   assert(N->isOriginal() && "tokens of modified nodes are not well-defined");
   if (auto *L = dyn_cast<syntax::Leaf>(N))
-    return llvm::makeArrayRef(L->getToken(), 1);
+    return llvm::makeArrayRef(STM.getToken(L->getTokenKey()), 1);
   auto *T = cast<syntax::Tree>(N);
-  return llvm::makeArrayRef(T->findFirstLeaf()->getToken(),
-                            T->findLastLeaf()->getToken() + 1);
+  return llvm::makeArrayRef(STM.getToken(T->findFirstLeaf()->getTokenKey()),
+                            STM.getToken(T->findLastLeaf()->getTokenKey()) + 1);
 }
 } // namespace
 
 std::vector<TestClangConfig> clang::syntax::allTestClangConfigs() {
   std::vector<TestClangConfig> all_configs;
   for (TestLanguage lang : {Lang_C89, Lang_C99, Lang_CXX03, Lang_CXX11,
                             Lang_CXX14, Lang_CXX17, Lang_CXX20}) {
     TestClangConfig config;

 syntax::TranslationUnit *
 SyntaxTreeTest::buildTree(StringRef Code, const TestClangConfig &ClangConfig) {
   // FIXME: this code is almost the identical to the one in TokensTest. Share
   //        it.
   class BuildSyntaxTree : public ASTConsumer {
   public:
     BuildSyntaxTree(syntax::TranslationUnit *&Root,
                     std::unique_ptr<syntax::TokenBuffer> &TB,
+                    std::unique_ptr<syntax::TokenBufferTokenManager> &TM,
                     std::unique_ptr<syntax::Arena> &Arena,
                     std::unique_ptr<syntax::TokenCollector> Tokens)
-        : Root(Root), TB(TB), Arena(Arena), Tokens(std::move(Tokens)) {
+        : Root(Root), TB(TB), TM(TM), Arena(Arena), Tokens(std::move(Tokens)) {
       assert(this->Tokens);
     }
 
     void HandleTranslationUnit(ASTContext &Ctx) override {
       TB = std::make_unique<syntax::TokenBuffer>(std::move(*Tokens).consume());
       Tokens = nullptr; // make sure we fail if this gets called twice.
-      Arena = std::make_unique<syntax::Arena>(Ctx.getSourceManager(),
-                                              Ctx.getLangOpts(), *TB);
-      Root = syntax::buildSyntaxTree(*Arena, Ctx);
+      TM = std::make_unique<syntax::TokenBufferTokenManager>(
+          *TB, Ctx.getLangOpts(), Ctx.getSourceManager());
+      Arena = std::make_unique<syntax::Arena>();
+      Root = syntax::buildSyntaxTree(*Arena, *TM, Ctx);
     }
 
   private:
     syntax::TranslationUnit *&Root;
     std::unique_ptr<syntax::TokenBuffer> &TB;
+    std::unique_ptr<syntax::TokenBufferTokenManager> &TM;
     std::unique_ptr<syntax::Arena> &Arena;
     std::unique_ptr<syntax::TokenCollector> Tokens;
   };
 
   class BuildSyntaxTreeAction : public ASTFrontendAction {
   public:
     BuildSyntaxTreeAction(syntax::TranslationUnit *&Root,
+                          std::unique_ptr<syntax::TokenBufferTokenManager> &TM,
                           std::unique_ptr<syntax::TokenBuffer> &TB,
                           std::unique_ptr<syntax::Arena> &Arena)
-        : Root(Root), TB(TB), Arena(Arena) {}
+        : Root(Root), TM(TM), TB(TB), Arena(Arena) {}
 
     std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                    StringRef InFile) override {
       // We start recording the tokens, ast consumer will take on the result.
       auto Tokens =
           std::make_unique<syntax::TokenCollector>(CI.getPreprocessor());
-      return std::make_unique<BuildSyntaxTree>(Root, TB, Arena,
+      return std::make_unique<BuildSyntaxTree>(Root, TB, TM, Arena,
                                                std::move(Tokens));
     }
 
   private:
     syntax::TranslationUnit *&Root;
+    std::unique_ptr<syntax::TokenBufferTokenManager> &TM;
     std::unique_ptr<syntax::TokenBuffer> &TB;
     std::unique_ptr<syntax::Arena> &Arena;
   };
 
   constexpr const char *FileName = "./input.cpp";
   FS->addFile(FileName, time_t(), llvm::MemoryBuffer::getMemBufferCopy(""));
 
   if (!Diags->getClient())
Invocation->getPreprocessorOpts().addRemappedFile(
       FileName, llvm::MemoryBuffer::getMemBufferCopy(Code).release());
   CompilerInstance Compiler;
   Compiler.setInvocation(Invocation);
   Compiler.setDiagnostics(Diags.get());
   Compiler.setFileManager(FileMgr.get());
   Compiler.setSourceManager(SourceMgr.get());
 
   syntax::TranslationUnit *Root = nullptr;
-  BuildSyntaxTreeAction Recorder(Root, this->TB, this->Arena);
+  BuildSyntaxTreeAction Recorder(Root, this->TM, this->TB, this->Arena);
 
   // Action could not be executed but the frontend didn't identify any errors
   // in the code ==> problem in setting up the action.
   if (!Compiler.ExecuteAction(Recorder) &&
       Diags->getClient()->getNumErrors() == 0) {
     ADD_FAILURE() << "failed to run the frontend";
     std::abort();
   }
   return Root;
 }
 
 syntax::Node *SyntaxTreeTest::nodeByRange(llvm::Annotations::Range R,
                                           syntax::Node *Root) {
-  ArrayRef<syntax::Token> Toks = tokens(Root);
+  ArrayRef<syntax::Token> Toks = tokens(Root, *TM);
 
   if (Toks.front().location().isFileID() && Toks.back().location().isFileID() &&
       syntax::Token::range(*SourceMgr, Toks.front(), Toks.back()) ==
           syntax::FileRange(SourceMgr->getMainFileID(), R.Begin, R.End))
     return Root;
 
   auto *T = dyn_cast<syntax::Tree>(Root);
   if (!T)
     return nullptr;
   for (auto *C = T->getFirstChild(); C != nullptr; C = C->getNextSibling()) {
     if (auto *Result = nodeByRange(R, C))
       return Result;
   }
   return nullptr;
 }