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

[bkramer]

lg

include/clang/Frontend/ASTUnit.h

 class MemoryBufferCache;
 class Preprocessor;
 class PCHContainerOperations;
 class PCHContainerReader;
 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 {
 public:
   struct StandaloneFixIt {
     std::pair<unsigned, unsigned> RemoveRange;
     std::pair<unsigned, unsigned> InsertFromRange;
     std::string CodeToInsert;
private:
   
   ASTUnit(const ASTUnit &) = delete;
   void operator=(const ASTUnit &) = delete;
   
   explicit ASTUnit(bool MainFileIsAST);
 
   void CleanTemporaryFiles();
   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;
     std::unique_ptr<llvm::MemoryBuffer> Owner;
     unsigned Size;
     bool PreambleEndsAtStartOfLine;
     ComputedPreamble(llvm::MemoryBuffer *Buffer,
                      std::unique_ptr<llvm::MemoryBuffer> Owner, unsigned Size,
                      bool PreambleEndsAtStartOfLine)
         : Buffer(Buffer), Owner(std::move(Owner)), Size(Size),
           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();
 
   /// \brief Transfers ownership of the objects (like SourceManager) from
   /// \param CI to this ASTUnit.
   void transferASTDataFromCompilerInstance(CompilerInstance &CI);
 
   /// \brief Allows us to assert that ASTUnit is not being used concurrently,

 private:
   /// \brief Helper function for \c LoadFromCompilerInvocation() and
   /// \c LoadFromCommandLine(), which loads an AST from a compiler invocation.
   ///
   /// \param PrecompilePreambleAfterNParses After how many parses the preamble
   /// 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:
   
   /// \brief Create an ASTUnit from a source file, via a CompilerInvocation
   /// object, by invoking the optionally provided ASTFrontendAction. 
   ///
   /// \param CI - The compiler invocation to use; it must have exactly one input
   /// source file. The ASTUnit takes ownership of the CompilerInvocation object.
public:
   /// \param ResourceFilesPath - The path to the compiler resource files.
   ///
   /// \param ModuleFormat - If provided, uses the specific module format.
   ///
   /// \param ErrAST - If non-null and parsing failed without any AST to return
   /// (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(
       const char **ArgBegin, const char **ArgEnd,
       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
       IntrusiveRefCntPtr<DiagnosticsEngine> Diags, StringRef ResourceFilesPath,
       bool OnlyLocalDecls = false, bool CaptureDiagnostics = false,
       ArrayRef<RemappedFile> RemappedFiles = None,
       bool RemappedFilesKeepOriginalName = true,
       unsigned PrecompilePreambleAfterNParses = 0,
       TranslationUnitKind TUKind = TU_Complete,
       bool CacheCodeCompletionResults = false,
       bool IncludeBriefCommentsInCodeCompletion = false,
       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.
   ///
   /// \param File The file in which code completion will occur.
   ///
   /// \param Line The line at which code completion will occur.
   ///

include/clang/Frontend/CompilerInvocation.h

 namespace vfs {
   class FileSystem;
 }
 
 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

lib/Frontend/ASTUnit.cpp

struct OnDiskData {
     void CleanTemporaryFiles();
 
     /// \brief Erase the preamble file.
     void CleanPreambleFile();
 
     /// \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() {
   static llvm::sys::SmartMutex<false> M(/* recursive = */ true);
   return M;
 }
 
 static void cleanupOnDiskMapAtExit();

 }
 
 /// 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) {
+                    std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
+                    IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   SavedMainFileBuffer.reset();
 
   if (!Invocation)
     return true;
 
   // 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>
     CICleanup(Clang.get());
 
   Clang->setInvocation(std::make_shared<CompilerInvocation>(*Invocation));
   OriginalSourceFile = Clang->getFrontendOpts().Inputs[0].getFile();
     
static std::string GetPreamblePCHPath() {
 
   return Path.str();
 }
 
 /// \brief Compute the preamble for the main file, providing the source buffer
 /// 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();
   
   // Try to determine if the main file has been remapped, either from the 
   // command line (to another file) or directly through the compiler invocation
   // (to a memory buffer).
   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();
           Buffer = const_cast<llvm::MemoryBuffer *>(RB.second);
         }
       }
     }
   }
   
   // 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);
   }
 
   if (!Buffer)
     Buffer = BufferOwner.get();
   auto Pre = Lexer::ComputePreamble(Buffer->getBuffer(),
                                     *Invocation.getLangOpts(), MaxLines);

 /// can occur within the preamble.
 ///
 /// \returns If the precompiled preamble can be used, returns a newly-allocated
 /// buffer that should be used in place of the main file when doing so.
 /// Otherwise, returns a NULL pointer.
 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);
   FrontendOptions &FrontendOpts = PreambleInvocation->getFrontendOpts();
   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
     // preamble, if we have one. It's obviously no good any more.
     Preamble.clear();
     erasePreambleFile(this);
 
     // The next time we actually see a preamble, precompile it.
if (Preamble.size() == NewPreamble.Size &&
       // First, make a record of those files that have been overridden via
       // remapping or unsaved_files.
       std::map<llvm::sys::fs::UniqueID, PreambleFileHash> OverriddenFiles;
       for (const auto &R : PreprocessorOpts.RemappedFiles) {
         if (AnyFileChanged)
           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;
           break;
         }
 
         OverriddenFiles[Status.getUniqueID()] = PreambleFileHash::createForFile(
             Status.getSize(),
             llvm::sys::toTimeT(Status.getLastModificationTime()));
       }
 
       for (const auto &RB : PreprocessorOpts.RemappedFileBuffers) {
         if (AnyFileChanged)
           break;
 
         vfs::Status Status;
-        if (FileMgr->getNoncachedStatValue(RB.first, Status)) {
+        if (!moveOnNoError(VFS->status(RB.first), Status)) {
           AnyFileChanged = true;
           break;
         }
 
         OverriddenFiles[Status.getUniqueID()] =
             PreambleFileHash::createForMemoryBuffer(RB.second);
       }
        
       // Check whether anything has changed.
       for (llvm::StringMap<PreambleFileHash>::iterator
              F = FilesInPreamble.begin(), FEnd = FilesInPreamble.end();
            !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;
         }
 
         std::map<llvm::sys::fs::UniqueID, PreambleFileHash>::iterator Overridden
           = OverriddenFiles.find(Status.getUniqueID());
         if (Overridden != OverriddenFiles.end()) {
ASTUnit::getMainBufferWithPrecompiledPreamble(
   // Clear out old caches and data.
   getDiagnostics().Reset();
   ProcessWarningOptions(getDiagnostics(), Clang->getDiagnosticOpts());
   checkAndRemoveNonDriverDiags(StoredDiagnostics);
   TopLevelDecls.clear();
   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()));
 
   auto PreambleDepCollector = std::make_shared<DependencyCollector>();
   Clang->addDependencyCollector(PreambleDepCollector);
 
   std::unique_ptr<PrecompilePreambleAction> Act;
ASTUnit *ASTUnit::LoadFromCompilerInvocationAction(
   if (OwnAST)
     return OwnAST.release();
   else
     return AST;
 }
 
 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;
   getDiagnostics().Reset();
   ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
 
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
   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(
     std::shared_ptr<CompilerInvocation> CI,
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     IntrusiveRefCntPtr<DiagnosticsEngine> Diags, FileManager *FileMgr,
     bool OnlyLocalDecls, bool CaptureDiagnostics,
     unsigned PrecompilePreambleAfterNParses, TranslationUnitKind TUKind,
std::unique_ptr<ASTUnit> ASTUnit::LoadFromCompilerInvocation(
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<ASTUnit>
     ASTUnitCleanup(AST.get());
   llvm::CrashRecoveryContextCleanupRegistrar<DiagnosticsEngine,
     llvm::CrashRecoveryContextReleaseRefCleanup<DiagnosticsEngine> >
     DiagCleanup(Diags.get());
 
   if (AST->LoadFromCompilerInvocation(std::move(PCHContainerOps),
-                                      PrecompilePreambleAfterNParses))
+                                      PrecompilePreambleAfterNParses,
+                                      AST->FileMgr->getVirtualFileSystem()))
     return nullptr;
   return AST;
 }
 
 ASTUnit *ASTUnit::LoadFromCommandLine(
     const char **ArgBegin, const char **ArgEnd,
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     IntrusiveRefCntPtr<DiagnosticsEngine> Diags, StringRef ResourceFilesPath,
     bool OnlyLocalDecls, bool CaptureDiagnostics,
     ArrayRef<RemappedFile> RemappedFiles, bool RemappedFilesKeepOriginalName,
     unsigned PrecompilePreambleAfterNParses, TranslationUnitKind TUKind,
     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;
 
   std::shared_ptr<CompilerInvocation> CI;
 
   {
 
if (ModuleFormat)
     CI->getHeaderSearchOpts().ModuleFormat = ModuleFormat.getValue();
 
   // Create the AST unit.
   std::unique_ptr<ASTUnit> AST;
   AST.reset(new ASTUnit(false));
   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);
   AST->PCMCache = new MemoryBufferCache;
   AST->OnlyLocalDecls = OnlyLocalDecls;
   AST->CaptureDiagnostics = CaptureDiagnostics;
   AST->TUKind = TUKind;
   AST->ShouldCacheCodeCompletionResults = CacheCodeCompletionResults;
ASTUnit *ASTUnit::LoadFromCommandLine(
   CI = nullptr;
   Diags = nullptr;
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<ASTUnit>
     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) {
       AST->StoredDiagnostics.swap(AST->FailedParseDiagnostics);
       ErrAST->swap(AST);
     }
     return nullptr;
   }
 
   return AST.release();
 }
 
 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);
   ParsingTimer.setOutput("Reparsing " + getMainFileName());
 
   // Remap files.
   PreprocessorOptions &PPOpts = Invocation->getPreprocessorOpts();
   for (const auto &RB : PPOpts.RemappedFileBuffers)
     delete RB.second;
 
   Invocation->getPreprocessorOpts().clearRemappedFiles();
   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 (!getPreambleFile(this).empty() || PreambleRebuildCounter > 0)
     OverrideMainBuffer =
-        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation);
+        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
+
 
   // Clear out the diagnostics state.
   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));
+      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.
   if (!Result && ShouldCacheCodeCompletionResults &&
       CurrentTopLevelHashValue != CompletionCacheTopLevelHashValue)
     CacheCodeCompletionResults();
 
   // We now need to clear out the completion info related to this translation
void ASTUnit::CodeComplete(
 
   // If we have a precompiled preamble, try to use it. We only allow
   // the use of the precompiled preamble if we're if the completion
   // point is within the main file, after the end of the precompiled
   // preamble.
   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);
       }
     }
   }
 
   // If the main file has been overridden due to the use of a preamble,
   // make that override happen and introduce the preamble.
   if (OverrideMainBuffer) {
     PreprocessorOpts.addRemappedFile(OriginalSourceFile,

lib/Frontend/CompilerInvocation.cpp

void BuryPointer(const void *Ptr) {
   if (Idx >= kGraveYardMaxSize)
     return;
   GraveYard[Idx] = 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>();
     }
 
     IntrusiveRefCntPtr<vfs::FileSystem> FS = vfs::getVFSFromYAML(
         std::move(Buffer.get()), /*DiagHandler*/ nullptr, File);
     if (!FS.get()) {
       Diags.Report(diag::err_invalid_vfs_overlay) << File;
       return IntrusiveRefCntPtr<vfs::FileSystem>();
     }
     Overlay->pushOverlay(FS);
   }
   return Overlay;
 }
 } // end namespace clang