swiftlang · ahoppen · Mar 25, 2024 · Mar 16, 2024
@@ -171,7 +171,7 @@ struct ResolvedValueRefCursorInfo : public ResolvedCursorInfo {
   TypeDecl *CtorTyRef = nullptr;
   ExtensionDecl *ExtTyRef = nullptr;
   bool IsRef = true;
-  Type Ty;
+  Type SolutionSpecificInterfaceType;
   Type ContainerType;
   std::optional<std::pair<const CustomAttr *, Decl *>> CustomAttrRef =
       std::nullopt;
@@ -196,13 +196,15 @@ struct ResolvedValueRefCursorInfo : public ResolvedCursorInfo {
   ResolvedValueRefCursorInfo() = default;
   explicit ResolvedValueRefCursorInfo(
       SourceFile *SF, SourceLoc Loc, ValueDecl *ValueD, TypeDecl *CtorTyRef,
-      ExtensionDecl *ExtTyRef, bool IsRef, Type Ty, Type ContainerType,
+      ExtensionDecl *ExtTyRef, bool IsRef, Type SolutionSpecificInterfaceType,
+      Type ContainerType,
       std::optional<std::pair<const CustomAttr *, Decl *>> CustomAttrRef,
       bool IsKeywordArgument, bool IsDynamic,
       SmallVector<NominalTypeDecl *> ReceiverTypes,
       SmallVector<ValueDecl *> ShorthandShadowedDecls)
       : ResolvedCursorInfo(CursorInfoKind::ValueRef, SF, Loc), ValueD(ValueD),
-        CtorTyRef(CtorTyRef), ExtTyRef(ExtTyRef), IsRef(IsRef), Ty(Ty),
+        CtorTyRef(CtorTyRef), ExtTyRef(ExtTyRef), IsRef(IsRef),
+        SolutionSpecificInterfaceType(SolutionSpecificInterfaceType),
         ContainerType(ContainerType), CustomAttrRef(CustomAttrRef),
         IsKeywordArgument(IsKeywordArgument), IsDynamic(IsDynamic),
         ReceiverTypes(ReceiverTypes),
@@ -216,7 +218,9 @@ struct ResolvedValueRefCursorInfo : public ResolvedCursorInfo {
 
   bool isRef() const { return IsRef; }
 
-  Type getType() const { return Ty; }
+  Type getSolutionSpecificInterfaceType() const {
+    return SolutionSpecificInterfaceType;
+  }
 
   Type getContainerType() const { return ContainerType; }
 

@@ -282,11 +282,17 @@ class CursorInfoTypeCheckSolutionCallback : public TypeCheckCompletionCallback {
     bool IsDynamicRef;
     /// The declaration that is being referenced. Will never be \c nullptr.
     ValueDecl *ReferencedDecl;
+    /// The interface type of the referenced declaration. This might not be
+    /// stored in `ReferencedDecl->getInterfaceType()` if the declaration's
+    /// type hasn't been applied to the AST.
+    Type SolutionSpecificInterfaceType;
 
     bool operator==(const CursorInfoDeclReference &Other) const {
       return nullableTypesEqual(BaseType, Other.BaseType) &&
              IsDynamicRef == Other.IsDynamicRef &&
-             ReferencedDecl == Other.ReferencedDecl;
+             ReferencedDecl == Other.ReferencedDecl &&
+             nullableTypesEqual(SolutionSpecificInterfaceType,
+                                Other.SolutionSpecificInterfaceType);
     }
   };
 
@@ -302,24 +308,43 @@ class CursorInfoTypeCheckSolutionCallback : public TypeCheckCompletionCallback {
 
   SmallVector<CursorInfoDeclReference, 1> Results;
 
-  Expr *getExprToResolve() {
+  void sawSolutionImpl(const Solution &S) override {
     NodeFinder Finder(DC, ResolveLoc);
     Finder.resolve();
     auto Result = Finder.takeResult();
-    if (!Result || Result->getKind() != NodeFinderResultKind::Expr) {
-      return nullptr;
+    if (!Result) {
+      return;
+    }
+    switch (Result->getKind()) {
+    case NodeFinderResultKind::Decl: {
+      ValueDecl *DeclToResolve =
+          cast<NodeFinderDeclResult>(Result.get())->getDecl();
+      addCursorInfoResultForDecl(DeclToResolve, S);
+      break;
+    }
+    case NodeFinderResultKind::Expr: {
+      Expr *ExprToResolve = cast<NodeFinderExprResult>(Result.get())->getExpr();
+      addCursorInfoResultForExpr(ExprToResolve, S);
+      break;
+    }
     }
-    return cast<NodeFinderExprResult>(Result.get())->getExpr();
   }
 
-  void sawSolutionImpl(const Solution &S) override {
-    auto &CS = S.getConstraintSystem();
-    auto ResolveExpr = getExprToResolve();
-    if (!ResolveExpr) {
+  void addCursorInfoResultForDecl(ValueDecl *DeclToResolve, const Solution &S) {
+    if (!S.hasType(DeclToResolve)) {
       return;
     }
+    Type SolutionInterfaceTy =
+        S.simplifyType(S.getType(DeclToResolve))->mapTypeOutOfContext();
+
+    addResult({/*BaseType=*/nullptr, /*IsDynamicRef=*/false, DeclToResolve,
+               SolutionInterfaceTy});
+  }
+
+  void addCursorInfoResultForExpr(Expr *ExprToResolve, const Solution &S) {
+    auto &CS = S.getConstraintSystem();
 
-    auto Locator = CS.getConstraintLocator(ResolveExpr);
+    auto Locator = CS.getConstraintLocator(ExprToResolve);
     auto CalleeLocator = S.getCalleeLocator(Locator);
     auto OverloadInfo = getSelectedOverloadInfo(S, CalleeLocator);
     if (!OverloadInfo.ValueRef) {
@@ -337,9 +362,11 @@ class CursorInfoTypeCheckSolutionCallback : public TypeCheckCompletionCallback {
                             [&S](Expr *E) { return S.getResolvedType(E); });
     }
 
-    CursorInfoDeclReference NewResult = {OverloadInfo.BaseTy, IsDynamicRef,
-                                         OverloadInfo.getValue()};
+    addResult({OverloadInfo.BaseTy, IsDynamicRef, OverloadInfo.getValue(),
+               /*SolutionSpecificInterfaceType=*/Type()});
+  }
 
+  void addResult(const CursorInfoDeclReference &NewResult) {
     if (llvm::any_of(Results, [&](const CursorInfoDeclReference &R) {
           return R == NewResult;
         })) {
@@ -367,47 +394,38 @@ class CursorInfoDoneParsingCallback : public DoneParsingCallback {
                                 SourceLoc RequestedLoc)
       : DoneParsingCallback(), Consumer(Consumer), RequestedLoc(RequestedLoc) {}
 
-  std::vector<ResolvedCursorInfoPtr>
-  getDeclResult(NodeFinderDeclResult *DeclResult, SourceFile *SrcFile,
-                NodeFinder &Finder) const {
-    typeCheckDeclAndParentClosures(DeclResult->getDecl());
-    auto CursorInfo = new ResolvedValueRefCursorInfo(
-        SrcFile, RequestedLoc, DeclResult->getDecl(),
-        /*CtorTyRef=*/nullptr,
-        /*ExtTyRef=*/nullptr, /*IsRef=*/false, /*Ty=*/Type(),
-        /*ContainerType=*/Type(),
-        /*CustomAttrRef=*/std::nullopt,
-        /*IsKeywordArgument=*/false,
-        /*IsDynamic=*/false,
-        /*ReceiverTypes=*/{},
-        Finder.getShorthandShadowedDecls(DeclResult->getDecl()));
-    return {CursorInfo};
-  }
-
-  std::vector<ResolvedCursorInfoPtr>
-  getExprResult(NodeFinderExprResult *ExprResult, SourceFile *SrcFile,
-                NodeFinder &Finder) const {
-    Expr *E = ExprResult->getExpr();
-    DeclContext *DC = ExprResult->getDeclContext();
-
+private:
+  /// Shared core of `getExprResult` and `getDeclResult`.
+  std::vector<ResolvedCursorInfoPtr> getResult(ASTNode Node, DeclContext *DC,
+                                               SourceFile *SrcFile,
+                                               NodeFinder &Finder) const {
     // Type check the statemnt containing E and listen for solutions.
     CursorInfoTypeCheckSolutionCallback Callback(*DC, RequestedLoc);
     {
       llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
           DC->getASTContext().SolutionCallback, &Callback);
-      typeCheckASTNodeAtLoc(TypeCheckASTNodeAtLocContext::declContext(DC),
-                            E->getLoc());
+      if (ValueDecl *VD = getAsDecl<ValueDecl>(Node)) {
+        typeCheckDeclAndParentClosures(VD);
+      } else {
+        typeCheckASTNodeAtLoc(TypeCheckASTNodeAtLocContext::declContext(DC),
+                              Node.getStartLoc());
+      }
     }
 
     if (Callback.getResults().empty()) {
       // No results.
       return {};
     }
 
-    for (auto Info : Callback.getResults()) {
-      // Type check the referenced decls so that all their parent closures are
-      // type-checked (see comment in typeCheckDeclAndParentClosures).
-      typeCheckDeclAndParentClosures(Info.ReferencedDecl);
+    if (Node.is<Expr *>()) {
+      // If we are performing cursor info on an expression, type check the
+      // referenced decls so that all their parent closures are type-checked
+      // (see comment in typeCheckDeclAndParentClosures).
+      // When doing cursor info on a declaration, we already type checked the
+      // decl above while listening to the solution callbacks.
+      for (auto Info : Callback.getResults()) {
+        typeCheckDeclAndParentClosures(Info.ReferencedDecl);
+      }
     }
 
     // Deliver results
@@ -434,7 +452,8 @@ class CursorInfoDoneParsingCallback : public DoneParsingCallback {
       auto CursorInfo = new ResolvedValueRefCursorInfo(
           SrcFile, RequestedLoc, Res.ReferencedDecl,
           /*CtorTyRef=*/nullptr,
-          /*ExtTyRef=*/nullptr, /*IsRef=*/true, /*Ty=*/Type(),
+          /*ExtTyRef=*/nullptr, /*IsRef=*/true,
+          Res.SolutionSpecificInterfaceType,
           /*ContainerType=*/Res.BaseType,
           /*CustomAttrRef=*/std::nullopt,
           /*IsKeywordArgument=*/false, Res.IsDynamicRef, ReceiverTypes,
@@ -444,6 +463,43 @@ class CursorInfoDoneParsingCallback : public DoneParsingCallback {
     return Results;
   }
 
+public:
+  std::vector<ResolvedCursorInfoPtr>
+  getDeclResult(NodeFinderDeclResult *DeclResult, SourceFile *SrcFile,
+                NodeFinder &Finder) const {
+    std::vector<ResolvedCursorInfoPtr> Results =
+        getResult(DeclResult->getDecl(),
+                  DeclResult->getDecl()->getDeclContext(), SrcFile, Finder);
+
+    if (!Results.empty()) {
+      return Results;
+    }
+
+    // If we didn't get any solution from the constraint system, try getting the
+    // type from the decl itself. This may happen if the decl is in an inactive
+    // branch of a `#if` clause.
+    auto CursorInfo = new ResolvedValueRefCursorInfo(
+        SrcFile, RequestedLoc, DeclResult->getDecl(),
+        /*CtorTyRef=*/nullptr,
+        /*ExtTyRef=*/nullptr,
+        /*IsRef=*/false,
+        /*SolutionSpecificInterfaceType=*/Type(),
+        /*ContainerType=*/Type(),
+        /*CustomAttrRef=*/std::nullopt,
+        /*IsKeywordArgument=*/false,
+        /*IsDynamic=*/false,
+        /*ReceiverTypes=*/{},
+        Finder.getShorthandShadowedDecls(DeclResult->getDecl()));
+    return {CursorInfo};
+  }
+
+  std::vector<ResolvedCursorInfoPtr>
+  getExprResult(NodeFinderExprResult *ExprResult, SourceFile *SrcFile,
+                NodeFinder &Finder) const {
+    return getResult(ExprResult->getExpr(), ExprResult->getDeclContext(),
+                     SrcFile, Finder);
+  }
+
   void doneParsing(SourceFile *SrcFile) override {
     if (!SrcFile) {
       return;

@@ -165,7 +165,8 @@ bool CursorInfoResolver::tryResolve(ValueDecl *D, TypeDecl *CtorTyRef,
 
   CursorInfo = new ResolvedValueRefCursorInfo(
       CursorInfo->getSourceFile(), CursorInfo->getLoc(), D, CtorTyRef, ExtTyRef,
-      IsRef, Ty, ContainerType, CustomAttrRef,
+      IsRef, /*SolutionSpecificInterfaceType=*/Type(), ContainerType,
+      CustomAttrRef,
       /*IsKeywordArgument=*/false, IsDynamic, ReceiverTypes,
       /*ShorthandShadowedDecls=*/{});
 

@@ -178,7 +178,7 @@ ASTWalker::PreWalkAction SemaAnnotator::walkToDeclPreProper(Decl *D) {
         SourceLoc loc = parsedName.second;
         if (auto assocTypeDecl = proto->getAssociatedType(name)) {
           auto Continue = passReference(
-              assocTypeDecl, assocTypeDecl->getDeclaredInterfaceType(),
+              assocTypeDecl, assocTypeDecl->getInterfaceType(),
               DeclNameLoc(loc),
               ReferenceMetaData(SemaReferenceKind::TypeRef, std::nullopt));
           if (!Continue)

@@ -0,0 +1,12 @@
+func withString(body: (String) -> Void) {}
+
+func test(array: [String]) {
+  withString { element in
+    // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):9 %s -- %s | %FileCheck %s
+    let refToElement = element
+
+    _ = invalid
+  }
+}
+
+// CHECK: <Declaration>let refToElement: <Type usr="s:SS">String</Type></Declaration>
@@ -16,7 +16,19 @@ func testAmbiguousFunctionReference() {
   // LOCAL_FUNC: SECONDARY SYMBOLS END
 }
 
-
+func testAmbiguousFunctionResult() {
+  func foo() -> Int {}
+  func foo() -> String {}
+
+  // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):7 %s -- %s | %FileCheck %s --check-prefix AMBIGUOUS_FUNC_RESULT
+  let value = foo()
+  // AMBIGUOUS_FUNC_RESULT: source.lang.swift.ref.var.local
+  // AMBIGUOUS_FUNC_RESULT: <Declaration>let value: <Type usr="s:Si">Int</Type></Declaration>
+  // AMBIGUOUS_FUNC_RESULT: SECONDARY SYMBOLS BEGIN
+  // AMBIGUOUS_FUNC_RESULT: source.lang.swift.ref.var.local
+  // AMBIGUOUS_FUNC_RESULT: <Declaration>let value: <Type usr="s:SS">String</Type></Declaration>
+  // AMBIGUOUS_FUNC_RESULT: SECONDARY SYMBOLS END
+}
 
 struct TestDeduplicateResults {
   // The constraints system produces multiple solutions here for the argument type but