[CodeCompletion] Compute type relations for the function calls selected by ArgumentCompletion

Author: ahoppen

include/swift/IDE/ArgumentCompletion.h

@@ -25,6 +25,8 @@ class ArgumentTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
   struct Result {
     /// The type associated with the code completion expression itself.
     Type ExpectedType;
+    /// The expected return type of the function call.
+    Type ExpectedCallType;
     /// True if this is a subscript rather than a function call.
     bool IsSubscript;
     /// The FuncDecl or SubscriptDecl associated with the call.

lib/IDE/ArgumentCompletion.cpp

@@ -78,6 +78,44 @@ bool ArgumentTypeCheckCompletionCallback::addPossibleParams(
   return ShowGlobalCompletions;
 }
 
+/// Applies heuristic to determine whether the result type of \p E is
+/// unconstrained, that is if the constraint system is satisfiable for any
+/// result type of \p E.
+static bool isExpressionResultTypeUnconstrained(const Solution &S, Expr *E) {
+  ConstraintSystem &CS = S.getConstraintSystem();
+  if (auto ParentExpr = CS.getParentExpr(E)) {
+    if (auto Assign = dyn_cast<AssignExpr>(ParentExpr)) {
+      if (isa<DiscardAssignmentExpr>(Assign->getDest())) {
+        // _ = <expr> is unconstrained
+        return true;
+      }
+    } else if (isa<RebindSelfInConstructorExpr>(ParentExpr)) {
+      // super.init() is unconstrained (it always produces the correct result
+      // by definition)
+      return true;
+    }
+  }
+  auto targetIt = S.solutionApplicationTargets.find(E);
+  if (targetIt == S.solutionApplicationTargets.end()) {
+    return false;
+  }
+  auto target = targetIt->second;
+  assert(target.kind == SolutionApplicationTarget::Kind::expression);
+  switch (target.getExprContextualTypePurpose()) {
+  case CTP_Unused:
+    // If we aren't using the contextual type, its unconstrained by definition.
+    return true;
+  case CTP_Initialization: {
+    // let x = <expr> is unconstrained
+    auto contextualType = target.getExprContextualType();
+    return !contextualType || contextualType->is<UnresolvedType>();
+  }
+  default:
+    // Assume that it's constrained by default.
+    return false;
+  }
+}
+
 void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   Type ExpectedTy = getTypeForCompletion(S, CompletionExpr);
 
@@ -111,6 +149,11 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   }
   auto ArgIdx = ArgInfo->completionIdx;
 
+  Type ExpectedCallType;
+  if (!isExpressionResultTypeUnconstrained(S, ParentCall)) {
+    ExpectedCallType = getTypeForCompletion(S, ParentCall);
+  }
+
   auto *CallLocator = CS.getConstraintLocator(ParentCall);
   auto *CalleeLocator = S.getCalleeLocator(CallLocator);
 
@@ -189,10 +232,10 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   if (Info.ValueTy) {
     FuncTy = Info.ValueTy->lookThroughAllOptionalTypes()->getAs<AnyFunctionType>();
   }
-  Results.push_back({ExpectedTy, isa<SubscriptExpr>(ParentCall), Info.Value,
-                     FuncTy, ArgIdx, ParamIdx, std::move(ClaimedParams),
-                     IsNoninitialVariadic, Info.BaseTy, HasLabel, IsAsync,
-                     SolutionSpecificVarTypes});
+  Results.push_back({ExpectedTy, ExpectedCallType,
+                     isa<SubscriptExpr>(ParentCall), Info.Value, FuncTy, ArgIdx,
+                     ParamIdx, std::move(ClaimedParams), IsNoninitialVariadic,
+                     Info.BaseTy, HasLabel, IsAsync, SolutionSpecificVarTypes});
 }
 
 void ArgumentTypeCheckCompletionCallback::deliverResults(
@@ -205,10 +248,18 @@ void ArgumentTypeCheckCompletionCallback::deliverResults(
 
   // Perform global completion as a fallback if we don't have any results.
   bool shouldPerformGlobalCompletion = Results.empty();
+  SmallVector<Type, 4> ExpectedCallTypes;
+  for (auto &Result : Results) {
+    ExpectedCallTypes.push_back(Result.ExpectedCallType);
+  }
+
   SmallVector<Type, 8> ExpectedTypes;
 
   if (IncludeSignature && !Results.empty()) {
     Lookup.setHaveLParen(true);
+    Lookup.setExpectedTypes(ExpectedCallTypes,
+                            /*isImplicitSingleExpressionReturn=*/false);
+
     for (auto &Result : Results) {
       auto SemanticContext = SemanticContextKind::None;
       NominalTypeDecl *BaseNominal = nullptr;

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -623,6 +623,7 @@ bool TypeChecker::typeCheckForCodeCompletion(
 
     // If solve failed to generate constraints or with some other
     // issue, we need to fallback to type-checking a sub-expression.
+    cs.setSolutionApplicationTarget(target.getAsExpr(), target);
     if (!cs.solveForCodeCompletion(target, solutions))
       return CompletionResult::Fallback;
 

test/IDE/complete_call_arg.swift

@@ -1231,7 +1231,7 @@ private extension Sequence {
   func SubstitutableBaseTyOfSubscript<T: Comparable>(by keyPath: KeyPath<Element, T>) -> [Element] {
     return sorted { a, b in a[#^GENERICBASE_SUB^#] }
     // GENERICBASE_SUB: Begin completions, 1 item
-    // GENERICBASE_SUB: Pattern/CurrNominal/Flair[ArgLabels]: ['[']{#keyPath: KeyPath<Self.Element, Value>#}[']'][#Value#];
+    // GENERICBASE_SUB: Pattern/CurrNominal/Flair[ArgLabels]/TypeRelation[Convertible]: ['[']{#keyPath: KeyPath<Self.Element, Value>#}[']'][#Value#];
     // GENERICBASE_SUB: End completions
   }
 }

test/IDE/complete_subscript.swift

@@ -162,7 +162,7 @@ func testSubcscriptTuple(val: (x: Int, String)) {
 }
 
 struct HasSettableSub {
-    subscript(a: String) -> Any {
+    subscript(a: String) -> Int {
         get { return 1 }
         set { }
     }
@@ -174,6 +174,6 @@ func testSettableSub(x: inout HasSettableSub) {
 }
 // SETTABLE_SUBSCRIPT: Begin completions
 // SETTABLE_SUBSCRIPT-DAG: Pattern/CurrNominal/Flair[ArgLabels]: ['[']{#keyPath: KeyPath<HasSettableSub, Value>#}[']'][#Value#];
-// SETTABLE_SUBSCRIPT-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]: ['[']{#(a): String#}[']'][#@lvalue Any#];
+// SETTABLE_SUBSCRIPT-DAG: Decl[Subscript]/CurrNominal/Flair[ArgLabels]/TypeRelation[Convertible]: ['[']{#(a): String#}[']'][#@lvalue Int#];
 // SETTABLE_SUBSCRIPT-DAG: Decl[LocalVar]/Local/TypeRelation[Convertible]: local[#String#]; name=local
 // SETTABLE_SUBSCRIPT: End completions