[IDE] Two more preparation commits for solver-based cursor info

include/swift/AST/ASTContext.h

@@ -289,6 +289,11 @@ class ASTContext final {
 
   ide::TypeCheckCompletionCallback *CompletionCallback = nullptr;
 
+  /// A callback that will be called when the constraint system found a
+  /// solution. Called multiple times if the constraint system has ambiguous
+  /// solutions.
+  ide::TypeCheckCompletionCallback *SolutionCallback = nullptr;
+
   /// The request-evaluator that is used to process various requests.
   Evaluator evaluator;
 

include/swift/IDE/SelectedOverloadInfo.h

@@ -0,0 +1,45 @@
+//===--- SelectedOverloadInfo.h -------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_IDE_SWIFTCOMPLETIONINFO_H
+#define SWIFT_IDE_SWIFTCOMPLETIONINFO_H
+
+#include "swift/AST/Decl.h"
+#include "swift/Sema/ConstraintSystem.h"
+
+namespace swift {
+namespace ide {
+
+using namespace swift::constraints;
+
+/// Information that \c getSelectedOverloadInfo gathered about a
+/// \c SelectedOverload.
+struct SelectedOverloadInfo {
+  /// The function that is being called or the value that is being accessed.
+  ValueDecl *Value = nullptr;
+  /// For a function, type of the called function itself (not its result type),
+  /// for an arbitrary value the type of that value.
+  Type ValueTy;
+  /// The type on which the overload is being accessed. \c null if it does not
+  /// have a base type, e.g. for a free function.
+  Type BaseTy;
+};
+
+/// Extract additional information about the overload that is being called by
+/// \p CalleeLocator.
+SelectedOverloadInfo getSelectedOverloadInfo(const Solution &S,
+                                             ConstraintLocator *CalleeLocator);
+
+} // end namespace ide
+} // end namespace swift
+
+#endif // SWIFT_IDE_SWIFTCOMPLETIONINFO_H

include/swift/IDE/Utils.h

@@ -18,6 +18,7 @@
 #include "swift/AST/DeclNameLoc.h"
 #include "swift/AST/Effects.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/Expr.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/IDE/SourceEntityWalker.h"
 #include "swift/Parse/Token.h"
@@ -740,7 +741,7 @@ bool isDeclOverridable(ValueDecl *D);
 /// one in `SomeType`. Contrast that to `type(of: foo).classMethod()` where
 /// `classMethod` could be any `classMethod` up or down the hierarchy from the
 /// type of the \p Base expression.
-bool isDynamicRef(Expr *Base, ValueDecl *D);
+bool isDynamicRef(Expr *Base, ValueDecl *D, llvm::function_ref<Type(Expr *)> getType = [](Expr *E) { return E->getType(); });
 
 /// Adds the resolved nominal types of \p Base to \p Types.
 void getReceiverType(Expr *Base,

lib/IDE/ArgumentCompletion.cpp

@@ -13,6 +13,7 @@
 #include "swift/IDE/ArgumentCompletion.h"
 #include "swift/IDE/CodeCompletion.h"
 #include "swift/IDE/CompletionLookup.h"
+#include "swift/IDE/SelectedOverloadInfo.h"
 #include "swift/Sema/ConstraintSystem.h"
 #include "swift/Sema/IDETypeChecking.h"
 
@@ -78,98 +79,6 @@ bool ArgumentTypeCheckCompletionCallback::addPossibleParams(
   return ShowGlobalCompletions;
 }
 
-/// Information that \c getSelectedOverloadInfo gathered about a
-/// \c SelectedOverload.
-struct SelectedOverloadInfo {
-  /// The function that is being called.
-  ValueDecl *FuncD = nullptr;
-  /// The type of the called function itself (not its result type)
-  Type FuncTy;
-  /// The type on which the function is being called. \c null if the function is
-  /// a free function.
-  Type CallBaseTy;
-};
-
-/// Extract additional information about the overload that is being called by
-/// \p CalleeLocator.
-SelectedOverloadInfo getSelectedOverloadInfo(const Solution &S,
-                                             ConstraintLocator *CalleeLocator) {
-  auto &CS = S.getConstraintSystem();
-
-  SelectedOverloadInfo Result;
-
-  auto SelectedOverload = S.getOverloadChoiceIfAvailable(CalleeLocator);
-  if (!SelectedOverload) {
-    return Result;
-  }
-
-  switch (SelectedOverload->choice.getKind()) {
-  case OverloadChoiceKind::KeyPathApplication:
-  case OverloadChoiceKind::Decl:
-  case OverloadChoiceKind::DeclViaDynamic:
-  case OverloadChoiceKind::DeclViaBridge:
-  case OverloadChoiceKind::DeclViaUnwrappedOptional: {
-    Result.CallBaseTy = SelectedOverload->choice.getBaseType();
-    if (Result.CallBaseTy) {
-      Result.CallBaseTy = S.simplifyType(Result.CallBaseTy)->getRValueType();
-    }
-
-    Result.FuncD = SelectedOverload->choice.getDeclOrNull();
-    Result.FuncTy =
-        S.simplifyTypeForCodeCompletion(SelectedOverload->adjustedOpenedType);
-
-    // For completion as the arg in a call to the implicit [keypath: _]
-    // subscript the solver can't know what kind of keypath is expected without
-    // an actual argument (e.g. a KeyPath vs WritableKeyPath) so it ends up as a
-    // hole. Just assume KeyPath so we show the expected keypath's root type to
-    // users rather than '_'.
-    if (SelectedOverload->choice.getKind() ==
-        OverloadChoiceKind::KeyPathApplication) {
-      auto Params = Result.FuncTy->getAs<AnyFunctionType>()->getParams();
-      if (Params.size() == 1 &&
-          Params[0].getPlainType()->is<UnresolvedType>()) {
-        auto *KPDecl = CS.getASTContext().getKeyPathDecl();
-        Type KPTy =
-            KPDecl->mapTypeIntoContext(KPDecl->getDeclaredInterfaceType());
-        Type KPValueTy = KPTy->castTo<BoundGenericType>()->getGenericArgs()[1];
-        KPTy = BoundGenericType::get(KPDecl, Type(),
-                                     {Result.CallBaseTy, KPValueTy});
-        Result.FuncTy =
-            FunctionType::get({Params[0].withType(KPTy)}, KPValueTy);
-      }
-    }
-    break;
-  }
-  case OverloadChoiceKind::KeyPathDynamicMemberLookup: {
-    auto *fnType = SelectedOverload->adjustedOpenedType->castTo<FunctionType>();
-    assert(fnType->getParams().size() == 1 &&
-           "subscript always has one argument");
-    // Parameter type is KeyPath<T, U> where `T` is a root type
-    // and U is a leaf type (aka member type).
-    auto keyPathTy =
-        fnType->getParams()[0].getPlainType()->castTo<BoundGenericType>();
-
-    auto *keyPathDecl = keyPathTy->getAnyNominal();
-    assert(isKnownKeyPathType(keyPathTy) &&
-           "parameter is supposed to be a keypath");
-
-    auto KeyPathDynamicLocator = CS.getConstraintLocator(
-        CalleeLocator, LocatorPathElt::KeyPathDynamicMember(keyPathDecl));
-    Result = getSelectedOverloadInfo(S, KeyPathDynamicLocator);
-    break;
-  }
-  case OverloadChoiceKind::DynamicMemberLookup:
-  case OverloadChoiceKind::TupleIndex:
-    // If it's DynamicMemberLookup, we don't know which function is being
-    // called, so we can't extract any information from it.
-    // TupleIndex isn't a function call and is not relevant for argument
-    // completion because it doesn't take arguments.
-    break;
-  }
-
-  return Result;
-}
-
 void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   Type ExpectedTy = getTypeForCompletion(S, CompletionExpr);
 
@@ -246,9 +155,9 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
 
   // If this is a duplicate of any other result, ignore this solution.
   if (llvm::any_of(Results, [&](const Result &R) {
-        return R.FuncD == Info.FuncD &&
-               nullableTypesEqual(R.FuncTy, Info.FuncTy) &&
-               nullableTypesEqual(R.BaseType, Info.CallBaseTy) &&
+        return R.FuncD == Info.Value &&
+               nullableTypesEqual(R.FuncTy, Info.ValueTy) &&
+               nullableTypesEqual(R.BaseType, Info.BaseTy) &&
                R.ParamIdx == ParamIdx &&
                R.IsNoninitialVariadic == IsNoninitialVariadic;
       })) {
@@ -258,9 +167,9 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
   getSolutionSpecificVarTypes(S, SolutionSpecificVarTypes);
 
-  Results.push_back({ExpectedTy, isa<SubscriptExpr>(ParentCall), Info.FuncD,
-                     Info.FuncTy, ArgIdx, ParamIdx, std::move(ClaimedParams),
-                     IsNoninitialVariadic, Info.CallBaseTy, HasLabel, IsAsync,
+  Results.push_back({ExpectedTy, isa<SubscriptExpr>(ParentCall), Info.Value,
+                     Info.ValueTy, ArgIdx, ParamIdx, std::move(ClaimedParams),
+                     IsNoninitialVariadic, Info.BaseTy, HasLabel, IsAsync,
                      SolutionSpecificVarTypes});
 }
 

lib/IDE/CMakeLists.txt

@@ -29,6 +29,7 @@ add_swift_host_library(swiftIDE STATIC
   ModuleInterfacePrinting.cpp
   PostfixCompletion.cpp
   REPLCodeCompletion.cpp
+  SelectedOverloadInfo.cpp
   SourceEntityWalker.cpp
   SwiftSourceDocInfo.cpp
   SyntaxModel.cpp

lib/IDE/SelectedOverloadInfo.cpp

@@ -0,0 +1,94 @@
+//===--- SelectedOverloadInfo.cpp -----------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/IDE/SelectedOverloadInfo.h"
+
+using namespace swift::ide;
+
+SelectedOverloadInfo
+swift::ide::getSelectedOverloadInfo(const Solution &S,
+                                    ConstraintLocator *CalleeLocator) {
+  auto &CS = S.getConstraintSystem();
+
+  SelectedOverloadInfo Result;
+
+  auto SelectedOverload = S.getOverloadChoiceIfAvailable(CalleeLocator);
+  if (!SelectedOverload) {
+    return Result;
+  }
+
+  switch (SelectedOverload->choice.getKind()) {
+  case OverloadChoiceKind::KeyPathApplication:
+  case OverloadChoiceKind::Decl:
+  case OverloadChoiceKind::DeclViaDynamic:
+  case OverloadChoiceKind::DeclViaBridge:
+  case OverloadChoiceKind::DeclViaUnwrappedOptional: {
+    Result.BaseTy = SelectedOverload->choice.getBaseType();
+    if (Result.BaseTy) {
+      Result.BaseTy = S.simplifyType(Result.BaseTy)->getRValueType();
+    }
+
+    Result.Value = SelectedOverload->choice.getDeclOrNull();
+    Result.ValueTy =
+        S.simplifyTypeForCodeCompletion(SelectedOverload->adjustedOpenedType);
+
+    // For completion as the arg in a call to the implicit [keypath: _]
+    // subscript the solver can't know what kind of keypath is expected without
+    // an actual argument (e.g. a KeyPath vs WritableKeyPath) so it ends up as a
+    // hole. Just assume KeyPath so we show the expected keypath's root type to
+    // users rather than '_'.
+    if (SelectedOverload->choice.getKind() ==
+        OverloadChoiceKind::KeyPathApplication) {
+      auto Params = Result.ValueTy->getAs<AnyFunctionType>()->getParams();
+      if (Params.size() == 1 &&
+          Params[0].getPlainType()->is<UnresolvedType>()) {
+        auto *KPDecl = CS.getASTContext().getKeyPathDecl();
+        Type KPTy =
+            KPDecl->mapTypeIntoContext(KPDecl->getDeclaredInterfaceType());
+        Type KPValueTy = KPTy->castTo<BoundGenericType>()->getGenericArgs()[1];
+        KPTy =
+            BoundGenericType::get(KPDecl, Type(), {Result.BaseTy, KPValueTy});
+        Result.ValueTy =
+            FunctionType::get({Params[0].withType(KPTy)}, KPValueTy);
+      }
+    }
+    break;
+  }
+  case OverloadChoiceKind::KeyPathDynamicMemberLookup: {
+    auto *fnType = SelectedOverload->adjustedOpenedType->castTo<FunctionType>();
+    assert(fnType->getParams().size() == 1 &&
+           "subscript always has one argument");
+    // Parameter type is KeyPath<T, U> where `T` is a root type
+    // and U is a leaf type (aka member type).
+    auto keyPathTy =
+        fnType->getParams()[0].getPlainType()->castTo<BoundGenericType>();
+
+    auto *keyPathDecl = keyPathTy->getAnyNominal();
+    assert(isKnownKeyPathType(keyPathTy) &&
+           "parameter is supposed to be a keypath");
+
+    auto KeyPathDynamicLocator = CS.getConstraintLocator(
+        CalleeLocator, LocatorPathElt::KeyPathDynamicMember(keyPathDecl));
+    Result = getSelectedOverloadInfo(S, KeyPathDynamicLocator);
+    break;
+  }
+  case OverloadChoiceKind::DynamicMemberLookup:
+  case OverloadChoiceKind::TupleIndex:
+    // If it's DynamicMemberLookup, we don't know which function is being
+    // called, so we can't extract any information from it.
+    // TupleIndex isn't a function call and is not relevant for argument
+    // completion because it doesn't take arguments.
+    break;
+  }
+
+  return Result;
+}

lib/IDE/Utils.cpp

@@ -921,7 +921,7 @@ bool swift::ide::isDeclOverridable(ValueDecl *D) {
   return true;
 }
 
-bool swift::ide::isDynamicRef(Expr *Base, ValueDecl *D) {
+bool swift::ide::isDynamicRef(Expr *Base, ValueDecl *D, llvm::function_ref<Type(Expr *)> getType) {
   if (!isDeclOverridable(D))
     return false;
 
@@ -940,7 +940,7 @@ bool swift::ide::isDynamicRef(Expr *Base, ValueDecl *D) {
 
   // `type(of: foo).staticOrClassMethod()`. A static method may be "dynamic"
   // here, but not if the instance type is a struct/enum.
-  if (auto IT = Base->getType()->getAs<MetatypeType>()) {
+  if (auto IT = getType(Base)->getAs<MetatypeType>()) {
     auto InstanceType = IT->getInstanceType();
     if (InstanceType->getStructOrBoundGenericStruct() ||
         InstanceType->getEnumOrBoundGenericEnum())

lib/Sema/CSSolver.cpp

@@ -1353,6 +1353,24 @@ Optional<std::vector<Solution>> ConstraintSystem::solve(
     }
   };
 
+  auto reportSolutionsToSolutionCallback = [&](const SolutionResult &result) {
+    if (!getASTContext().SolutionCallback) {
+      return;
+    }
+    switch (result.getKind()) {
+    case SolutionResult::Success:
+      getASTContext().SolutionCallback->sawSolution(result.getSolution());
+      break;
+    case SolutionResult::Ambiguous:
+      for (auto &solution : result.getAmbiguousSolutions()) {
+        getASTContext().SolutionCallback->sawSolution(solution);
+      }
+      break;
+    default:
+      break;
+    }
+  };
+
   // Take up to two attempts at solving the system. The first attempts to
   // solve a system that is expected to be well-formed, the second kicks in
   // when there is an error and attempts to salvage an ill-formed program.
@@ -1365,6 +1383,7 @@ Optional<std::vector<Solution>> ConstraintSystem::solve(
     case SolutionResult::Success: {
       // Return the successful solution.
       dumpSolutions(solution);
+      reportSolutionsToSolutionCallback(solution);
       std::vector<Solution> result;
       result.push_back(std::move(solution).takeSolution());
       return std::move(result);
@@ -1394,6 +1413,7 @@ Optional<std::vector<Solution>> ConstraintSystem::solve(
       // If salvaging produced an ambiguous result, it has already been
       // diagnosed.
       if (stage == 1) {
+        reportSolutionsToSolutionCallback(solution);
         solution.markAsDiagnosed();
         return None;
       }
@@ -1412,12 +1432,14 @@ Optional<std::vector<Solution>> ConstraintSystem::solve(
 
     case SolutionResult::UndiagnosedError:
       if (shouldSuppressDiagnostics()) {
+        reportSolutionsToSolutionCallback(solution);
         solution.markAsDiagnosed();
         return None;
       }
 
       if (stage == 1) {
         diagnoseFailureFor(target);
+        reportSolutionsToSolutionCallback(solution);
         solution.markAsDiagnosed();
         return None;
       }