[SourceKit] Make sure we reuse ASTContext in function bodies for solver-based cursor info

The main problem that prevented us from reusing the ASTContext was that we weren’t remapping the `LocToResolve` in the temporary buffer that only contains the re-parsed function back to the original buffer. Thus `NodeFinder` couldn’t find the node that we want to get cursor info for.

Getting AST reuse to work for top-level items is harder because it currently heavily relies on the `HasCodeCompletion` state being set on the parser result. I’ll try that in a follow-up PR.

rdar://103251263

Author: ahoppen

include/swift/Basic/SourceManager.h

@@ -166,6 +166,17 @@ class SourceManager {
 
   SourceLoc getIDEInspectionTargetLoc() const;
 
+  /// Returns whether `Range` contains `Loc`. This also respects the
+  /// `ReplacedRanges`, i.e. if `Loc` is in a range that replaces a range which
+  /// overlaps `Range`, this also returns `true`.
+  bool containsRespectingReplacedRanges(SourceRange Range, SourceLoc Loc) const;
+
+  /// Returns whether `Enclosing` contains `Inner`. This also respects the
+  /// `ReplacedRanges`, i.e. if `Inner` is contained in a range that replaces a
+  /// range which overlaps `Range`, this also returns `true`.
+  bool rangeContainsRespectingReplacedRanges(SourceRange Enclosing,
+                                             SourceRange Inner) const;
+
   const llvm::DenseMap<SourceRange, SourceRange> &getReplacedRanges() const {
     return ReplacedRanges;
   }

lib/Basic/SourceLoc.cpp

@@ -45,6 +45,31 @@ SourceLoc SourceManager::getIDEInspectionTargetLoc() const {
       .getAdvancedLoc(IDEInspectionTargetOffset);
 }
 
+bool SourceManager::containsRespectingReplacedRanges(SourceRange Range,
+                                                     SourceLoc Loc) const {
+  if (Loc.isInvalid() || Range.isInvalid()) {
+    return false;
+  }
+
+  if (Range.contains(Loc)) {
+    return true;
+  }
+  for (const auto &pair : getReplacedRanges()) {
+    auto OriginalRange = pair.first;
+    auto NewRange = pair.second;
+    if (NewRange.contains(Loc) && Range.overlaps(OriginalRange)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool SourceManager::rangeContainsRespectingReplacedRanges(
+    SourceRange Enclosing, SourceRange Inner) const {
+  return containsRespectingReplacedRanges(Enclosing, Inner.Start) &&
+         containsRespectingReplacedRanges(Enclosing, Inner.End);
+}
+
 StringRef SourceManager::getDisplayNameForLoc(SourceLoc Loc, bool ForceGeneratedSourceToDisk) const {
   // Respect #line first
   if (auto VFile = getVirtualFile(Loc))

lib/IDE/CursorInfo.cpp

@@ -140,7 +140,7 @@ class NodeFinder : ASTWalker {
   DeclContext *getCurrentDeclContext() { return DeclContextStack.back(); }
 
   bool rangeContainsLocToResolve(SourceRange Range) const {
-    return Range.contains(LocToResolve);
+    return getSourceMgr().containsRespectingReplacedRanges(Range, LocToResolve);
   }
 
   PreWalkAction walkToDeclPre(Decl *D) override {

lib/IDETool/IDEInspectionInstance.cpp

@@ -842,10 +842,9 @@ void swift::ide::IDEInspectionInstance::cursorInfo(
       CancellationFlag,
       [&](CancellableResult<IDEInspectionInstanceResult> CIResult) {
         CIResult.mapAsync<CursorInfoResults>(
-            [&CancellationFlag, Offset](auto &Result, auto DeliverTransformed) {
+            [&CancellationFlag](auto &Result, auto DeliverTransformed) {
               auto &Mgr = Result.CI->getSourceMgr();
-              auto RequestedLoc = Mgr.getLocForOffset(
-                  Mgr.getIDEInspectionTargetBufferID(), Offset);
+              auto RequestedLoc = Mgr.getIDEInspectionTargetLoc();
               ConsumerToCallbackAdapter Consumer(
                   Result.DidReuseAST, CancellationFlag, DeliverTransformed);
               std::unique_ptr<IDEInspectionCallbacksFactory> callbacksFactory(

lib/Refactoring/Refactoring.cpp

@@ -52,9 +52,11 @@ class ContextFinder : public SourceEntityWalker {
   std::function<bool(ASTNode)> IsContext;
   SmallVector<ASTNode, 4> AllContexts;
   bool contains(ASTNode Enclosing) {
-    auto Result = SM.rangeContains(Enclosing.getSourceRange(), Target);
-    if (Result && IsContext(Enclosing))
+    auto Result = SM.rangeContainsRespectingReplacedRanges(
+        Enclosing.getSourceRange(), Target);
+    if (Result && IsContext(Enclosing)) {
       AllContexts.push_back(Enclosing);
+    }
     return Result;
   }
 public:

test/SourceKit/CursorInfo/cursor_reuses_astcontext.swift

@@ -0,0 +1,27 @@
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 2):7 %s -- %s == -req=cursor -pos=%(line + 3):7 %s -- %s | %FileCheck %s --check-prefix IN-FUNCTION
+func foo() {
+  let inFunctionA = 1
+  let inFunctionB = "hi"
+}
+
+// IN-FUNCTION: source.lang.swift.decl.var.local
+// IN-FUNCTION-NEXT: inFunctionA
+// IN-FUNCTION: DID REUSE AST CONTEXT: 0
+// IN-FUNCTION: source.lang.swift.decl.var.local
+// IN-FUNCTION-NEXT: inFunctionB
+// IN-FUNCTION: DID REUSE AST CONTEXT: 1
+
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 3):9 %s -- %s == -req=cursor -pos=%(line + 4):9 %s -- %s | %FileCheck %s --check-prefix IN-INSTANCE-METHOD
+struct MyStruct {
+  func test() {
+    let inInstanceMethod1 = 2
+    let inInstanceMethod2 = "hello"
+  }
+}
+
+// IN-INSTANCE-METHOD: source.lang.swift.decl.var.local
+// IN-INSTANCE-METHOD-NEXT: inInstanceMethod1
+// IN-INSTANCE-METHOD: DID REUSE AST CONTEXT: 0
+// IN-INSTANCE-METHOD: source.lang.swift.decl.var.local
+// IN-INSTANCE-METHOD-NEXT: inInstanceMethod2
+// IN-INSTANCE-METHOD: DID REUSE AST CONTEXT: 1

test/SourceKit/CursorInfo/refactoring_actions_with_ast_reuse.swift

@@ -0,0 +1,20 @@
+func materialize() {
+ let a = 2
+}
+
+func outer() {
+  func inner() {
+  }
+}
+
+// RUN: %sourcekitd-test -req=cursor -pos=2:6 %s -- %s == \
+// RUN: -req=cursor -req-opts=retrieve_refactor_actions=1 -pos=6:8 %s -- %s | %FileCheck %s
+
+// CHECK: source.lang.swift.decl.function.free (6:8-6:15)
+// CHECK: inner()
+// CHECK: ACTIONS BEGIN
+// CHECK-DAG: source.refactoring.kind.rename.local
+// CHECK-DAG: Local Rename
+// CHECK-DAG: source.refactoring.kind.convert.func-to-async
+// CHECK-DAG: Convert Function to Async
+// CHECK: ACTIONS END

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -363,6 +363,7 @@ class RequestResult {
     return RequestResult();
   }
 
+  bool isValue() const { return type == Value; }
   const T &value() const {
     assert(type == Value);
     return *data;
@@ -588,8 +589,14 @@ struct CursorInfoData {
   llvm::ArrayRef<CursorSymbolInfo> Symbols;
   /// All available actions on the code under cursor.
   llvm::ArrayRef<RefactoringInfo> AvailableActions;
-
-  void print(llvm::raw_ostream &OS, std::string Indentation) const {
+  /// Whether the ASTContext was reused for this cursor info.
+  bool DidReuseAST = false;
+
+  /// If \p ForSolverBasedCursorInfoVerification is \c true, fields that are
+  /// acceptable to differ between the AST-based and the solver-based result,
+  /// will be excluded.
+  void print(llvm::raw_ostream &OS, std::string Indentation,
+             bool ForSolverBasedCursorInfoVerification = false) const {
     OS << Indentation << "CursorInfoData" << '\n';
     OS << Indentation << "  Symbols:" << '\n';
     for (auto Symbol : Symbols) {
@@ -599,6 +606,9 @@ struct CursorInfoData {
     for (auto AvailableAction : AvailableActions) {
       AvailableAction.print(OS, Indentation + "    ");
     }
+    if (!ForSolverBasedCursorInfoVerification) {
+      OS << Indentation << "DidReuseAST: " << DidReuseAST << '\n';
+    }
   }
 
   SWIFT_DEBUG_DUMP { print(llvm::errs(), ""); }

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -911,6 +911,24 @@ static void setLocationInfo(const ValueDecl *VD,
     Location.Length = NameLen;
     std::tie(Location.Line, Location.Column) = SM.getLineAndColumnInBuffer(
         Loc, DeclBufID);
+    if (auto GeneratedSourceInfo = SM.getGeneratedSourceInfo(DeclBufID)) {
+      if (GeneratedSourceInfo->kind ==
+          GeneratedSourceInfo::ReplacedFunctionBody) {
+        // The location was in a temporary source buffer that just contains the
+        // function body and which we created while reusing the ASTContext for
+        // the rest of the file. Map the location back to the original file.
+        unsigned OriginalBufID = SM.findBufferContainingLoc(
+            GeneratedSourceInfo->originalSourceRange.Start);
+        auto OriginalStartOffset = SM.getLocOffsetInBuffer(
+            GeneratedSourceInfo->originalSourceRange.Start, OriginalBufID);
+        auto GeneratedStartOffset = SM.getLocOffsetInBuffer(
+            GeneratedSourceInfo->generatedSourceRange.Start, DeclBufID);
+        Location.Offset += OriginalStartOffset - GeneratedStartOffset;
+        assert(SM.findBufferContainingLoc(Loc) == DeclBufID);
+        std::tie(Location.Line, Location.Column) =
+            SM.getPresumedLineAndColumnForLoc(Loc, DeclBufID);
+      }
+    }
   } else if (ClangNode) {
     ClangImporter *Importer =
         static_cast<ClangImporter*>(Ctx.getClangModuleLoader());
@@ -1146,6 +1164,7 @@ static bool passCursorInfoForDecl(
     bool AddSymbolGraph, ArrayRef<RefactoringInfo> KnownRefactoringInfo,
     SwiftLangSupport &Lang, const CompilerInvocation &Invoc,
     std::string &Diagnostic, ArrayRef<ImmutableTextSnapshotRef> PreviousSnaps,
+    bool DidReuseAST,
     std::function<void(const RequestResult<CursorInfoData> &)> Receiver) {
   DeclInfo OrigInfo(Info.getValueD(), Info.getContainerType(), Info.isRef(),
                     Info.isDynamic(), Info.getReceiverTypes(), Invoc);
@@ -1215,6 +1234,7 @@ static bool passCursorInfoForDecl(
   CursorInfoData Data;
   Data.Symbols = llvm::makeArrayRef(Symbols);
   Data.AvailableActions = llvm::makeArrayRef(Refactorings);
+  Data.DidReuseAST = DidReuseAST;
   Receiver(RequestResult<CursorInfoData>::fromResult(Data));
   return true;
 }
@@ -1554,7 +1574,8 @@ static void resolveCursor(
         bool Success = passCursorInfoForDecl(
             cast<ResolvedValueRefCursorInfo>(CursorInfo), Actionables,
             SymbolGraph, Actions, Lang, CompInvok, Diagnostic,
-            getPreviousASTSnaps(), Receiver);
+            getPreviousASTSnaps(),
+            /*DidReuseAST=*/false, Receiver);
         if (!Success) {
           if (!getPreviousASTSnaps().empty()) {
             // Attempt again using the up-to-date AST.
@@ -1873,7 +1894,8 @@ static void deliverCursorInfoResults(
             Results.getResult().Result)) {
       std::string Diagnostic; // Unused
       passCursorInfoForDecl(*Result, AddRefactorings, AddSymbolGraph, {}, Lang,
-                            Invoc, Diagnostic, /*PreviousSnaps=*/{}, Receiver);
+                            Invoc, Diagnostic, /*PreviousSnaps=*/{},
+                            Results->DidReuseAST, Receiver);
     }
     break;
   }
@@ -1914,8 +1936,10 @@ void SwiftLangSupport::getCursorInfo(
           ResolvedValueRefCursorInfo Info;
           Info.setValueD(const_cast<ValueDecl *>(Entity.Dcl));
           Info.setIsRef(Entity.IsRef);
-          passCursorInfoForDecl(Info, Actionables, SymbolGraph, {}, *this,
-                                Invok, Diagnostic, {}, Receiver);
+          passCursorInfoForDecl(
+              Info, Actionables, SymbolGraph, {}, *this, Invok, Diagnostic,
+              /*PreviousSnaps=*/{},
+              /*DidReuseAST=*/false, Receiver);
         }
       } else {
         CursorInfoData Info;
@@ -1940,12 +1964,10 @@ void SwiftLangSupport::getCursorInfo(
 
   /// Counts how many symbols \p Res contains.
   auto ResultCount = [](const RequestResult<CursorInfoData> &Res) -> size_t {
-    if (Res.isCancelled()) {
-      return 0;
-    } else if (Res.isError()) {
-      return 0;
-    } else {
+    if (Res.isValue()) {
       return Res.value().Symbols.size();
+    } else {
+      return 0;
     }
   };
 
@@ -1959,7 +1981,8 @@ void SwiftLangSupport::getCursorInfo(
     } else {
       std::string Description;
       llvm::raw_string_ostream OS(Description);
-      Res.value().print(OS, /*Indentation=*/"");
+      Res.value().print(OS, /*Indentation=*/"",
+                        /*ForSolverBasedCursorInfoVerification=*/true);
       return OS.str();
     }
   };
@@ -1979,6 +2002,7 @@ void SwiftLangSupport::getCursorInfo(
   // AST-based result.
   std::string SolverBasedResultDescription;
   size_t SolverBasedResultCount = 0;
+  bool SolverBasedReusedAST = false;
   if (EnableSolverBasedCursorInfo) {
     std::string InputFileError;
     llvm::SmallString<64> RealInputFilePath;
@@ -1990,6 +2014,9 @@ void SwiftLangSupport::getCursorInfo(
       auto SolverBasedReceiver = [&](const RequestResult<CursorInfoData> &Res) {
         SolverBasedResultCount = ResultCount(Res);
         SolverBasedResultDescription = ResultDescription(Res);
+        if (Res.isValue()) {
+          SolverBasedReusedAST = Res.value().DidReuseAST;
+        }
       };
 
       CompilerInvocation CompInvok;
@@ -2053,15 +2080,24 @@ void SwiftLangSupport::getCursorInfo(
   // Thunk around `Receiver` that, if solver-based cursor info is enabled,
   // verifies that the solver-based cursor info result matches the AST-based
   // result.
-  auto ReceiverThunk = [Receiver, VerifySolverBasedResult](
-                           const RequestResult<CursorInfoData> &Res) {
-    auto VerificationError = VerifySolverBasedResult(Res);
-    if (VerificationError.empty()) {
-      Receiver(Res);
-    } else {
-      Receiver(RequestResult<CursorInfoData>::fromError(VerificationError));
-    }
-  };
+  auto ReceiverThunk =
+      [Receiver, VerifySolverBasedResult,
+       SolverBasedReusedAST](const RequestResult<CursorInfoData> &Res) {
+        auto VerificationError = VerifySolverBasedResult(Res);
+        if (VerificationError.empty()) {
+          if (Res.isValue()) {
+            // Report whether the solver-based implemenatation reused the AST so
+            // we can check it in test cases.
+            auto Value = Res.value();
+            Value.DidReuseAST = SolverBasedReusedAST;
+            Receiver(RequestResult<CursorInfoData>::fromResult(Value));
+          } else {
+            Receiver(Res);
+          }
+        } else {
+          Receiver(RequestResult<CursorInfoData>::fromError(VerificationError));
+        }
+      };
 
   resolveCursor(*this, InputFile, Offset, Length, Actionables, SymbolGraph,
                 Invok, /*TryExistingAST=*/true, CancelOnSubsequentRequest,
@@ -2253,10 +2289,11 @@ static void resolveCursorFromUSR(
         }
 
         std::string Diagnostic;
-        bool Success =
-            passCursorInfoForDecl(Info, /*AddRefactorings*/ false,
-                                  /*AddSymbolGraph*/ false, {}, Lang, CompInvok,
-                                  Diagnostic, PreviousASTSnaps, Receiver);
+        bool Success = passCursorInfoForDecl(
+            Info, /*AddRefactorings*/ false,
+            /*AddSymbolGraph*/ false, {}, Lang, CompInvok, Diagnostic,
+            PreviousASTSnaps,
+            /*DidReuseAST=*/false, Receiver);
         if (!Success) {
           if (!PreviousASTSnaps.empty()) {
             // Attempt again using the up-to-date AST.

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -1991,6 +1991,9 @@ static void printCursorInfo(sourcekitd_variant_t Info, StringRef FilenameIn,
     OS << "-----\n";
   }
   OS << "SECONDARY SYMBOLS END\n";
+  OS << "DID REUSE AST CONTEXT: "
+     << sourcekitd_variant_dictionary_get_bool(Info, KeyReusingASTContext)
+     << '\n';
 }
 
 static void printRangeInfo(sourcekitd_variant_t Info, StringRef FilenameIn,

tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp

@@ -2406,6 +2406,8 @@ static void reportCursorInfo(const RequestResult<CursorInfoData> &Result,
     }
   }
 
+  Elem.setBool(KeyReusingASTContext, Info.DidReuseAST);
+
   return Rec(RespBuilder.createResponse());
 }