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

ahoppen · ahoppen · commit be0b18d79cbf · 2023-01-30T23:48:39.000+01:00
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
diff --git a/include/swift/Basic/SourceManager.h b/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;
   }
diff --git a/lib/Basic/SourceLoc.cpp b/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))
diff --git a/lib/IDE/CursorInfo.cpp b/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 {
diff --git a/lib/IDETool/IDEInspectionInstance.cpp b/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(
diff --git a/lib/Refactoring/Refactoring.cpp b/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:
diff --git a/test/SourceKit/CursorInfo/cursor_reuses_astcontext.swift b/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
diff --git a/test/SourceKit/CursorInfo/cursor_with_file_replacement.swift b/test/SourceKit/CursorInfo/cursor_with_file_replacement.swift
@@ -0,0 +1,91 @@
+// BEGIN State1.swift
+
+func foo() {
+  let inFunctionA = 1
+  let inFunctionB = "hi"
+}
+
+// BEGIN State2.swift
+
+func foo() {
+  let newlyAddedMember = 3
+  let inFunctionA = 1
+  let inFunctionB = "hi"
+}
+
+// BEGIN State3.swift
+
+func foo() {
+  let inFunctionB = "hi"
+}
+
+// BEGIN State4.swift
+
+func foo() {
+  let myNewName = "hi"
+}
+
+// BEGIN State5.swift
+
+func foo(param: Int) {
+  let myNewName = "hi"
+}
+
+// BEGIN State6.swift
+
+func foo(param: Int) {
+  let myNewName = 7
+}
+
+
+// BEGIN Dummy.swift
+
+// RUN: %empty-directory(%t)
+// RUN: %{python} %utils/split_file.py -o %t %s
+
+// RUN: cp %t/State1.swift %t/file.swift
+
+// RUN: %sourcekitd-test \
+// RUN:   -shell -- echo '## State 1' == \
+// RUN:   -req=cursor -pos=4:7 %t/file.swift -- %t/file.swift == \
+// RUN:   -shell -- echo '## State 2' == \
+// RUN:   -shell -- cp %t/State2.swift %t/file.swift == \
+// RUN:   -req=cursor -pos=5:7 %t/file.swift -- %t/file.swift == \
+// RUN:   -shell -- echo '## State 3' == \
+// RUN:   -shell -- cp %t/State3.swift %t/file.swift == \
+// RUN:   -req=cursor -pos=3:7 %t/file.swift -- %t/file.swift == \
+// RUN:   -shell -- echo '## State 4' == \
+// RUN:   -shell -- cp %t/State4.swift %t/file.swift == \
+// RUN:   -req=cursor -pos=3:7 %t/file.swift -- %t/file.swift == \
+// RUN:   -shell -- echo '## State 5' == \
+// RUN:   -shell -- cp %t/State5.swift %t/file.swift == \
+// RUN:   -req=cursor -pos=3:7 %t/file.swift -- %t/file.swift == \
+// RUN:   -shell -- echo '## State 6' == \
+// RUN:   -shell -- cp %t/State6.swift %t/file.swift == \
+// RUN:   -req=cursor -pos=3:7 %t/file.swift -- %t/file.swift > %t/response.txt
+// RUN: %FileCheck %s < %t/response.txt
+
+// CHECK-LABEL: ## State 1
+// CHECK: source.lang.swift.decl.var.local (4:7-4:18)
+// CHECK: <Declaration>let inFunctionB: <Type usr="s:SS">String</Type></Declaration>
+// CHECK: DID REUSE AST CONTEXT: 0
+// CHECK-LABEL: ## State 2
+// CHECK: source.lang.swift.decl.var.local (5:7-5:18)
+// CHECK: <Declaration>let inFunctionB: <Type usr="s:SS">String</Type></Declaration>
+// CHECK: DID REUSE AST CONTEXT: 1
+// CHECK-LABEL: ## State 3
+// CHECK: source.lang.swift.decl.var.local (3:7-3:18)
+// CHECK: <Declaration>let inFunctionB: <Type usr="s:SS">String</Type></Declaration>
+// CHECK: DID REUSE AST CONTEXT: 1
+// CHECK-LABEL: ## State 4
+// CHECK: source.lang.swift.decl.var.local (3:7-3:16)
+// CHECK: <Declaration>let myNewName: <Type usr="s:SS">String</Type></Declaration>
+// CHECK: DID REUSE AST CONTEXT: 1
+// CHECK-LABEL: ## State 5
+// CHECK: source.lang.swift.decl.var.local (3:7-3:16)
+// CHECK: <Declaration>let myNewName: <Type usr="s:SS">String</Type></Declaration>
+// CHECK: DID REUSE AST CONTEXT: 0
+// CHECK-LABEL: ## State 6
+// CHECK: source.lang.swift.decl.var.local (3:7-3:16)
+// CHECK: <Declaration>let myNewName: <Type usr="s:Si">Int</Type></Declaration>
+// CHECK: DID REUSE AST CONTEXT: 1
diff --git a/test/SourceKit/CursorInfo/refactoring_actions_with_ast_reuse.swift b/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
diff --git a/tools/SourceKit/include/SourceKit/Core/LangSupport.h b/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(), ""); }
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp b/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
diff --git a/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp b/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp
diff --git a/tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp b/tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp

Original file line number	Diff line number	Diff line change
`@@ -140,7 +140,7 @@ class NodeFinder : ASTWalker {`
`140`	`140`	`DeclContext *getCurrentDeclContext() { return DeclContextStack.back(); }`
`141`	`141`
`142`	`142`	`bool rangeContainsLocToResolve(SourceRange Range) const {`
`143`		`- return Range.contains(LocToResolve);`
	`143`	`+ return getSourceMgr().containsRespectingReplacedRanges(Range, LocToResolve);`
`144`	`144`	`}`
`145`	`145`
`146`	`146`	`PreWalkAction walkToDeclPre(Decl *D) override {`