1 files changed, 22 insertions, 15 deletions

diff --git a/clang/lib/Tooling/Syntax/ComputeReplacements.cpp b/clang/lib/Tooling/Syntax/ComputeReplacements.cpp
index 31e1a40c74b6..08e09e4ebdbf 100644
--- a/clang/lib/Tooling/Syntax/ComputeReplacements.cpp
+++ b/clang/lib/Tooling/Syntax/ComputeReplacements.cpp
@@ -7,7 +7,9 @@
 //===----------------------------------------------------------------------===//
 #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;
@@ -16,10 +18,13 @@ 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) {
@@ -39,7 +44,8 @@ void enumerateTokenSpans(const syntax::Tree *Root, ProcessTokensFn Callback) {
       }
 
       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;
@@ -48,24 +54,25 @@ void enumerateTokenSpans(const syntax::Tree *Root, ProcessTokensFn Callback) {
       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());
@@ -83,10 +90,10 @@ syntax::FileRange rangeOfExpanded(const syntax::Arena &A,
 } // 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.
@@ -95,13 +102,13 @@ syntax::computeReplacements(const syntax::Arena &A,
     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);