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

include/swift/AST/TypeCheckRequests.h

@@ -3102,7 +3102,8 @@ class LookupAllConformancesInContextRequest
 
 class CheckRedeclarationRequest
     : public SimpleRequest<
-          CheckRedeclarationRequest, evaluator::SideEffect(ValueDecl *),
+          CheckRedeclarationRequest,
+          evaluator::SideEffect(ValueDecl *, NominalTypeDecl *),
           RequestFlags::SeparatelyCached | RequestFlags::DependencySource |
               RequestFlags::DependencySink> {
 public:
@@ -3112,10 +3113,13 @@ class CheckRedeclarationRequest
   friend SimpleRequest;
 
    // Evaluation.
-  evaluator::SideEffect
-  evaluate(Evaluator &evaluator, ValueDecl *VD) const;
+  /// \p SelfNominalType is \c VD->getDeclContext()->getSelfNominalType().
+  /// Passed as a parameter in here so this request doesn't tigger self nominal
+  /// type computation.
+  evaluator::SideEffect evaluate(Evaluator &evaluator, ValueDecl *VD,
+                                 NominalTypeDecl *SelfNominalType) const;
 
- public:
+public:
   // Separate caching.
   bool isCached() const { return true; }
   Optional<evaluator::SideEffect> getCachedResult() const;

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/AST/TypeCheckRequests.cpp

@@ -1324,7 +1324,7 @@ void CheckRedeclarationRequest::writeDependencySink(
     return;
 
   if (currentDC->isTypeContext()) {
-    if (auto nominal = currentDC->getSelfNominalTypeDecl()) {
+    if (auto nominal = std::get<1>(getStorage())) {
       tracker.addUsedMember(nominal, current->getBaseName());
     }
   } else {

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

@@ -50,6 +50,13 @@ void typeCheckDeclAndParentClosures(ValueDecl *VD) {
   typeCheckASTNodeAtLoc(
       TypeCheckASTNodeAtLocContext::declContext(VD->getDeclContext()),
       VD->getLoc());
+  if (auto VarD = dyn_cast<VarDecl>(VD)) {
+    // Type check any attached property wrappers so the annotated declaration
+    // can refer to their USRs.
+    (void)VarD->getPropertyWrapperBackingPropertyType();
+    // Visit emitted accessors so we generated accessors from property wrappers.
+    VarD->visitEmittedAccessors([&](AccessorDecl *accessor) {});
+  }
 }
 
 // MARK: - NodeFinderResults
@@ -133,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 {
@@ -150,7 +157,10 @@ class NodeFinder : ASTWalker {
     }
 
     if (auto VD = dyn_cast<ValueDecl>(D)) {
-      if (VD->hasName()) {
+      // FIXME: ParamDecls might be closure parameters that can have ambiguous
+      // types. The current infrastructure of just asking for the VD's type
+      // doesn't work here. We need to inspect the constraints system solution.
+      if (VD->hasName() && !isa<ParamDecl>(D)) {
         assert(Result == nullptr);
         Result = std::make_unique<NodeFinderDeclResult>(VD);
         return Action::Stop();

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:

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -542,7 +542,8 @@ static void checkRedeclaration(PrecedenceGroupDecl *group) {
 
 /// Check whether \c current is a redeclaration.
 evaluator::SideEffect
-CheckRedeclarationRequest::evaluate(Evaluator &eval, ValueDecl *current) const {
+CheckRedeclarationRequest::evaluate(Evaluator &eval, ValueDecl *current,
+                                    NominalTypeDecl *SelfNominalType) const {
   // Ignore invalid and anonymous declarations.
   if (current->isInvalid() || !current->hasName())
     return std::make_tuple<>();
@@ -1872,8 +1873,11 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
       // Force some requests, which can produce diagnostics.
 
       // Check redeclaration.
-      (void) evaluateOrDefault(Context.evaluator,
-                               CheckRedeclarationRequest{VD}, {});
+      (void)evaluateOrDefault(
+          Context.evaluator,
+          CheckRedeclarationRequest{
+              VD, VD->getDeclContext()->getSelfNominalTypeDecl()},
+          {});
 
       // Compute access level.
       (void) VD->getFormalAccess();

lib/Sema/TypeCheckStmt.cpp

@@ -281,6 +281,11 @@ namespace {
         // Otherwise, it'll have been separately type-checked.
         if (!CE->isSeparatelyTypeChecked()) {
           SetLocalDiscriminators innerVisitor;
+          if (auto params = CE->getParameters()) {
+            for (auto *param : *params) {
+              innerVisitor.setLocalDiscriminator(param);
+            }
+          }
           CE->getBody()->walk(innerVisitor);
         }
 
@@ -332,16 +337,22 @@ namespace {
 
     /// Set the local discriminator for a named declaration.
     void setLocalDiscriminator(ValueDecl *valueDecl) {
-      if (valueDecl->hasLocalDiscriminator() &&
-          valueDecl->getRawLocalDiscriminator() ==
-              ValueDecl::InvalidDiscriminator) {
-        // Assign the next discriminator.
-        Identifier name = valueDecl->getBaseIdentifier();
-        auto &discriminator = DeclDiscriminators[name];
-        if (discriminator < InitialDiscriminator)
-          discriminator = InitialDiscriminator;
-
-        valueDecl->setLocalDiscriminator(discriminator++);
+      if (valueDecl->hasLocalDiscriminator()) {
+        if (valueDecl->getRawLocalDiscriminator() ==
+            ValueDecl::InvalidDiscriminator) {
+          // Assign the next discriminator.
+          Identifier name = valueDecl->getBaseIdentifier();
+          auto &discriminator = DeclDiscriminators[name];
+          if (discriminator < InitialDiscriminator)
+            discriminator = InitialDiscriminator;
+
+          valueDecl->setLocalDiscriminator(discriminator++);
+        } else {
+          // Assign the next discriminator.
+          Identifier name = valueDecl->getBaseIdentifier();
+          auto &discriminator = DeclDiscriminators[name];
+          discriminator = std::max(discriminator, std::max(InitialDiscriminator, valueDecl->getLocalDiscriminator() + 1));
+        }
       }
 
       // If this is a property wrapper, check for projected storage.

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -547,6 +547,21 @@ void SymbolGraph::serialize(llvm::json::OStream &OS) {
       }
     });
 
+#ifndef NDEBUG
+    // FIXME (solver-based-verification-sorting): In assert builds sort the
+    // edges so we get consistent symbol graph output. This allows us to compare
+    // the string representation of the symbolgraph between the solver-based
+    // and AST-based result.
+    // This can be removed once the AST-based cursor info has been removed.
+    SmallVector<Edge> Edges(this->Edges.begin(), this->Edges.end());
+    std::sort(Edges.begin(), Edges.end(), [](const Edge &LHS, const Edge &RHS) {
+      SmallString<256> LHSTargetUSR, RHSTargetUSR;
+      LHS.Target.getUSR(LHSTargetUSR);
+      RHS.Target.getUSR(RHSTargetUSR);
+      return LHSTargetUSR < RHSTargetUSR;
+    });
+#endif
+
     OS.attributeArray("relationships", [&](){
       for (const auto &Relationship : Edges) {
         Relationship.serialize(OS);

test/SourceKit/CursorInfo/cursor_infer_nonmutating_from_property_wrapper.swift

@@ -0,0 +1,17 @@
+func test() {
+  struct MyStruct {
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):16 %s -- %s | %FileCheck %s
+    @State var myState
+  }
+}
+
+@propertyWrapper struct State {
+    init(initialValue value: Int) {}
+
+    public var wrappedValue: Int {
+        get { return 1 }
+        nonmutating set { }
+    }
+}
+
+// CHECK: <Declaration>@<Type usr="s:46cursor_infer_nonmutating_from_property_wrapper5StateV">State</Type> var myState: &lt;&lt;error type&gt;&gt; { get nonmutating set }</Declaration>

test/SourceKit/CursorInfo/cursor_local_var_with_property_wrapper.swift

@@ -0,0 +1,9 @@
+@propertyWrapper struct State {
+    public init(initialValue value: Int)
+    public var wrappedValue: Int { get nonmutating set }
+}
+
+func loadPage() {
+  // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):14 %s -- %s
+  @State var pageListener: Int
+}

test/SourceKit/CursorInfo/cursor_on_closure_parameter.swift

@@ -0,0 +1,8 @@
+func save(_ record: Int, _ x: (String) -> Void) {}
+
+func test() {
+  let record = 2
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):19 %s -- %s
+  save(record) { (record) in
+  }
+}

test/SourceKit/CursorInfo/cursor_on_var_in_extension_in_func.swift

@@ -0,0 +1,6 @@
+func foo() {
+  extension XXX {
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):9 %s -- %s
+    var hex 
+  }
+}

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