Add option to reparse an ASTUnit keeping its SourceManager

[giulianobelinassi]

Summary

When issuing a reparse, the SourceManager is then rebuilt from the
FileMgr object it had in ASTUnit. This, however, issues an reload
of the files MemoryBuffer which could have been changed by some
analysis process.

For example, an tooling application can call overrideFileContents with
a new MemoryBuffer with modified pieces of code, and then issue a
reparse to either rebuild the tree or update the SourceLocations. This
is useful when used with the Rewriter object because one can do:

std::string *modified_str = new std::string(rewritebuf->begin(), rewritebuf->end());
auto new_membuff = llvm::MemoryBuffer::getMemBuffer(*modified_str);
SM.overrideFileContents(file_entry, std::move(new_membuff));

AU->Reparse(..., /*KeepSourceManager=*/true);

And then the modifications to the file is loaded into the AU instance
of ASTUnit object. The test case in this patch reflects this use case.

Signed-off-by: Giuliano Belinassi <[e-mail]>

clang/include/clang/Frontend/ASTUnit.h

private:
 
   /// Clear out and deallocate
   void ClearCachedCompletionResults();
 
   explicit ASTUnit(bool MainFileIsAST);
 
   bool Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
              std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
-             IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS);
+             IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+             bool KeepSourceMgr = false);
 
   std::unique_ptr<llvm::MemoryBuffer> getMainBufferWithPrecompiledPreamble(
       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
       CompilerInvocation &PreambleInvocationIn,
       IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS, bool AllowRebuild = true,
       unsigned MaxLines = 0);
   void RealizeTopLevelDeclsFromPreamble();
 
public:
   /// were originally used to produce this translation unit.
   ///
   /// \param VFS - A llvm::vfs::FileSystem to be used for all file accesses.
   /// Note that preamble is saved to a temporary directory on a RealFileSystem,
   /// so in order for it to be loaded correctly, VFS should give an access to
   /// this(i.e. be an overlay over RealFileSystem).
   /// FileMgr->getVirtualFileSystem() will be used if \p VFS is nullptr.
   ///
+  /// \param KeepSourceMgr - Keep the current SourceMgr object instead of
+  /// creating a new one.  This is useful if you replaced the file buffers with
+  /// new code and want to reparse the source with this buffer.
+  ///
   /// \returns True if a failure occurred that causes the ASTUnit not to
   /// contain any translation-unit information, false otherwise.
   bool Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
                ArrayRef<RemappedFile> RemappedFiles = None,
-               IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS = nullptr);
+               IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS = nullptr,
+               bool KeepSourceMgr = false);
 
   /// Free data that will be re-generated on the next parse.
   ///
   /// Preamble-related data is not affected.
-  void ResetForParse();
+  ///
+  /// \param KeepSourceMgr - Keep the SourceManager object of this class.
+  void ResetForParse(bool KeepSourceMgr = false);
 
   /// Perform code completion at the given file, line, and
   /// column within this translation unit.
   ///
   /// \param File The file in which code completion will occur.
   ///
   /// \param Line The line at which code completion will occur.
   ///

clang/lib/Frontend/ASTUnit.cpp

 
 /// Parse the source file into a translation unit using the given compiler
 /// invocation, replacing the current translation unit.
 ///
 /// \returns True if a failure occurred that causes the ASTUnit not to
 /// contain any translation-unit information, false otherwise.
 bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
                     std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
-                    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS) {
+                    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+                    bool KeepSourceMgr) {
+
   if (!Invocation)
     return true;
 
   if (VFS && FileMgr)
     assert(VFS == &FileMgr->getVirtualFileSystem() &&
            "VFS passed to Parse and VFS in FileMgr are different");
 
   auto CCInvocation = std::make_shared<CompilerInvocation>(*Invocation);
auto CleanOnError = llvm::make_scope_exit([&]() {
     FailedParseDiagnostics.swap(StoredDiagnostics);
     StoredDiagnostics.clear();
     NumStoredDiagnosticsFromDriver = 0;
   });
 
   // Ensure that Clang has a FileManager with the right VFS, which may have
   // changed above in AddImplicitPreamble.  If VFS is nullptr, rely on
   // createFileManager to create one.
-  if (VFS && FileMgr && &FileMgr->getVirtualFileSystem() == VFS)
+  if (FileMgr &&
+      (KeepSourceMgr || (VFS && &FileMgr->getVirtualFileSystem() == VFS)))
     Clang->setFileManager(&*FileMgr);
   else
     FileMgr = Clang->createFileManager(std::move(VFS));
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<CompilerInstance>
     CICleanup(Clang.get());
 
bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   assert(Clang->getFrontendOpts().Inputs[0].getKind().getLanguage() !=
              Language::LLVM_IR &&
          "IR inputs not support here!");
 
   // Configure the various subsystems.
   LangOpts = Clang->getInvocation().LangOpts;
   FileSystemOpts = Clang->getFileSystemOpts();
 
-  ResetForParse();
+  ResetForParse(KeepSourceMgr);
+
+  if (!KeepSourceMgr && !SourceMgr)
+    SourceMgr =
+        new SourceManager(getDiagnostics(), *FileMgr, UserFilesAreVolatile);
 
-  SourceMgr = new SourceManager(getDiagnostics(), *FileMgr,
-                                UserFilesAreVolatile);
   if (!OverrideMainBuffer) {
     checkAndRemoveNonDriverDiags(StoredDiagnostics);
     TopLevelDeclsInPreamble.clear();
   }
 
   // Create the source manager.
   Clang->setSourceManager(&getSourceManager());
 
if (AST->LoadFromCompilerInvocation(std::move(PCHContainerOps),
     return nullptr;
   }
 
   return AST.release();
 }
 
 bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
                       ArrayRef<RemappedFile> RemappedFiles,
-                      IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS) {
+                      IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+                      bool KeepSourceMgr) {
   if (!Invocation)
     return true;
 
   if (!VFS) {
     assert(FileMgr && "FileMgr is null on Reparse call");
     VFS = &FileMgr->getVirtualFileSystem();
   }
 
bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   for (const auto &RemappedFile : RemappedFiles) {
     Invocation->getPreprocessorOpts().addRemappedFile(RemappedFile.first,
                                                       RemappedFile.second);
   }
 
   // If we have a preamble file lying around, or if we might try to
   // build a precompiled preamble, do so now.
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
-  if (Preamble || PreambleRebuildCountdown > 0)
+  if (!KeepSourceMgr && (Preamble || PreambleRebuildCountdown > 0))
     OverrideMainBuffer =
         getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
 
   // Clear out the diagnostics state.
+  if (!KeepSourceMgr)
     FileMgr.reset();
+
   getDiagnostics().Reset();
   ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
   if (OverrideMainBuffer)
     getDiagnostics().setNumWarnings(NumWarningsInPreamble);
 
   // Parse the sources
-  bool Result =
-      Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer), VFS);
+  bool Result = Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer),
+                      VFS, KeepSourceMgr);
 
   // If we're caching global code-completion results, and the top-level
   // declarations have changed, clear out the code-completion cache.
   if (!Result && ShouldCacheCodeCompletionResults &&
       CurrentTopLevelHashValue != CompletionCacheTopLevelHashValue)
     CacheCodeCompletionResults();
 
   // We now need to clear out the completion info related to this translation
   // unit; it'll be recreated if necessary.
   CCTUInfo.reset();
 
   return Result;
 }
 
-void ASTUnit::ResetForParse() {
+void ASTUnit::ResetForParse(bool KeepSourceMgr) {
   SavedMainFileBuffer.reset();
 
+  if (!KeepSourceMgr)
     SourceMgr.reset();
+
   TheSema.reset();
   Ctx.reset();
   PP.reset();
   Reader.reset();
 
   TopLevelDecls.clear();
   clearFileLevelDecls();
 }

clang/unittests/Frontend/ASTUnitTest.cpp

ASTUnit *AST = ASTUnit::LoadFromCommandLine(
       nullptr);
 
   ASSERT_EQ(AST, nullptr);
   ASSERT_NE(ErrUnit, nullptr);
   ASSERT_TRUE(Diags->hasErrorOccurred());
   ASSERT_NE(ErrUnit->stored_diag_size(), 0U);
 }
 
+TEST_F(ASTUnitTest, ReparseWithOldSourceMgr) {
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> InMemoryFs =
+      new llvm::vfs::InMemoryFileSystem();
+  InMemoryFs->addFile("test.cpp", 0, llvm::MemoryBuffer::getMemBuffer(R"cpp(
+      void f() {}
+    )cpp"));
+
+  const char *Args[] = {"clang", "test.cpp"};
+  Diags = CompilerInstance::createDiagnostics(new DiagnosticOptions());
+  CreateInvocationOptions CIOpts;
+  CIOpts.Diags = Diags;
+  CInvok = createInvocation(Args, std::move(CIOpts));
+  ASSERT_TRUE(CInvok);
+
+  FileManager *FileMgr = new FileManager(FileSystemOptions(), InMemoryFs);
+  PCHContainerOps = std::make_shared<PCHContainerOperations>();
+
+  auto AU = ASTUnit::LoadFromCompilerInvocation(
+      CInvok, PCHContainerOps, Diags, FileMgr, false, CaptureDiagsKind::None, 1,
+      TU_Complete, false, false, false);
+
+  ASSERT_TRUE(AU);
+
+  const FileEntry *FEntry = AU->getSourceManager().getFileEntryForID(
+      AU->getSourceManager().getMainFileID());
+  std::string NewCode = "void g() {}";
+
+  AU->getSourceManager().overrideFileContents(
+      FEntry, llvm::MemoryBuffer::getMemBuffer(NewCode));
+
+  AU->Reparse(std::make_shared<PCHContainerOperations>(), None, nullptr,
+              /*KeepFileMgr=*/true);
+
+  auto FileMemBuf = AU->getSourceManager().getMemoryBufferForFileOrFake(FEntry);
+
+  // Verify that we still get the same memory buffer after Reparse.
+  EXPECT_TRUE(FileMemBuf.getBuffer() == NewCode);
+}
+
 } // anonymous namespace