Allow to use vfs::FileSystem for file accesses inside ASTUnit.

Reviewers: bkramer, krasimir, arphaman, akyrtzi

Reviewed By: bkramer

Subscribers: klimek, cfe-commits

Differential Revision: https://reviews.llvm.org/D33397

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@303630 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 134 insertions, 48 deletions

diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index 46395cf6e8..7f70609efc 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/include/clang/Frontend/ASTUnit.h
@@ -59,6 +59,10 @@ class TargetInfo;
 class FrontendAction;
 class ASTDeserializationListener;
 
+namespace vfs {
+class FileSystem;
+}
+
 /// \brief Utility class for loading a ASTContext from an AST file.
 ///
 class ASTUnit : public ModuleLoader {
@@ -420,7 +424,8 @@ private:
   explicit ASTUnit(bool MainFileIsAST);
 
   bool Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-             std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer);
+             std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
+             IntrusiveRefCntPtr<vfs::FileSystem> VFS);
 
   struct ComputedPreamble {
     llvm::MemoryBuffer *Buffer;
@@ -434,11 +439,13 @@ private:
           PreambleEndsAtStartOfLine(PreambleEndsAtStartOfLine) {}
   };
   ComputedPreamble ComputePreamble(CompilerInvocation &Invocation,
-                                   unsigned MaxLines);
+                                   unsigned MaxLines,
+                                   IntrusiveRefCntPtr<vfs::FileSystem> VFS);
 
   std::unique_ptr<llvm::MemoryBuffer> getMainBufferWithPrecompiledPreamble(
       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-      const CompilerInvocation &PreambleInvocationIn, bool AllowRebuild = true,
+      const CompilerInvocation &PreambleInvocationIn,
+      IntrusiveRefCntPtr<vfs::FileSystem> VFS, bool AllowRebuild = true,
       unsigned MaxLines = 0);
   void RealizeTopLevelDeclsFromPreamble();
 
@@ -731,11 +738,17 @@ private:
   /// of this translation unit should be precompiled, to improve the performance
   /// of reparsing. Set to zero to disable preambles.
   ///
+  /// \param VFS - A 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 have access to it(i.e., be an
+  /// overlay over RealFileSystem).
+  ///
   /// \returns \c true if a catastrophic failure occurred (which means that the
   /// \c ASTUnit itself is invalid), or \c false otherwise.
   bool LoadFromCompilerInvocation(
       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-      unsigned PrecompilePreambleAfterNParses);
+      unsigned PrecompilePreambleAfterNParses,
+      IntrusiveRefCntPtr<vfs::FileSystem> VFS);
 
 public:
   
@@ -826,6 +839,11 @@ public:
   /// (e.g. because the PCH could not be loaded), this accepts the ASTUnit
   /// mainly to allow the caller to see the diagnostics.
   ///
+  /// \param VFS - A 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 have access to it(i.e., be an
+  /// overlay over RealFileSystem). RealFileSystem will be used if \p VFS is nullptr.
+  ///
   // FIXME: Move OnlyLocalDecls, UseBumpAllocator to setters on the ASTUnit, we
   // shouldn't need to specify them at construction time.
   static ASTUnit *LoadFromCommandLine(
@@ -842,15 +860,23 @@ public:
       bool AllowPCHWithCompilerErrors = false, bool SkipFunctionBodies = false,
       bool UserFilesAreVolatile = false, bool ForSerialization = false,
       llvm::Optional<StringRef> ModuleFormat = llvm::None,
-      std::unique_ptr<ASTUnit> *ErrAST = nullptr);
+      std::unique_ptr<ASTUnit> *ErrAST = nullptr,
+      IntrusiveRefCntPtr<vfs::FileSystem> VFS = nullptr);
 
   /// \brief Reparse the source files using the same command-line options that
   /// were originally used to produce this translation unit.
   ///
+  /// \param VFS - A 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.
+  ///
   /// \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);
+               ArrayRef<RemappedFile> RemappedFiles = None,
+               IntrusiveRefCntPtr<vfs::FileSystem> VFS = nullptr);
 
   /// \brief Perform code completion at the given file, line, and
   /// column within this translation unit.
diff --git a/include/clang/Frontend/CompilerInvocation.h b/include/clang/Frontend/CompilerInvocation.h
index cef7f73eca..8c4c932190 100644
--- a/include/clang/Frontend/CompilerInvocation.h
+++ b/include/clang/Frontend/CompilerInvocation.h
@@ -225,6 +225,11 @@ IntrusiveRefCntPtr<vfs::FileSystem>
 createVFSFromCompilerInvocation(const CompilerInvocation &CI,
                                 DiagnosticsEngine &Diags);
 
+IntrusiveRefCntPtr<vfs::FileSystem>
+createVFSFromCompilerInvocation(const CompilerInvocation &CI,
+                                DiagnosticsEngine &Diags,
+                                IntrusiveRefCntPtr<vfs::FileSystem> BaseFS);
+
 } // end namespace clang
 
 #endif
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index 6ee211c2de..d660638a1e 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -90,6 +90,21 @@ namespace {
     /// \brief Erase temporary files and the preamble file.
     void Cleanup();
   };
+
+  template <class T>
+  std::unique_ptr<T> valueOrNull(llvm::ErrorOr<std::unique_ptr<T>> Val) {
+    if (!Val)
+      return nullptr;
+    return std::move(*Val);
+  }
+
+  template <class T>
+  bool moveOnNoError(llvm::ErrorOr<T> Val, T &Output) {
+    if (!Val)
+      return false;
+    Output = std::move(*Val);
+    return true;
+  }
 }
 
 static llvm::sys::SmartMutex<false> &getOnDiskMutex() {
@@ -1019,7 +1034,8 @@ static void checkAndSanitizeDiags(SmallVectorImpl<StoredDiagnostic> &
 /// \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) {
+                    std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
+                    IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   SavedMainFileBuffer.reset();
 
   if (!Invocation)
@@ -1028,6 +1044,12 @@ bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   // Create the compiler instance to use for building the AST.
   std::unique_ptr<CompilerInstance> Clang(
       new CompilerInstance(std::move(PCHContainerOps)));
+  if (FileMgr && VFS) {
+    assert(VFS == FileMgr->getVirtualFileSystem() &&
+           "VFS passed to Parse and VFS in FileMgr are different");
+  } else if (VFS) {
+    Clang->setVirtualFileSystem(VFS);
+  }
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<CompilerInstance>
@@ -1170,7 +1192,8 @@ static std::string GetPreamblePCHPath() {
 /// that corresponds to the main file along with a pair (bytes, start-of-line)
 /// that describes the preamble.
 ASTUnit::ComputedPreamble
-ASTUnit::ComputePreamble(CompilerInvocation &Invocation, unsigned MaxLines) {
+ASTUnit::ComputePreamble(CompilerInvocation &Invocation, unsigned MaxLines,
+                         IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   FrontendOptions &FrontendOpts = Invocation.getFrontendOpts();
   PreprocessorOptions &PreprocessorOpts = Invocation.getPreprocessorOpts();
   
@@ -1180,28 +1203,32 @@ ASTUnit::ComputePreamble(CompilerInvocation &Invocation, unsigned MaxLines) {
   llvm::MemoryBuffer *Buffer = nullptr;
   std::unique_ptr<llvm::MemoryBuffer> BufferOwner;
   std::string MainFilePath(FrontendOpts.Inputs[0].getFile());
-  llvm::sys::fs::UniqueID MainFileID;
-  if (!llvm::sys::fs::getUniqueID(MainFilePath, MainFileID)) {
+  auto MainFileStatus = VFS->status(MainFilePath);
+  if (MainFileStatus) {
+    llvm::sys::fs::UniqueID MainFileID = MainFileStatus->getUniqueID();
+
     // Check whether there is a file-file remapping of the main file
     for (const auto &RF : PreprocessorOpts.RemappedFiles) {
       std::string MPath(RF.first);
-      llvm::sys::fs::UniqueID MID;
-      if (!llvm::sys::fs::getUniqueID(MPath, MID)) {
+      auto MPathStatus = VFS->status(MPath);
+      if (MPathStatus) {
+        llvm::sys::fs::UniqueID MID = MPathStatus->getUniqueID();
         if (MainFileID == MID) {
           // We found a remapping. Try to load the resulting, remapped source.
-          BufferOwner = getBufferForFile(RF.second);
+          BufferOwner = valueOrNull(VFS->getBufferForFile(RF.second));
           if (!BufferOwner)
             return ComputedPreamble(nullptr, nullptr, 0, true);
         }
       }
     }
-    
+
     // Check whether there is a file-buffer remapping. It supercedes the
     // file-file remapping.
     for (const auto &RB : PreprocessorOpts.RemappedFileBuffers) {
       std::string MPath(RB.first);
-      llvm::sys::fs::UniqueID MID;
-      if (!llvm::sys::fs::getUniqueID(MPath, MID)) {
+      auto MPathStatus = VFS->status(MPath);
+      if (MPathStatus) {
+        llvm::sys::fs::UniqueID MID = MPathStatus->getUniqueID();
         if (MainFileID == MID) {
           // We found a remapping.
           BufferOwner.reset();
@@ -1213,7 +1240,7 @@ ASTUnit::ComputePreamble(CompilerInvocation &Invocation, unsigned MaxLines) {
   
   // If the main source file was not remapped, load it now.
   if (!Buffer && !BufferOwner) {
-    BufferOwner = getBufferForFile(FrontendOpts.Inputs[0].getFile());
+    BufferOwner = valueOrNull(VFS->getBufferForFile(FrontendOpts.Inputs[0].getFile()));
     if (!BufferOwner)
       return ComputedPreamble(nullptr, nullptr, 0, true);
   }
@@ -1324,8 +1351,10 @@ makeStandaloneDiagnostic(const LangOptions &LangOpts,
 std::unique_ptr<llvm::MemoryBuffer>
 ASTUnit::getMainBufferWithPrecompiledPreamble(
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-    const CompilerInvocation &PreambleInvocationIn, bool AllowRebuild,
+    const CompilerInvocation &PreambleInvocationIn,
+    IntrusiveRefCntPtr<vfs::FileSystem> VFS, bool AllowRebuild,
     unsigned MaxLines) {
+  assert(VFS && "VFS is null");
 
   auto PreambleInvocation =
       std::make_shared<CompilerInvocation>(PreambleInvocationIn);
@@ -1333,7 +1362,8 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
   PreprocessorOptions &PreprocessorOpts
     = PreambleInvocation->getPreprocessorOpts();
 
-  ComputedPreamble NewPreamble = ComputePreamble(*PreambleInvocation, MaxLines);
+  ComputedPreamble NewPreamble =
+      ComputePreamble(*PreambleInvocation, MaxLines, VFS);
 
   if (!NewPreamble.Size) {
     // We couldn't find a preamble in the main source. Clear out the current
@@ -1369,7 +1399,7 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
           break;
 
         vfs::Status Status;
-        if (FileMgr->getNoncachedStatValue(R.second, Status)) {
+        if (!moveOnNoError(VFS->status(R.second), Status)) {
           // If we can't stat the file we're remapping to, assume that something
           // horrible happened.
           AnyFileChanged = true;
@@ -1386,7 +1416,7 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
           break;
 
         vfs::Status Status;
-        if (FileMgr->getNoncachedStatValue(RB.first, Status)) {
+        if (!moveOnNoError(VFS->status(RB.first), Status)) {
           AnyFileChanged = true;
           break;
         }
@@ -1401,7 +1431,7 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
            !AnyFileChanged && F != FEnd; 
            ++F) {
         vfs::Status Status;
-        if (FileMgr->getNoncachedStatValue(F->first(), Status)) {
+        if (!moveOnNoError(VFS->status(F->first()), Status)) {
           // If we can't stat the file, assume that something horrible happened.
           AnyFileChanged = true;
           break;
@@ -1546,14 +1576,14 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
   TopLevelDeclsInPreamble.clear();
   PreambleDiagnostics.clear();
 
-  IntrusiveRefCntPtr<vfs::FileSystem> VFS =
-      createVFSFromCompilerInvocation(Clang->getInvocation(), getDiagnostics());
+  VFS = createVFSFromCompilerInvocation(Clang->getInvocation(),
+                                        getDiagnostics(), VFS);
   if (!VFS)
     return nullptr;
 
   // Create a file manager object to provide access to and cache the filesystem.
   Clang->setFileManager(new FileManager(Clang->getFileSystemOpts(), VFS));
-  
+
   // Create the source manager.
   Clang->setSourceManager(new SourceManager(getDiagnostics(),
                                             Clang->getFileManager()));
@@ -1863,10 +1893,13 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocationAction(
 
 bool ASTUnit::LoadFromCompilerInvocation(
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-    unsigned PrecompilePreambleAfterNParses) {
+    unsigned PrecompilePreambleAfterNParses,
+    IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   if (!Invocation)
     return true;
-  
+
+  assert(VFS && "VFS is null");
+
   // We'll manage file buffers ourselves.
   Invocation->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   Invocation->getFrontendOpts().DisableFree = false;
@@ -1877,19 +1910,19 @@ bool ASTUnit::LoadFromCompilerInvocation(
   if (PrecompilePreambleAfterNParses > 0) {
     PreambleRebuildCounter = PrecompilePreambleAfterNParses;
     OverrideMainBuffer =
-        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation);
+        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
     getDiagnostics().Reset();
     ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
   }
-  
+
   SimpleTimer ParsingTimer(WantTiming);
   ParsingTimer.setOutput("Parsing " + getMainFileName());
-  
+
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<llvm::MemoryBuffer>
     MemBufferCleanup(OverrideMainBuffer.get());
 
-  return Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer));
+  return Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer), VFS);
 }
 
 std::unique_ptr<ASTUnit> ASTUnit::LoadFromCompilerInvocation(
@@ -1923,7 +1956,8 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromCompilerInvocation(
     DiagCleanup(Diags.get());
 
   if (AST->LoadFromCompilerInvocation(std::move(PCHContainerOps),
-                                      PrecompilePreambleAfterNParses))
+                                      PrecompilePreambleAfterNParses,
+                                      AST->FileMgr->getVirtualFileSystem()))
     return nullptr;
   return AST;
 }
@@ -1938,7 +1972,8 @@ ASTUnit *ASTUnit::LoadFromCommandLine(
     bool CacheCodeCompletionResults, bool IncludeBriefCommentsInCodeCompletion,
     bool AllowPCHWithCompilerErrors, bool SkipFunctionBodies,
     bool UserFilesAreVolatile, bool ForSerialization,
-    llvm::Optional<StringRef> ModuleFormat, std::unique_ptr<ASTUnit> *ErrAST) {
+    llvm::Optional<StringRef> ModuleFormat, std::unique_ptr<ASTUnit> *ErrAST,
+    IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   assert(Diags.get() && "no DiagnosticsEngine was provided");
 
   SmallVector<StoredDiagnostic, 4> StoredDiagnostics;
@@ -1979,8 +2014,9 @@ ASTUnit *ASTUnit::LoadFromCommandLine(
   ConfigureDiags(Diags, *AST, CaptureDiagnostics);
   AST->Diagnostics = Diags;
   AST->FileSystemOpts = CI->getFileSystemOpts();
-  IntrusiveRefCntPtr<vfs::FileSystem> VFS =
-      createVFSFromCompilerInvocation(*CI, *Diags);
+  if (!VFS)
+    VFS = vfs::getRealFileSystem();
+  VFS = createVFSFromCompilerInvocation(*CI, *Diags, VFS);
   if (!VFS)
     return nullptr;
   AST->FileMgr = new FileManager(AST->FileSystemOpts, VFS);
@@ -2006,7 +2042,8 @@ ASTUnit *ASTUnit::LoadFromCommandLine(
     ASTUnitCleanup(AST.get());
 
   if (AST->LoadFromCompilerInvocation(std::move(PCHContainerOps),
-                                      PrecompilePreambleAfterNParses)) {
+                                      PrecompilePreambleAfterNParses,
+                                      VFS)) {
     // Some error occurred, if caller wants to examine diagnostics, pass it the
     // ASTUnit.
     if (ErrAST) {
@@ -2020,10 +2057,16 @@ ASTUnit *ASTUnit::LoadFromCommandLine(
 }
 
 bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-                      ArrayRef<RemappedFile> RemappedFiles) {
+                      ArrayRef<RemappedFile> RemappedFiles,
+                      IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   if (!Invocation)
     return true;
 
+  if (!VFS) {
+    assert(FileMgr && "FileMgr is null on Reparse call");
+    VFS = FileMgr->getVirtualFileSystem();
+  }
+
   clearFileLevelDecls();
   
   SimpleTimer ParsingTimer(WantTiming);
@@ -2045,7 +2088,8 @@ bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
   if (!getPreambleFile(this).empty() || PreambleRebuildCounter > 0)
     OverrideMainBuffer =
-        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation);
+        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
+
 
   // Clear out the diagnostics state.
   FileMgr.reset();
@@ -2056,7 +2100,7 @@ bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
 
   // Parse the sources
   bool Result =
-      Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer));
+      Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer), VFS);
 
   // If we're caching global code-completion results, and the top-level 
   // declarations have changed, clear out the code-completion cache.
@@ -2414,15 +2458,19 @@ void ASTUnit::CodeComplete(
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
   if (!getPreambleFile(this).empty()) {
     std::string CompleteFilePath(File);
-    llvm::sys::fs::UniqueID CompleteFileID;
 
-    if (!llvm::sys::fs::getUniqueID(CompleteFilePath, CompleteFileID)) {
+    auto VFS = FileMgr.getVirtualFileSystem();
+    auto CompleteFileStatus = VFS->status(CompleteFilePath);
+    if (CompleteFileStatus) {
+      llvm::sys::fs::UniqueID CompleteFileID = CompleteFileStatus->getUniqueID();
+
       std::string MainPath(OriginalSourceFile);
-      llvm::sys::fs::UniqueID MainID;
-      if (!llvm::sys::fs::getUniqueID(MainPath, MainID)) {
+      auto MainStatus = VFS->status(MainPath);
+      if (MainStatus) {
+        llvm::sys::fs::UniqueID MainID = MainStatus->getUniqueID();
         if (CompleteFileID == MainID && Line > 1)
           OverrideMainBuffer = getMainBufferWithPrecompiledPreamble(
-              PCHContainerOps, Inv, false, Line - 1);
+              PCHContainerOps, Inv, VFS, false, Line - 1);
       }
     }
   }
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
index 29cdcd7df1..859e8ba18b 100644
--- a/lib/Frontend/CompilerInvocation.cpp
+++ b/lib/Frontend/CompilerInvocation.cpp
@@ -2750,15 +2750,22 @@ void BuryPointer(const void *Ptr) {
 IntrusiveRefCntPtr<vfs::FileSystem>
 createVFSFromCompilerInvocation(const CompilerInvocation &CI,
                                 DiagnosticsEngine &Diags) {
+  return createVFSFromCompilerInvocation(CI, Diags, vfs::getRealFileSystem());
+}
+
+IntrusiveRefCntPtr<vfs::FileSystem>
+createVFSFromCompilerInvocation(const CompilerInvocation &CI,
+                                DiagnosticsEngine &Diags,
+                                IntrusiveRefCntPtr<vfs::FileSystem> BaseFS) {
   if (CI.getHeaderSearchOpts().VFSOverlayFiles.empty())
-    return vfs::getRealFileSystem();
+    return BaseFS;
 
-  IntrusiveRefCntPtr<vfs::OverlayFileSystem>
-    Overlay(new vfs::OverlayFileSystem(vfs::getRealFileSystem()));
+  IntrusiveRefCntPtr<vfs::OverlayFileSystem> Overlay(
+      new vfs::OverlayFileSystem(BaseFS));
   // earlier vfs files are on the bottom
   for (const std::string &File : CI.getHeaderSearchOpts().VFSOverlayFiles) {
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> Buffer =
-        llvm::MemoryBuffer::getFile(File);
+        BaseFS->getBufferForFile(File);
     if (!Buffer) {
       Diags.Report(diag::err_missing_vfs_overlay_file) << File;
       return IntrusiveRefCntPtr<vfs::FileSystem>();