[SourceKit] Allow generation of cursor info for declarations from solutions that haven’t aren’t applied to the AST ye

ahoppen · ahoppen · commit 3ef654352848 · 2024-03-15T20:48:21.000-07:00
This has two benefits:
1. We can now report ambiguous variable types
2. We are more robust in the generation of results for declarations inside closures. If the closure has an error, we won’t apply the solution to the AST and thus any cursor info that tried to get types out of the AST would fail.

rdar://123845208
diff --git a/include/swift/IDE/Utils.h b/include/swift/IDE/Utils.h
@@ -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; }
 
diff --git a/lib/IDE/CursorInfo.cpp b/lib/IDE/CursorInfo.cpp
@@ -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);
     }
   };
 
@@ -313,13 +319,42 @@ class CursorInfoTypeCheckSolutionCallback : public TypeCheckCompletionCallback {
   }
 
   void sawSolutionImpl(const Solution &S) override {
-    auto &CS = S.getConstraintSystem();
-    auto ResolveExpr = getExprToResolve();
-    if (!ResolveExpr) {
+    NodeFinder Finder(DC, ResolveLoc);
+    Finder.resolve();
+    auto Result = Finder.takeResult();
+    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;
+    }
+    }
+  }
+
+  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});
+  }
 
-    auto Locator = CS.getConstraintLocator(ResolveExpr);
+  void addCursorInfoResultForExpr(Expr *ExprToResolve, const Solution &S) {
+    auto &CS = S.getConstraintSystem();
+
+    auto Locator = CS.getConstraintLocator(ExprToResolve);
     auto CalleeLocator = S.getCalleeLocator(Locator);
     auto OverloadInfo = getSelectedOverloadInfo(S, CalleeLocator);
     if (!OverloadInfo.ValueRef) {
@@ -337,9 +372,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,36 +404,22 @@ 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()) {
@@ -434,7 +457,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 +468,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;
diff --git a/lib/IDE/IDERequests.cpp b/lib/IDE/IDERequests.cpp
@@ -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=*/{});
 
diff --git a/lib/IDE/SourceEntityWalker.cpp b/lib/IDE/SourceEntityWalker.cpp
@@ -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)
diff --git a/test/SourceKit/CursorInfo/closure_with_invalid_var_reference.swift b/test/SourceKit/CursorInfo/closure_with_invalid_var_reference.swift
@@ -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>
diff --git a/test/SourceKit/CursorInfo/cursor_ambiguous.swift b/test/SourceKit/CursorInfo/cursor_ambiguous.swift
@@ -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
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp b/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
@@ -770,7 +770,13 @@ static StringRef getModuleName(const ValueDecl *VD,
 }
 
 struct DeclInfo {
-  const ValueDecl *VD;
+  ValueDecl *VD;
+  /// If not null, a solution-specific interface type of `VD`. This allows
+  /// us to produce results for declarations where the soluion has not been
+  /// applied to the AST, eg. because the declaration has an ambiguous type
+  /// or because it occurs inside a closure that has an error, which prevents
+  /// the constraint system solution to be applied to the AST.
+  Type SolutionSpecificInterfaceType;
   Type ContainerType;
   bool IsRef;
   bool IsDynamic;
@@ -785,11 +791,13 @@ struct DeclInfo {
   /// Whether the \c VD is in a synthesized extension of \c BaseType
   bool InSynthesizedExtension = false;
 
-  DeclInfo(const ValueDecl *VD, Type ContainerType, bool IsRef, bool IsDynamic,
+  DeclInfo(ValueDecl *VD, Type SolutionSpecificInterfaceType,
+           Type ContainerType, bool IsRef, bool IsDynamic,
            ArrayRef<NominalTypeDecl *> ReceiverTypes,
            const CompilerInvocation &Invoc)
-      : VD(VD), ContainerType(ContainerType), IsRef(IsRef),
-        IsDynamic(IsDynamic), ReceiverTypes(ReceiverTypes) {
+      : VD(VD), SolutionSpecificInterfaceType(SolutionSpecificInterfaceType),
+        ContainerType(ContainerType), IsRef(IsRef), IsDynamic(IsDynamic),
+        ReceiverTypes(ReceiverTypes) {
     if (VD == nullptr)
       return;
 
@@ -939,6 +947,18 @@ fillSymbolInfo(CursorSymbolInfo &Symbol, const DeclInfo &DInfo,
                const CompilerInvocation &Invoc,
                ArrayRef<ImmutableTextSnapshotRef> PreviousSnaps,
                llvm::BumpPtrAllocator &Allocator) {
+  // Override the type of `DInfo.VD` while retrieving symbol information.
+  // Ideall, we would pass a customizable `getInterfaceType` function to all
+  // functions that operate on VD but doing so is not viable - USR generation,
+  // AST printing, symbol graph and probably more have been designed with the
+  // assumption that VD->getInterfaceType() returns the fixed interface type
+  // of the declaration.
+  Type OriginalInterfaceType = DInfo.VD->getInterfaceType();
+  if (DInfo.SolutionSpecificInterfaceType) {
+    DInfo.VD->setInterfaceType(DInfo.SolutionSpecificInterfaceType);
+  }
+  SWIFT_DEFER { DInfo.VD->setInterfaceType(OriginalInterfaceType); };
+
   SmallString<256> Buffer;
   SmallVector<StringRef, 4> Strings;
   llvm::raw_svector_ostream OS(Buffer);
@@ -1199,7 +1219,7 @@ static bool addCursorInfoForLiteral(
   }
 
   auto &Symbol = Data.Symbols.emplace_back();
-  DeclInfo Info(Decl, nullptr, true, false, {}, CompInvoc);
+  DeclInfo Info(Decl, nullptr, nullptr, true, false, {}, CompInvoc);
   auto Err = fillSymbolInfo(Symbol, Info, CursorLoc, false, Lang, CompInvoc,
                             PreviousSnaps, Data.Allocator);
 
@@ -1218,9 +1238,10 @@ addCursorInfoForDecl(CursorInfoData &Data, ResolvedValueRefCursorInfoPtr Info,
                      SwiftLangSupport &Lang, const CompilerInvocation &Invoc,
                      std::string &Diagnostic,
                      ArrayRef<ImmutableTextSnapshotRef> PreviousSnaps) {
-  DeclInfo OrigInfo(Info->getValueD(), Info->getContainerType(), Info->isRef(),
-                    Info->isDynamic(), Info->getReceiverTypes(), Invoc);
-  DeclInfo CtorTypeInfo(Info->getCtorTyRef(), Type(), true, false,
+  DeclInfo OrigInfo(Info->getValueD(), Info->getSolutionSpecificInterfaceType(),
+                    Info->getContainerType(), Info->isRef(), Info->isDynamic(),
+                    Info->getReceiverTypes(), Invoc);
+  DeclInfo CtorTypeInfo(Info->getCtorTyRef(), Type(), Type(), true, false,
                         ArrayRef<NominalTypeDecl *>(), Invoc);
   DeclInfo &MainInfo = CtorTypeInfo.VD ? CtorTypeInfo : OrigInfo;
   if (MainInfo.Unavailable) {
@@ -1257,7 +1278,7 @@ addCursorInfoForDecl(CursorInfoData &Data, ResolvedValueRefCursorInfoPtr Info,
   if (!Info->isRef()) {
     for (auto D : Info->getShorthandShadowedDecls()) {
       CursorSymbolInfo &SymbolInfo = Data.Symbols.emplace_back();
-      DeclInfo DInfo(D, Type(), /*IsRef=*/true, /*IsDynamic=*/false,
+      DeclInfo DInfo(D, Type(), Type(), /*IsRef=*/true, /*IsDynamic=*/false,
                      ArrayRef<NominalTypeDecl *>(), Invoc);
       if (auto Err =
               fillSymbolInfo(SymbolInfo, DInfo, Info->getLoc(), AddSymbolGraph,
@@ -2082,7 +2103,7 @@ void SwiftLangSupport::getCursorInfo(
               const_cast<ValueDecl *>(Entity.Dcl),
               /*CtorTyRef=*/nullptr,
               /*ExtTyRef=*/nullptr, Entity.IsRef,
-              /*Ty=*/Type(),
+              /*SolutionSpecificInterfaceType=*/Type(),
               /*ContainerType=*/Type(),
               /*CustomAttrRef=*/std::nullopt,
               /*IsKeywordArgument=*/false,
@@ -2143,7 +2164,13 @@ void SwiftLangSupport::getCursorInfo(
     // AST based completion *always* produces a result
     bool NoResults = Res.isError();
     if (Res.isValue()) {
-      NoResults = Res.value().Symbols.empty();
+      NoResults |= Res.value().Symbols.empty();
+      bool AllResultsHaveErrorType =
+          llvm::all_of(Res.value().Symbols, [](const CursorSymbolInfo &Symbol) {
+            // USR of the AST's error type if no type could be resolved.
+            return Symbol.TypeUSR == "$sXeD";
+          });
+      NoResults |= AllResultsHaveErrorType;
     }
     if (!NoResults || !InputBuffer) {
       Receiver(Res);
@@ -2383,7 +2410,7 @@ static void resolveCursorFromUSR(
             /*CtorTyRef=*/nullptr,
             /*ExtTyRef=*/nullptr,
             /*IsRef=*/false,
-            /*Ty=*/Type(), ContainerType,
+            /*SolutionSpecificInterfaceType=*/Type(), ContainerType,
             /*CustomAttrRef=*/std::nullopt,
             /*IsKeywordArgument=*/false,
             /*IsDynamic=*/false,
