[clang][Syntax] syntax::Arena doesnt own TokenBuffer

[kadircet]

Summary

Currently an Arena can only be built while consuming a TokenBuffer,
some users (like clangd) might want to share a TokenBuffer with multiple
compenents. This patch changes Arena's TokenBuffer member to be a reference so
that it can be created with read-only token buffers.

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

 namespace clang {
 namespace syntax {
 
 /// A memory arena for syntax trees. Also tracks the underlying token buffers,
 /// source manager, etc.
 class Arena {
 public:
   Arena(SourceManager &SourceMgr, const LangOptions &LangOpts,
-        TokenBuffer Tokens);
+        const TokenBuffer &Tokens);
 
   const SourceManager &sourceManager() const { return SourceMgr; }
   const LangOptions &langOptions() const { return LangOpts; }
 
   const TokenBuffer &tokenBuffer() const;
   llvm::BumpPtrAllocator &allocator() { return Allocator; }
 
   /// Add \p Buffer to the underlying source manager, tokenize it and store the
   /// resulting tokens. Useful when there is a need to materialize tokens that
   /// were not written in user code.
   std::pair<FileID, llvm::ArrayRef<syntax::Token>>
   lexBuffer(std::unique_ptr<llvm::MemoryBuffer> Buffer);
 
 private:
   SourceManager &SourceMgr;
   const LangOptions &LangOpts;
-  TokenBuffer Tokens;
+  const TokenBuffer &Tokens;
   /// IDs and storage for additional tokenized files.
   llvm::DenseMap<FileID, std::vector<syntax::Token>> ExtraTokens;
   /// Keeps all the allocated nodes and their intermediate data structures.
   llvm::BumpPtrAllocator Allocator;
 };
 
 class Tree;
 class TreeBuilder;

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,
-                     TokenBuffer Tokens)
-    : SourceMgr(SourceMgr), LangOpts(LangOpts), Tokens(std::move(Tokens)) {}
+                     const TokenBuffer &Tokens)
+    : SourceMgr(SourceMgr), LangOpts(LangOpts), Tokens(Tokens) {}
 
 const clang::syntax::TokenBuffer &syntax::Arena::tokenBuffer() const {
   return Tokens;
 }
 
 std::pair<FileID, llvm::ArrayRef<syntax::Token>>
 syntax::Arena::lexBuffer(std::unique_ptr<llvm::MemoryBuffer> Input) {
   auto FID = SourceMgr.createFileID(std::move(Input));

clang/unittests/Tooling/Syntax/TreeTest.cpp

 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Testing/Support/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <cstdlib>
+#include <memory>
 
 using namespace clang;
 
 namespace {
 static llvm::ArrayRef<syntax::Token> tokens(syntax::Node *N) {
   assert(N->isOriginal() && "tokens of modified nodes are not well-defined");
   if (auto *L = dyn_cast<syntax::Leaf>(N))
     return llvm::makeArrayRef(L->token(), 1);
   auto *T = cast<syntax::Tree>(N);
   return llvm::makeArrayRef(T->firstLeaf()->token(),
                             T->lastLeaf()->token() + 1);
 }
 
 class SyntaxTreeTest : public ::testing::Test,
                        public ::testing::WithParamInterface<TestClangConfig> {
 protected:
   // Build a syntax tree for the code.
   syntax::TranslationUnit *buildTree(llvm::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::Arena> &Arena,
                       std::unique_ptr<syntax::TokenCollector> Tokens)
-          : Root(Root), Arena(Arena), Tokens(std::move(Tokens)) {
+          : Root(Root), TB(TB), Arena(Arena), Tokens(std::move(Tokens)) {
         assert(this->Tokens);
       }
 
       void HandleTranslationUnit(ASTContext &Ctx) override {
-        Arena = std::make_unique<syntax::Arena>(Ctx.getSourceManager(),
-                                                Ctx.getLangOpts(),
-                                                std::move(*Tokens).consume());
+        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.getTranslationUnitDecl());
       }
 
     private:
       syntax::TranslationUnit *&Root;
+      std::unique_ptr<syntax::TokenBuffer> &TB;
       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::TokenBuffer> &TB,
                             std::unique_ptr<syntax::Arena> &Arena)
-          : Root(Root), Arena(Arena) {}
+          : Root(Root), 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, Arena,
+        return std::make_unique<BuildSyntaxTree>(Root, TB, Arena,
                                                  std::move(Tokens));
       }
 
     private:
       syntax::TranslationUnit *&Root;
+      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())
       Diags->setClient(new TextDiagnosticPrinter(llvm::errs(), DiagOpts.get()));
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->Arena);
+    BuildSyntaxTreeAction Recorder(Root, 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();
     }
protected:
   IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> FS =
       new llvm::vfs::InMemoryFileSystem;
   llvm::IntrusiveRefCntPtr<FileManager> FileMgr =
       new FileManager(FileSystemOptions(), FS);
   llvm::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::Arena> Arena;
 };
 
 TEST_P(SyntaxTreeTest, Simple) {
   EXPECT_TRUE(treeDumpEqual(
       R"cpp(
 int main() {}
 void foo() {}