[SourceKit] Make sure to use overlays when reading files

Two paths missed setting up overlays:
  - `CompletionInstance` when checking files from dependencies
  - `SwiftASTManager` when reading in files that it would later replace
    all inputs with

(1) would cause the AST context not to be re-used, even though nothing
had changed. (2) caused all non-completion functionality to fail for any
symbols within files only specified by the overlay.

Resolves rdar://85508213.

Author: bnbarham

include/swift/AST/SearchPathOptions.h

@@ -16,7 +16,9 @@
 #include "swift/Basic/ArrayRefView.h"
 #include "swift/Basic/PathRemapper.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/VirtualFileSystem.h"
 
 #include <string>
@@ -357,6 +359,14 @@ class SearchPathOptions {
     return Lookup.searchPathsContainingFile(this, Filenames, FS, IsOSDarwin);
   }
 
+  /// Creates a filesystem taking into account any overlays specified in
+  /// \c VFSOverlayFiles. Returns \p BaseFS if there were no overlays and
+  /// \c FileError(s) if any error occurred while attempting to parse the
+  /// overlay files.
+  llvm::Expected<llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
+  makeOverlayFileSystem(
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> BaseFS) const;
+
 private:
   static StringRef
   pathStringFromFrameworkSearchPath(const FrameworkSearchPath &next) {

include/swift/IDE/CompletionInstance.h

@@ -109,6 +109,7 @@ class CompletionInstance {
   bool performCachedOperationIfPossible(
       llvm::hash_code ArgsHash,
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+      const SearchPathOptions &SearchPathOpts,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC,
       std::shared_ptr<std::atomic<bool>> CancellationFlag,

lib/AST/SearchPathOptions.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/AST/SearchPathOptions.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/Support/Errc.h"
 
 using namespace swift;
 
@@ -114,3 +115,47 @@ ModuleSearchPathLookup::searchPathsContainingFile(
                         const ModuleSearchPath *Rhs) { return *Lhs < *Rhs; });
   return Result;
 }
+
+/// Loads a VFS YAML file located at \p File using \p BaseFS and adds it to
+/// \p OverlayFS. Returns an error if either loading the \p File failed or it
+/// is invalid.
+static llvm::Error loadAndValidateVFSOverlay(
+    const std::string &File,
+    const llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> &BaseFS,
+    const llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> &OverlayFS) {
+  auto Buffer = BaseFS->getBufferForFile(File);
+  if (!Buffer)
+    return llvm::createFileError(File, Buffer.getError());
+
+  auto VFS = llvm::vfs::getVFSFromYAML(std::move(Buffer.get()), nullptr, File);
+  if (!VFS)
+    return llvm::createFileError(File, llvm::errc::invalid_argument);
+
+  OverlayFS->pushOverlay(std::move(VFS));
+  return llvm::Error::success();
+}
+
+llvm::Expected<llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
+SearchPathOptions::makeOverlayFileSystem(
+    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> BaseFS) const {
+  // TODO: This implementation is different to how Clang reads overlays in.
+  // Expose a helper in Clang rather than doing this ourselves.
+
+  auto OverlayFS =
+      llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(BaseFS);
+
+  llvm::Error AllErrors = llvm::Error::success();
+  bool hasOverlays = false;
+  for (const auto &File : VFSOverlayFiles) {
+    hasOverlays = true;
+    if (auto Err = loadAndValidateVFSOverlay(File, BaseFS, OverlayFS))
+      AllErrors = llvm::joinErrors(std::move(AllErrors), std::move(Err));
+  }
+
+  if (AllErrors)
+    return std::move(AllErrors);
+
+  if (hasOverlays)
+    return OverlayFS;
+  return BaseFS;
+}

lib/Frontend/Frontend.cpp

@@ -405,49 +405,28 @@ bool CompilerInstance::setup(const CompilerInvocation &Invok,
   return false;
 }
 
-static bool loadAndValidateVFSOverlay(
-    const std::string &File,
-    const llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> &BaseFS,
-    const llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> &OverlayFS,
-    DiagnosticEngine &Diag) {
-  auto Buffer = BaseFS->getBufferForFile(File);
-  if (!Buffer) {
-    Diag.diagnose(SourceLoc(), diag::cannot_open_file, File,
-                         Buffer.getError().message());
+bool CompilerInstance::setUpVirtualFileSystemOverlays() {
+  auto ExpectedOverlay =
+      Invocation.getSearchPathOptions().makeOverlayFileSystem(
+          SourceMgr.getFileSystem());
+  if (!ExpectedOverlay) {
+    llvm::handleAllErrors(
+        ExpectedOverlay.takeError(), [&](const llvm::FileError &FE) {
+          if (FE.convertToErrorCode() == std::errc::no_such_file_or_directory) {
+            Diagnostics.diagnose(SourceLoc(), diag::cannot_open_file,
+                                 FE.getFileName(), FE.messageWithoutFileInfo());
+          } else {
+            Diagnostics.diagnose(SourceLoc(), diag::invalid_vfs_overlay_file,
+                                 FE.getFileName());
+          }
+        });
     return true;
   }
 
-  auto VFS = llvm::vfs::getVFSFromYAML(std::move(Buffer.get()),
-                                        nullptr, File);
-  if (!VFS) {
-    Diag.diagnose(SourceLoc(), diag::invalid_vfs_overlay_file, File);
-    return true;
-  }
-  OverlayFS->pushOverlay(std::move(VFS));
+  SourceMgr.setFileSystem(*ExpectedOverlay);
   return false;
 }
 
-bool CompilerInstance::setUpVirtualFileSystemOverlays() {
-  auto BaseFS = SourceMgr.getFileSystem();
-  auto OverlayFS = llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem>(
-                    new llvm::vfs::OverlayFileSystem(BaseFS));
-  bool hadAnyFailure = false;
-  bool hasOverlays = false;
-  for (const auto &File : Invocation.getSearchPathOptions().VFSOverlayFiles) {
-    hasOverlays = true;
-    hadAnyFailure |=
-        loadAndValidateVFSOverlay(File, BaseFS, OverlayFS, Diagnostics);
-  }
-
-  // If we successfully loaded all the overlays, let the source manager and
-  // diagnostic engine take advantage of the overlay file system.
-  if (!hadAnyFailure && hasOverlays) {
-    SourceMgr.setFileSystem(OverlayFS);
-  }
-
-  return hadAnyFailure;
-}
-
 void CompilerInstance::setUpLLVMArguments() {
   // Honor -Xllvm.
   if (!Invocation.getFrontendOptions().LLVMArgs.empty()) {

lib/IDE/CompletionInstance.cpp

@@ -177,6 +177,7 @@ static DeclContext *getEquivalentDeclContextFromSourceFile(DeclContext *DC,
 bool CompletionInstance::performCachedOperationIfPossible(
     llvm::hash_code ArgsHash,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+    const SearchPathOptions &SearchPathOpts,
     llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
     DiagnosticConsumer *DiagC,
     std::shared_ptr<std::atomic<bool>> CancellationFlag,
@@ -210,6 +211,16 @@ bool CompletionInstance::performCachedOperationIfPossible(
     return false;
 
   if (shouldCheckDependencies()) {
+    // The passed in FileSystem does not have any overlays resolved. Make sure
+    // to do so before checking dependencies (otherwise we might decide we need
+    // to run the slow path due to a missing/different file).
+    auto ExpectedOverlay = SearchPathOpts.makeOverlayFileSystem(FileSystem);
+    if (ExpectedOverlay) {
+      FileSystem = std::move(ExpectedOverlay.get());
+    } else {
+      llvm::consumeError(ExpectedOverlay.takeError());
+    }
+
     if (areAnyDependentFilesInvalidated(
             CI, *FileSystem, *oldSF->getBufferID(),
             DependencyCheckedTimestamp, InMemoryDependencyHash))
@@ -543,9 +554,10 @@ void swift::ide::CompletionInstance::performOperation(
   // the cached completion instance.
   std::lock_guard<std::mutex> lock(mtx);
 
-  if (performCachedOperationIfPossible(ArgsHash, FileSystem, completionBuffer,
-                                       Offset, DiagC, CancellationFlag,
-                                       Callback)) {
+  if (performCachedOperationIfPossible(ArgsHash, FileSystem,
+                                       Invocation.getSearchPathOptions(),
+                                       completionBuffer, Offset, DiagC,
+                                       CancellationFlag, Callback)) {
     // We were able to reuse a cached AST. Callback has already been invoked
     // and we don't need to build a new AST. We are done.
     return;

test/SourceKit/CodeComplete/complete_swift_overlay.swift

@@ -0,0 +1,44 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+// RUN: sed -e "s@NAME_DIR@%{/t:regex_replacement}/A@g" -e "s@EXTERNAL_DIR@%{/t:regex_replacement}/B@g" %t/base.yaml > %t/overlay.yaml
+
+// RUN: %sourcekitd-test \
+// RUN:   -req=global-config -req-opts=completion_check_dependency_interval=0 == \
+// RUN:   -req=complete -pos=2:5 %t/A/a.swift -- -vfsoverlay %t/overlay.yaml %t/A/a.swift %t/A/b.swift == \
+// RUN:   -req=complete -pos=2:5 %t/A/a.swift -- -vfsoverlay %t/overlay.yaml %t/A/a.swift %t/A/b.swift \
+// RUN:   | %FileCheck %s
+
+// CHECK-LABEL: key.results: [
+// CHECK-DAG: key.name: "method()"
+// CHECK: ]
+// CHECK-NOT: key.reusingastcontext: 1
+
+// CHECK-LABEL: key.results: [
+// CHECK-DAG: key.name: "method()"
+// CHECK: ]
+// CHECK: key.reusingastcontext: 1
+
+
+//--- A/a.swift
+func a(b: B) {
+  b.
+}
+
+//--- B/b.swift
+struct B {
+  func method() {}
+}
+
+//--- base.yaml
+{
+  version: 0,
+  redirecting-with: "fallback",
+  use-external-names: true,
+  roots: [
+    {
+      type: "directory-remap",
+      name: "NAME_DIR",
+      external-contents: "EXTERNAL_DIR"
+    }
+  ]
+}

test/SourceKit/CursorInfo/cursor_swift_overlay.swift

@@ -0,0 +1,31 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file --leading-lines %s %t
+// RUN: sed -e "s@NAME_DIR@%{/t:regex_replacement}/A@g" -e "s@EXTERNAL_DIR@%{/t:regex_replacement}/B@g" %t/base.yaml > %t/overlay.yaml
+
+//--- A/a.swift
+func a() {
+// RUN: %sourcekitd-test -req=cursor -pos=%(line+1):3 %t/A/a.swift -- -vfsoverlay %t/overlay.yaml %t/A/a.swift %t/A/b.swift | %FileCheck %s
+  b()
+}
+
+//--- B/b.swift
+// TODO: This should be B/b.swift, but there's currently a bug with multiple overlays.
+//       See rdar://90578880 or https://github.com/llvm/llvm-project/issues/53306.
+//       Replace with CHECK-FIXED when that's fixed.
+// CHECK: source.lang.swift.ref.function.free ({{.*}}{{/|\\}}A{{/|\\}}b.swift:*missing file*)
+// CHECK-FIXED: source.lang.swift.ref.function.free ({{.*}}{{/|\\}}B{{/|\\}}b.swift:[[@LINE+1]]:6-[[@LINE+1]]:9)
+func b() {}
+
+//--- base.yaml
+{
+  version: 0,
+  redirecting-with: "fallback",
+  use-external-names: true,
+  roots: [
+    {
+      type: "directory-remap",
+      name: "NAME_DIR",
+      external-contents: "EXTERNAL_DIR"
+    }
+  ]
+}

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -1208,15 +1208,28 @@ void ASTProducer::enqueueConsumer(
   // this synchronously since all results will be delivered async anyway.
   auto This = shared_from_this();
   BuildOperationsQueue.dispatch([Consumer, FileSystem, Mgr, This]() {
+    // The passed in filesystem does not have overlays resolved. Make sure to
+    // do so before performing any file operations.
+    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS = FileSystem;
+    const InvocationOptions &InvocOpts = This->InvokRef->Impl.Opts;
+    const CompilerInvocation &ActualInvoc = InvocOpts.Invok;
+    auto ExpectedOverlay =
+        ActualInvoc.getSearchPathOptions().makeOverlayFileSystem(FileSystem);
+    if (ExpectedOverlay) {
+      FS = std::move(ExpectedOverlay.get());
+    } else {
+      llvm::consumeError(ExpectedOverlay.takeError());
+    }
+
     if (auto BuildOp =
-            This->getBuildOperationForConsumer(Consumer, FileSystem, Mgr)) {
+            This->getBuildOperationForConsumer(Consumer, FS, Mgr)) {
       bool WasAdded = BuildOp->addConsumer(Consumer);
       if (!WasAdded) {
         // The build operation was cancelled after the call to
         // getBuildOperationForConsumer but before the consumer could be
         // added. This should be an absolute edge case. Let's just try
         // again.
-        This->enqueueConsumer(Consumer, FileSystem, Mgr);
+        This->enqueueConsumer(Consumer, FS, Mgr);
       }
     } else {
       auto WeakThis = std::weak_ptr<ASTProducer>(This);
@@ -1229,8 +1242,9 @@ void ASTProducer::enqueueConsumer(
           Mgr->Impl.ASTCache.set(This->InvokRef->Impl.Key, This);
         }
       };
+
       ASTBuildOperationRef NewBuildOp = std::make_shared<ASTBuildOperation>(
-          FileSystem, This->InvokRef, Mgr, DidFinishCallback);
+          FS, This->InvokRef, Mgr, DidFinishCallback);
       This->BuildOperations.push_back(NewBuildOp);
       bool WasAdded = NewBuildOp->addConsumer(Consumer);
       assert(WasAdded && "Consumer wasn't added to a new build operation "