[NFC][CAS] Cleanup CAS-based file system initialization

Author: cachemeifyoucan

include/swift/Frontend/CachingUtils.h

@@ -56,8 +56,11 @@ llvm::Error storeCachedCompilerOutput(llvm::cas::ObjectStore &CAS,
                                       file_types::ID OutputKind);
 
 llvm::Expected<llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
-createCASFileSystem(llvm::cas::ObjectStore &CAS, ArrayRef<std::string> FSRoots,
-                    ArrayRef<std::string> IncludeTreeRoots);
+createCASFileSystem(llvm::cas::ObjectStore &CAS, ArrayRef<std::string> FSRoots);
+
+llvm::Expected<llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
+createIncludeTreeFileSystem(llvm::cas::ObjectStore &CAS,
+                            ArrayRef<std::string> IncludeTreeRoots);
 
 namespace cas {
 /// Helper class to manage CAS/Caching from libSwiftScan C APIs.

lib/ClangImporter/ClangImporter.cpp

@@ -1166,8 +1166,10 @@ ClangImporter::create(ASTContext &ctx,
   auto fileMapping = getClangInvocationFileMapping(ctx);
   // Wrap Swift's FS to allow Clang to override the working directory
   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS =
-      llvm::vfs::RedirectingFileSystem::create(
-          fileMapping.redirectedFiles, true, *ctx.SourceMgr.getFileSystem());
+      ctx.SourceMgr.getFileSystem();
+  if (!fileMapping.redirectedFiles.empty())
+    VFS = llvm::vfs::RedirectingFileSystem::create(
+        fileMapping.redirectedFiles, true, *ctx.SourceMgr.getFileSystem());
   if (!fileMapping.overridenFiles.empty()) {
     llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> overridenVFS =
         new llvm::vfs::InMemoryFileSystem();

lib/Frontend/CachingUtils.cpp

@@ -377,29 +377,35 @@ static Expected<ObjectRef> mergeCASFileSystem(ObjectStore &CAS,
 }
 
 Expected<IntrusiveRefCntPtr<vfs::FileSystem>>
-createCASFileSystem(ObjectStore &CAS, ArrayRef<std::string> FSRoots,
-                    ArrayRef<std::string> IncludeTrees) {
-  assert(!FSRoots.empty() || !IncludeTrees.empty() && "no root ID provided");
-  if (FSRoots.size() == 1 && IncludeTrees.empty()) {
+createCASFileSystem(ObjectStore &CAS, ArrayRef<std::string> FSRoots) {
+  if (FSRoots.empty())
+    return nullptr;
+
+  if (FSRoots.size() == 1) {
     auto ID = CAS.parseID(FSRoots.front());
     if (!ID)
       return ID.takeError();
-    auto Ref = CAS.getReference(*ID);
-    if (!Ref)
-      return createCASObjectNotFoundError(*ID);
+
+    return createCASFileSystem(CAS, *ID);
   }
 
   auto NewRoot = mergeCASFileSystem(CAS, FSRoots);
   if (!NewRoot)
     return NewRoot.takeError();
 
-  auto FS = createCASFileSystem(CAS, CAS.getID(*NewRoot));
-  if (!FS)
-    return FS.takeError();
+  return createCASFileSystem(CAS, CAS.getID(*NewRoot));
+}
 
-  auto CASFS = makeIntrusiveRefCnt<vfs::OverlayFileSystem>(std::move(*FS));
-  // Push all Include File System onto overlay.
-  for (auto &Tree : IncludeTrees) {
+Expected<IntrusiveRefCntPtr<vfs::FileSystem>>
+createIncludeTreeFileSystem(ObjectStore &CAS,
+                            ArrayRef<std::string> IncludeTreeRoots) {
+  if (IncludeTreeRoots.empty())
+    return nullptr;
+
+  IntrusiveRefCntPtr<clang::cas::IncludeTreeFileSystem> IncludeTreeFS =
+      new clang::cas::IncludeTreeFileSystem(CAS);
+
+  for (const auto &Tree : IncludeTreeRoots) {
     auto ID = CAS.parseID(Tree);
     if (!ID)
       return ID.takeError();
@@ -411,13 +417,11 @@ createCASFileSystem(ObjectStore &CAS, ArrayRef<std::string> FSRoots,
     if (!IT)
       return IT.takeError();
 
-    auto ITFS = clang::cas::createIncludeTreeFileSystem(*IT);
-    if (!ITFS)
-      return ITFS.takeError();
-    CASFS->pushOverlay(std::move(*ITFS));
-  }
+    if (auto E = IncludeTreeFS->addFilesFromIncludeTree(*IT))
+      return std::move(E);
+  };
 
-  return CASFS;
+  return IncludeTreeFS;
 }
 
 namespace cas {

lib/Frontend/Frontend.cpp

@@ -523,16 +523,33 @@ bool CompilerInstance::setUpVirtualFileSystemOverlays() {
   if (Invocation.getFrontendOptions().EnableCaching &&
       (!Invocation.getFrontendOptions().CASFSRootIDs.empty() ||
        !Invocation.getFrontendOptions().ClangIncludeTrees.empty())) {
-    // Set up CASFS as BaseFS.
+    // Set up FileSystem from CAS.
     const auto &Opts = getInvocation().getFrontendOptions();
-    auto FS =
-        createCASFileSystem(*CAS, Opts.CASFSRootIDs, Opts.ClangIncludeTrees);
-    if (!FS) {
+    auto CASFS =
+        createCASFileSystem(*CAS, Opts.CASFSRootIDs);
+    if (!CASFS) {
       Diagnostics.diagnose(SourceLoc(), diag::error_cas,
-                           toString(FS.takeError()));
+                           toString(CASFS.takeError()));
       return true;
     }
-    SourceMgr.setFileSystem(std::move(*FS));
+    auto IncludeTreeFS =
+        createIncludeTreeFileSystem(*CAS, Opts.ClangIncludeTrees);
+    if (!IncludeTreeFS) {
+      Diagnostics.diagnose(SourceLoc(), diag::error_cas,
+                           toString(IncludeTreeFS.takeError()));
+      return true;
+    }
+    // If both FS exists, overlay them, otherwise, use one or the other as
+    // BaseFS.
+    if (*CASFS && *IncludeTreeFS) {
+      llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> OverlayVFS =
+          new llvm::vfs::OverlayFileSystem(std::move(*CASFS));
+      OverlayVFS->pushOverlay(*IncludeTreeFS);
+      SourceMgr.setFileSystem(std::move(OverlayVFS));
+    } else if (*CASFS)
+      SourceMgr.setFileSystem(std::move(*CASFS));
+    else
+      SourceMgr.setFileSystem(*IncludeTreeFS);
   }
 
   // If we have a bridging header cache key, try load it now and overlay it.