Merge pull request #64456 from bnbarham/nested-cursor

[SourceKit] Update requests to handle locations within generated buffers

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -636,7 +636,8 @@ mapOffsetToNewerSnapshot(unsigned Offset,
 static void mapLocToLatestSnapshot(
     SwiftLangSupport &Lang, LocationInfo &Location,
     ArrayRef<ImmutableTextSnapshotRef> PreviousASTSnaps) {
-  auto EditorDoc = Lang.getEditorDocuments()->findByPath(Location.Filename);
+  auto EditorDoc = Lang.getEditorDocuments()->findByPath(Location.Filename,
+                                                         /*IsRealpath=*/true);
   if (!EditorDoc)
     return;
 
@@ -1434,9 +1435,9 @@ public:
   CursorRangeInfoConsumer(StringRef InputFile, unsigned Offset, unsigned Length,
                           SwiftLangSupport &Lang, SwiftInvocationRef ASTInvok,
                           bool TryExistingAST, bool CancelOnSubsequentRequest)
-    : Lang(Lang), ASTInvok(ASTInvok),InputFile(InputFile.str()), Offset(Offset),
-      Length(Length), TryExistingAST(TryExistingAST),
-      CancelOnSubsequentRequest(CancelOnSubsequentRequest) {}
+      : Lang(Lang), ASTInvok(ASTInvok), InputFile(InputFile.str()),
+        Offset(Offset), Length(Length), TryExistingAST(TryExistingAST),
+        CancelOnSubsequentRequest(CancelOnSubsequentRequest) {}
 
   bool canUseASTWithSnapshots(ArrayRef<ImmutableTextSnapshotRef> Snapshots) override {
     if (!TryExistingAST) {
@@ -1451,7 +1452,8 @@ public:
     // blocked waiting on the AST to be fully typechecked.
 
     ImmutableTextSnapshotRef InputSnap;
-    if (auto EditorDoc = Lang.getEditorDocuments()->findByPath(InputFile))
+    if (auto EditorDoc = Lang.getEditorDocuments()->findByPath(
+            InputFile, /*IsRealpath=*/true))
       InputSnap = EditorDoc->getLatestSnapshot();
     if (!InputSnap)
       return false;
@@ -1492,6 +1494,33 @@ public:
   }
 };
 
+static SourceFile *retrieveInputFile(StringRef inputBufferName,
+                                     const CompilerInstance &CI,
+                                     bool haveRealPath = false) {
+  // Don't bother looking up if we have the same file as the primary file or
+  // we weren't given a separate input file
+  if (inputBufferName.empty() ||
+      CI.getPrimarySourceFile()->getFilename() == inputBufferName)
+    return CI.getPrimarySourceFile();
+
+  // Otherwise, try to find the given buffer identifier
+  const SourceManager &SM = CI.getSourceMgr();
+  Optional<unsigned> inputBufferID =
+      SM.getIDForBufferIdentifier(inputBufferName);
+  if (!inputBufferID) {
+    // If that failed, try again with symlinks resolved (unless we've already
+    // done that)
+    if (haveRealPath)
+      return nullptr;
+    std::string realPath =
+        SwiftLangSupport::resolvePathSymlinks(inputBufferName);
+    return retrieveInputFile(realPath, CI, /*haveRealPath=*/true);
+  }
+
+  return CI.getMainModule()->getSourceFileContainingLocation(
+      SM.getRangeForBuffer(*inputBufferID).getStart());
+}
+
 static void resolveCursor(
     SwiftLangSupport &Lang, StringRef InputFile, unsigned Offset,
     unsigned Length, bool Actionables, bool SymbolGraph,
@@ -1523,13 +1552,21 @@ static void resolveCursor(
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       auto &CompIns = AstUnit->getCompilerInstance();
+
+      SourceFile *SF = retrieveInputFile(InputFile, CompIns);
+      if (!SF) {
+        Receiver(RequestResult<CursorInfoData>::fromError(
+            "Unable to find input file"));
+        return;
+      }
+
       SourceManager &SM = CompIns.getSourceMgr();
-      unsigned BufferID = AstUnit->getPrimarySourceFile().getBufferID().value();
+      unsigned BufferID = SF->getBufferID().value();
       SourceLoc Loc =
         Lexer::getLocForStartOfToken(SM, BufferID, Offset);
       if (Loc.isInvalid()) {
         Receiver(RequestResult<CursorInfoData>::fromError(
-          "Unable to resolve the start of the token."));
+            "Unable to find initial lookup location"));
         return;
       }
 
@@ -1555,8 +1592,8 @@ static void resolveCursor(
         Range.Column = Pair.second;
         Range.Length = Length;
         bool CollectRangeStartRefactorings = false;
-        collectAvailableRefactorings(&AstUnit->getPrimarySourceFile(), Range,
-                                     CollectRangeStartRefactorings, Kinds, {});
+        collectAvailableRefactorings(SF, Range, CollectRangeStartRefactorings,
+                                     Kinds, {});
         for (RefactoringKind Kind : Kinds) {
           Actions.emplace_back(SwiftLangSupport::getUIDForRefactoringKind(Kind),
                                getDescriptiveRefactoringKindName(Kind),
@@ -1574,10 +1611,9 @@ static void resolveCursor(
         // Fall through to collect cursor based refactorings
       }
 
-      auto *File = &AstUnit->getPrimarySourceFile();
       ResolvedCursorInfoPtr CursorInfo =
-          evaluateOrDefault(File->getASTContext().evaluator,
-                            CursorInfoRequest{CursorInfoOwner(File, Loc)},
+          evaluateOrDefault(CompIns.getASTContext().evaluator,
+                            CursorInfoRequest{CursorInfoOwner(SF, Loc)},
                             new ResolvedCursorInfo());
 
       CompilerInvocation CompInvok;
@@ -1670,7 +1706,7 @@ static void resolveCursor(
 }
 
 static void computeDiagnostics(
-    SwiftLangSupport &Lang, StringRef InputFile, SwiftInvocationRef Invok,
+    SwiftLangSupport &Lang, SwiftInvocationRef Invok,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
     SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<DiagnosticsResult> &)> Receiver) {
@@ -1684,8 +1720,7 @@ static void computeDiagnostics(
         : Receiver(Receiver) {}
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
-      unsigned BufferID =
-          AstUnit->getPrimarySourceFile().getBufferID().value();
+      unsigned BufferID = *AstUnit->getPrimarySourceFile().getBufferID();
       auto &DiagConsumer = AstUnit->getEditorDiagConsumer();
       auto Diagnostics = DiagConsumer.getDiagnosticsForBuffer(BufferID);
       Receiver(RequestResult<DiagnosticsResult>::fromResult(Diagnostics));
@@ -1731,19 +1766,24 @@ static void resolveName(
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       auto &CompIns = AstUnit->getCompilerInstance();
 
-      unsigned BufferID = AstUnit->getPrimarySourceFile().getBufferID().value();
-      SourceLoc Loc =
-        Lexer::getLocForStartOfToken(CompIns.getSourceMgr(), BufferID, Offset);
+      SourceFile *SF = retrieveInputFile(InputFile, CompIns);
+      if (!SF) {
+        Receiver(RequestResult<NameTranslatingInfo>::fromError(
+            "Unable to find input file"));
+        return;
+      }
+
+      SourceLoc Loc = Lexer::getLocForStartOfToken(CompIns.getSourceMgr(),
+                                                   *SF->getBufferID(), Offset);
       if (Loc.isInvalid()) {
         Receiver(RequestResult<NameTranslatingInfo>::fromError(
           "Unable to resolve the start of the token."));
         return;
       }
 
-      auto *File = &AstUnit->getPrimarySourceFile();
       ResolvedCursorInfoPtr CursorInfo =
-          evaluateOrDefault(File->getASTContext().evaluator,
-                            CursorInfoRequest{CursorInfoOwner(File, Loc)},
+          evaluateOrDefault(CompIns.getASTContext().evaluator,
+                            CursorInfoRequest{CursorInfoOwner(SF, Loc)},
                             new ResolvedCursorInfo());
       if (CursorInfo->isInvalid()) {
         NameTranslatingInfo Info;
@@ -1835,10 +1875,19 @@ resolveRange(SwiftLangSupport &Lang, StringRef InputFile, unsigned Offset,
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       // FIXME: Implement tracing
-      auto *File = &AstUnit->getPrimarySourceFile();
-      ResolvedRangeInfo Info = evaluateOrDefault(File->getASTContext().evaluator,
-        RangeInfoRequest(RangeInfoOwner({File, Offset, Length})),
-                                                 ResolvedRangeInfo());
+      auto &CompIns = AstUnit->getCompilerInstance();
+
+      SourceFile *SF = retrieveInputFile(InputFile, CompIns);
+      if (!SF) {
+        Receiver(
+            RequestResult<RangeInfo>::fromError("Unable to find input file"));
+        return;
+      }
+
+      ResolvedRangeInfo Info = evaluateOrDefault(
+          CompIns.getASTContext().evaluator,
+          RangeInfoRequest(RangeInfoOwner({SF, Offset, Length})),
+          ResolvedRangeInfo());
 
       CompilerInvocation CompInvok;
       ASTInvok->applyTo(CompInvok);
@@ -1938,17 +1987,18 @@ static void deliverCursorInfoResults(
 }
 
 void SwiftLangSupport::getCursorInfo(
-    StringRef InputFile, unsigned Offset, unsigned Length, bool Actionables,
-    bool SymbolGraph, bool CancelOnSubsequentRequest,
-    ArrayRef<const char *> Args, Optional<VFSOptions> vfsOptions,
+    StringRef PrimaryFilePath, StringRef InputBufferName, unsigned Offset,
+    unsigned Length, bool Actionables, bool SymbolGraph,
+    bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
+    Optional<VFSOptions> vfsOptions,
     SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<CursorInfoData> &)> Receiver) {
   std::string error;
-  auto fileSystem = getFileSystem(vfsOptions, InputFile, error);
+  auto fileSystem = getFileSystem(vfsOptions, PrimaryFilePath, error);
   if (!fileSystem)
     return Receiver(RequestResult<CursorInfoData>::fromError(error));
 
-  if (auto IFaceGenRef = IFaceGenContexts.get(InputFile)) {
+  if (auto IFaceGenRef = IFaceGenContexts.get(PrimaryFilePath)) {
     IFaceGenRef->accessASTAsync([this, IFaceGenRef, Offset, Actionables,
                                  SymbolGraph, Receiver] {
       SwiftInterfaceGenContext::ResolvedEntity Entity;
@@ -1990,7 +2040,7 @@ void SwiftLangSupport::getCursorInfo(
 
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, InputFile, fileSystem, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, fileSystem, Error);
   if (!Error.empty()) {
     LOG_WARN_FUNC("error creating ASTInvocation: " << Error);
   }
@@ -2000,14 +2050,18 @@ void SwiftLangSupport::getCursorInfo(
   }
 
   bool SolverBasedProducedResult = false;
-  std::string InputFileError;
-  llvm::SmallString<64> RealInputFilePath;
-  fileSystem->getRealPath(InputFile, RealInputFilePath);
-  std::unique_ptr<llvm::MemoryBuffer> UnresolvedInputFile =
-      getASTManager()->getMemoryBuffer(RealInputFilePath, fileSystem,
-                                       InputFileError);
-  // The solver-based implementation doesn't support range based cursor info.
-  if (UnresolvedInputFile && Length == 0) {
+  // Solver based cursor info cannot handle generated buffers or range based
+  // cursor info.
+  std::unique_ptr<llvm::MemoryBuffer> InputBuffer;
+  if (InputBufferName.empty() && Length == 0) {
+    std::string InputFileError;
+    llvm::SmallString<128> RealInputFilePath;
+    fileSystem->getRealPath(PrimaryFilePath, RealInputFilePath);
+    InputBuffer = getASTManager()->getMemoryBuffer(RealInputFilePath,
+                                                   fileSystem, InputFileError);
+  }
+
+  if (InputBuffer) {
     auto SolverBasedReceiver = [&](const RequestResult<CursorInfoData> &Res) {
       SolverBasedProducedResult = true;
       Receiver(Res);
@@ -2017,7 +2071,7 @@ void SwiftLangSupport::getCursorInfo(
     Invok->applyTo(CompInvok);
 
     performWithParamsToCompletionLikeOperation(
-        UnresolvedInputFile.get(), Offset,
+        InputBuffer.get(), Offset,
         /*InsertCodeCompletionToken=*/false, Args, fileSystem,
         CancellationToken,
         [&](CancellableResult<CompletionLikeOperationParams> ParmsResult) {
@@ -2036,19 +2090,20 @@ void SwiftLangSupport::getCursorInfo(
   }
 
   if (!SolverBasedProducedResult) {
-    resolveCursor(*this, InputFile, Offset, Length, Actionables, SymbolGraph,
-                  Invok, /*TryExistingAST=*/true, CancelOnSubsequentRequest,
-                  fileSystem, CancellationToken, Receiver);
+    resolveCursor(*this, InputBufferName, Offset, Length, Actionables,
+                  SymbolGraph, Invok, /*TryExistingAST=*/true,
+                  CancelOnSubsequentRequest, fileSystem, CancellationToken,
+                  Receiver);
   }
 }
 
 void SwiftLangSupport::getDiagnostics(
-    StringRef InputFile, ArrayRef<const char *> Args,
+    StringRef PrimaryFilePath, ArrayRef<const char *> Args,
     Optional<VFSOptions> VfsOptions,
     SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<DiagnosticsResult> &)> Receiver) {
   std::string FileSystemError;
-  auto FileSystem = getFileSystem(VfsOptions, InputFile, FileSystemError);
+  auto FileSystem = getFileSystem(VfsOptions, PrimaryFilePath, FileSystemError);
   if (!FileSystem) {
     Receiver(RequestResult<DiagnosticsResult>::fromError(FileSystemError));
     return;
@@ -2056,7 +2111,7 @@ void SwiftLangSupport::getDiagnostics(
 
   std::string InvocationError;
   SwiftInvocationRef Invok = ASTMgr->getTypecheckInvocation(
-      Args, InputFile, FileSystem, InvocationError);
+      Args, PrimaryFilePath, FileSystem, InvocationError);
   if (!InvocationError.empty()) {
     LOG_WARN_FUNC("error creating ASTInvocation: " << InvocationError);
   }
@@ -2065,16 +2120,15 @@ void SwiftLangSupport::getDiagnostics(
     return;
   }
 
-  computeDiagnostics(*this, InputFile, Invok, FileSystem, CancellationToken,
-                     Receiver);
+  computeDiagnostics(*this, Invok, FileSystem, CancellationToken, Receiver);
 }
 
 void SwiftLangSupport::getRangeInfo(
-    StringRef InputFile, unsigned Offset, unsigned Length,
-    bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
-    SourceKitCancellationToken CancellationToken,
+    StringRef PrimaryFilePath, StringRef InputBufferName, unsigned Offset,
+    unsigned Length, bool CancelOnSubsequentRequest,
+    ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<RangeInfo> &)> Receiver) {
-  if (IFaceGenContexts.get(InputFile)) {
+  if (IFaceGenContexts.get(InputBufferName)) {
     // FIXME: return range info for generated interfaces.
     Receiver(RequestResult<RangeInfo>::fromError(
         "Range info for generated interfaces is not implemented."));
@@ -2082,7 +2136,7 @@ void SwiftLangSupport::getRangeInfo(
   }
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, InputFile, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<RangeInfo>::fromError(Error));
@@ -2092,8 +2146,9 @@ void SwiftLangSupport::getRangeInfo(
     Receiver(RequestResult<RangeInfo>::fromError("Invalid range length."));
     return;
   }
-  resolveRange(*this, InputFile, Offset, Length, Invok, /*TryExistingAST=*/true,
-               CancelOnSubsequentRequest, CancellationToken, Receiver);
+  resolveRange(*this, InputBufferName, Offset, Length, Invok,
+               /*TryExistingAST=*/true, CancelOnSubsequentRequest,
+               CancellationToken, Receiver);
 }
 
 void SwiftLangSupport::getNameInfo(
@@ -2282,17 +2337,18 @@ static void resolveCursorFromUSR(
 }
 
 void SwiftLangSupport::getCursorInfoFromUSR(
-    StringRef filename, StringRef USR, bool CancelOnSubsequentRequest,
-    ArrayRef<const char *> Args, Optional<VFSOptions> vfsOptions,
+    StringRef PrimaryFilePath, StringRef InputBufferName, StringRef USR,
+    bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
+    Optional<VFSOptions> vfsOptions,
     SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<CursorInfoData> &)> Receiver) {
   std::string error;
 
-  auto fileSystem = getFileSystem(vfsOptions, filename, error);
+  auto fileSystem = getFileSystem(vfsOptions, PrimaryFilePath, error);
   if (!fileSystem)
     return Receiver(RequestResult<CursorInfoData>::fromError(error));
 
-  if (auto IFaceGenRef = IFaceGenContexts.get(filename)) {
+  if (auto IFaceGenRef = IFaceGenContexts.get(PrimaryFilePath)) {
     LOG_WARN_FUNC("Info from usr for generated interface not implemented yet.");
     CursorInfoData Info;
     Info.InternalDiagnostic = "Info for generated interfaces not implemented.";
@@ -2302,16 +2358,16 @@ void SwiftLangSupport::getCursorInfoFromUSR(
 
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, filename, fileSystem, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, fileSystem, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<CursorInfoData>::fromError(Error));
     return;
   }
 
-  resolveCursorFromUSR(*this, filename, USR, Invok, /*TryExistingAST=*/true,
-                       CancelOnSubsequentRequest, fileSystem, CancellationToken,
-                       Receiver);
+  resolveCursorFromUSR(*this, InputBufferName, USR, Invok,
+                       /*TryExistingAST=*/true, CancelOnSubsequentRequest,
+                       fileSystem, CancellationToken, Receiver);
 }
 
 //===----------------------------------------------------------------------===//
@@ -2456,13 +2512,14 @@ private:
 } // end anonymous namespace
 
 void SwiftLangSupport::findRelatedIdentifiersInFile(
-    StringRef InputFile, unsigned Offset, bool CancelOnSubsequentRequest,
-    ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
+    StringRef PrimaryFilePath, StringRef InputBufferName, unsigned Offset,
+    bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
+    SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<RelatedIdentsInfo> &)> Receiver) {
 
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, InputFile, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<RelatedIdentsInfo>::fromError(Error));
@@ -2470,34 +2527,42 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
   }
 
   class RelatedIdConsumer : public SwiftASTConsumer {
+    std::string InputFile;
     unsigned Offset;
     std::function<void(const RequestResult<RelatedIdentsInfo> &)> Receiver;
     SwiftInvocationRef Invok;
 
   public:
-    RelatedIdConsumer(unsigned Offset,
-                      std::function<void(const RequestResult<RelatedIdentsInfo> &)> Receiver,
-                      SwiftInvocationRef Invok)
-      : Offset(Offset), Receiver(std::move(Receiver)), Invok(Invok) { }
+    RelatedIdConsumer(
+        StringRef InputFile, unsigned Offset,
+        std::function<void(const RequestResult<RelatedIdentsInfo> &)> Receiver,
+        SwiftInvocationRef Invok)
+        : InputFile(InputFile.str()), Offset(Offset),
+          Receiver(std::move(Receiver)), Invok(Invok) {}
 
     // FIXME: Don't silently eat errors here.
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       auto &CompInst = AstUnit->getCompilerInstance();
-      auto &SrcFile = AstUnit->getPrimarySourceFile();
 
+      auto *SrcFile = retrieveInputFile(InputFile, CompInst);
+      if (!SrcFile) {
+        Receiver(RequestResult<RelatedIdentsInfo>::fromError(
+            "Unable to find input file"));
+        return;
+      }
 
       SmallVector<std::pair<unsigned, unsigned>, 8> Ranges;
 
       auto Action = [&]() {
-        unsigned BufferID = SrcFile.getBufferID().value();
+        unsigned BufferID = SrcFile->getBufferID().value();
         SourceLoc Loc =
           Lexer::getLocForStartOfToken(CompInst.getSourceMgr(), BufferID, Offset);
         if (Loc.isInvalid())
           return;
 
         ResolvedCursorInfoPtr CursorInfo =
-            evaluateOrDefault(SrcFile.getASTContext().evaluator,
-                              CursorInfoRequest{CursorInfoOwner(&SrcFile, Loc)},
+            evaluateOrDefault(CompInst.getASTContext().evaluator,
+                              CursorInfoRequest{CursorInfoOwner(SrcFile, Loc)},
                               new ResolvedCursorInfo());
         auto ValueRefCursorInfo =
             dyn_cast<ResolvedValueRefCursorInfo>(CursorInfo);
@@ -2538,7 +2603,8 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
             continue;
           }
 
-          RelatedIdScanner Scanner(SrcFile, BufferID, Dcl, RangesSet, Worklist);
+          RelatedIdScanner Scanner(*SrcFile, BufferID, Dcl, RangesSet,
+                                   Worklist);
 
           if (auto *Case = getCaseStmtOfCanonicalVar(Dcl)) {
             Scanner.walk(Case);
@@ -2549,7 +2615,7 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
                          Dcl->getDeclContext()->getLocalContext()) {
             Scanner.walk(LocalDC);
           } else {
-            Scanner.walk(SrcFile);
+            Scanner.walk(*SrcFile);
           }
         }
 
@@ -2591,7 +2657,8 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
     }
   };
 
-  auto Consumer = std::make_shared<RelatedIdConsumer>(Offset, Receiver, Invok);
+  auto Consumer = std::make_shared<RelatedIdConsumer>(InputBufferName, Offset,
+                                                      Receiver, Invok);
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
@@ -2606,16 +2673,14 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
 
 namespace {
 class IfConfigScanner : public SourceEntityWalker {
-  SmallVectorImpl<IfConfigInfo> &Infos;
-  SourceManager &SourceMgr;
   unsigned BufferID = -1;
+  SmallVectorImpl<IfConfigInfo> &Infos;
   bool Cancelled = false;
 
 public:
-  explicit IfConfigScanner(SourceFile &SrcFile, unsigned BufferID,
+  explicit IfConfigScanner(unsigned BufferID,
                            SmallVectorImpl<IfConfigInfo> &Infos)
-      : Infos(Infos), SourceMgr(SrcFile.getASTContext().SourceMgr),
-        BufferID(BufferID) {}
+      : BufferID(BufferID), Infos(Infos) {}
 
 private:
   bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
@@ -2624,7 +2689,8 @@ private:
 
     if (auto *IfDecl = dyn_cast<IfConfigDecl>(D)) {
       for (auto &Clause : IfDecl->getClauses()) {
-        unsigned Offset = SourceMgr.getLocOffsetInBuffer(Clause.Loc, BufferID);
+        unsigned Offset = D->getASTContext().SourceMgr.getLocOffsetInBuffer(
+            Clause.Loc, BufferID);
         Infos.emplace_back(Offset, Clause.isActive);
       }
     }
@@ -2636,13 +2702,13 @@ private:
 } // end anonymous namespace
 
 void SwiftLangSupport::findActiveRegionsInFile(
-    StringRef InputFile, ArrayRef<const char *> Args,
-    SourceKitCancellationToken CancellationToken,
+    StringRef PrimaryFilePath, StringRef InputBufferName,
+    ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<ActiveRegionsInfo> &)> Receiver) {
 
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, InputFile, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<ActiveRegionsInfo>::fromError(Error));
@@ -2650,23 +2716,29 @@ void SwiftLangSupport::findActiveRegionsInFile(
   }
 
   class IfConfigConsumer : public SwiftASTConsumer {
+    std::string InputFile;
     std::function<void(const RequestResult<ActiveRegionsInfo> &)> Receiver;
     SwiftInvocationRef Invok;
 
   public:
     IfConfigConsumer(
+        StringRef InputFile,
         std::function<void(const RequestResult<ActiveRegionsInfo> &)> Receiver,
         SwiftInvocationRef Invok)
-        : Receiver(std::move(Receiver)), Invok(Invok) {}
+        : InputFile(InputFile.str()), Receiver(std::move(Receiver)),
+          Invok(Invok) {}
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
-      auto &SrcFile = AstUnit->getPrimarySourceFile();
-      SmallVector<IfConfigInfo> Configs;
-      auto BufferID = SrcFile.getBufferID();
-      if (!BufferID)
+      auto *SF = retrieveInputFile(InputFile, AstUnit->getCompilerInstance());
+      if (!SF) {
+        Receiver(RequestResult<ActiveRegionsInfo>::fromError(
+            "Unable to find input file"));
         return;
-      IfConfigScanner Scanner(SrcFile, *BufferID, Configs);
-      Scanner.walk(SrcFile);
+      }
+
+      SmallVector<IfConfigInfo> Configs;
+      IfConfigScanner Scanner(*SF->getBufferID(), Configs);
+      Scanner.walk(SF);
 
       // Sort by offset so nested decls are reported
       // in source order (not tree order).
@@ -2689,7 +2761,8 @@ void SwiftLangSupport::findActiveRegionsInFile(
     }
   };
 
-  auto Consumer = std::make_shared<IfConfigConsumer>(Receiver, Invok);
+  auto Consumer =
+      std::make_shared<IfConfigConsumer>(InputBufferName, Receiver, Invok);
   ASTMgr->processASTAsync(Invok, std::move(Consumer),
                           /*OncePerASTToken=*/nullptr, CancellationToken,
                           llvm::vfs::getRealFileSystem());
@@ -2709,12 +2782,12 @@ static RefactoringKind getIDERefactoringKind(SemanticRefactoringInfo Info) {
 }
 
 void SwiftLangSupport::semanticRefactoring(
-    StringRef PrimaryFile, SemanticRefactoringInfo Info,
+    StringRef PrimaryFilePath, SemanticRefactoringInfo Info,
     ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
     CategorizedEditsReceiver Receiver) {
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, PrimaryFile, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<ArrayRef<CategorizedEdits>>::fromError(Error));
@@ -2735,20 +2808,17 @@ void SwiftLangSupport::semanticRefactoring(
       RequestRefactoringEditConsumer EditConsumer(Receiver);
 
       auto &CompIns = AstUnit->getCompilerInstance();
-      auto &SM = CompIns.getSourceMgr();
 
       RefactoringOptions Opts(getIDERefactoringKind(Info));
 
-      Optional<unsigned> BufferID;
-      if (Info.SourceFile.empty()) {
-        BufferID = AstUnit->getPrimarySourceFile().getBufferID();
-      } else {
-        BufferID = SM.getIDForBufferIdentifier(Info.SourceFile);
-      }
-      if (!BufferID)
+      SourceFile *SF = retrieveInputFile(Info.InputBufferName, CompIns);
+      if (!SF) {
+        Receiver(RequestResult<ArrayRef<CategorizedEdits>>::fromError(
+            "Unable to find input file"));
         return;
+      }
 
-      Opts.Range.BufferID = *BufferID;
+      Opts.Range.BufferID = *SF->getBufferID();
       Opts.Range.Line = Info.Line;
       Opts.Range.Column = Info.Column;
       Opts.Range.Length = Info.Length;
@@ -2777,14 +2847,15 @@ void SwiftLangSupport::semanticRefactoring(
 }
 
 void SwiftLangSupport::collectExpressionTypes(
-    StringRef FileName, ArrayRef<const char *> Args,
-    ArrayRef<const char *> ExpectedProtocols, bool FullyQualified,
-    bool CanonicalType, SourceKitCancellationToken CancellationToken,
+    StringRef PrimaryFilePath, StringRef InputBufferName,
+    ArrayRef<const char *> Args, ArrayRef<const char *> ExpectedProtocols,
+    bool FullyQualified, bool CanonicalType,
+    SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<ExpressionTypesInFile> &)>
         Receiver) {
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, FileName, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<ExpressionTypesInFile>::fromError(Error));
@@ -2793,6 +2864,7 @@ void SwiftLangSupport::collectExpressionTypes(
   assert(Invok);
   class ExpressionTypeCollector: public SwiftASTConsumer {
     std::function<void(const RequestResult<ExpressionTypesInFile> &)> Receiver;
+    std::string InputFile;
     std::vector<const char *> ExpectedProtocols;
     bool FullyQualified;
     bool CanonicalType;
@@ -2800,13 +2872,20 @@ void SwiftLangSupport::collectExpressionTypes(
     ExpressionTypeCollector(
         std::function<void(const RequestResult<ExpressionTypesInFile> &)>
             Receiver,
-        ArrayRef<const char *> ExpectedProtocols, bool FullyQualified,
-        bool CanonicalType)
-        : Receiver(std::move(Receiver)),
+        StringRef InputFile, ArrayRef<const char *> ExpectedProtocols,
+        bool FullyQualified, bool CanonicalType)
+        : Receiver(std::move(Receiver)), InputFile(InputFile.str()),
           ExpectedProtocols(ExpectedProtocols.vec()),
           FullyQualified(FullyQualified), CanonicalType(CanonicalType) {}
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
-      auto *SF = AstUnit->getCompilerInstance().getPrimarySourceFile();
+      SourceFile *SF =
+          retrieveInputFile(InputFile, AstUnit->getCompilerInstance());
+      if (!SF) {
+        Receiver(RequestResult<ExpressionTypesInFile>::fromError(
+            "Unable to find input file"));
+        return;
+      }
+
       std::vector<ExpressionTypeInfo> Scratch;
       llvm::SmallString<256> TypeBuffer;
       llvm::raw_svector_ostream OS(TypeBuffer);
@@ -2832,7 +2911,8 @@ void SwiftLangSupport::collectExpressionTypes(
     }
   };
   auto Collector = std::make_shared<ExpressionTypeCollector>(
-      Receiver, ExpectedProtocols, FullyQualified, CanonicalType);
+      Receiver, InputBufferName, ExpectedProtocols, FullyQualified,
+      CanonicalType);
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
@@ -2842,13 +2922,14 @@ void SwiftLangSupport::collectExpressionTypes(
 }
 
 void SwiftLangSupport::collectVariableTypes(
-    StringRef FileName, ArrayRef<const char *> Args, Optional<unsigned> Offset,
+    StringRef PrimaryFilePath, StringRef InputBufferName,
+    ArrayRef<const char *> Args, Optional<unsigned> Offset,
     Optional<unsigned> Length, bool FullyQualified,
     SourceKitCancellationToken CancellationToken,
     std::function<void(const RequestResult<VariableTypesInFile> &)> Receiver) {
   std::string Error;
   SwiftInvocationRef Invok =
-      ASTMgr->getTypecheckInvocation(Args, FileName, Error);
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
   if (!Invok) {
     LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
     Receiver(RequestResult<VariableTypesInFile>::fromError(Error));
@@ -2859,6 +2940,7 @@ void SwiftLangSupport::collectVariableTypes(
   class VariableTypeCollectorASTConsumer : public SwiftASTConsumer {
   private:
     std::function<void(const RequestResult<VariableTypesInFile> &)> Receiver;
+    std::string InputFile;
     Optional<unsigned> Offset;
     Optional<unsigned> Length;
     bool FullyQualified;
@@ -2867,14 +2949,19 @@ void SwiftLangSupport::collectVariableTypes(
     VariableTypeCollectorASTConsumer(
         std::function<void(const RequestResult<VariableTypesInFile> &)>
             Receiver,
-        Optional<unsigned> Offset, Optional<unsigned> Length,
-        bool FullyQualified)
-        : Receiver(std::move(Receiver)), Offset(Offset), Length(Length),
-          FullyQualified(FullyQualified) {}
+        StringRef InputFile, Optional<unsigned> Offset,
+        Optional<unsigned> Length, bool FullyQualified)
+        : Receiver(std::move(Receiver)), InputFile(InputFile), Offset(Offset),
+          Length(Length), FullyQualified(FullyQualified) {}
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       auto &CompInst = AstUnit->getCompilerInstance();
-      auto *SF = CompInst.getPrimarySourceFile();
+      SourceFile *SF = retrieveInputFile(InputFile, CompInst);
+      if (!SF) {
+        Receiver(RequestResult<VariableTypesInFile>::fromError(
+            "Unable to find input file"));
+        return;
+      }
 
       // Construct the range for which variable types are to be queried. If
       // offset/length are unset, the (default) range will be used, which
@@ -2913,7 +3000,7 @@ void SwiftLangSupport::collectVariableTypes(
   };
 
   auto Collector = std::make_shared<VariableTypeCollectorASTConsumer>(
-      Receiver, Offset, Length, FullyQualified);
+      Receiver, InputBufferName, Offset, Length, FullyQualified);
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;