[CodeCompletion] Migrate expression completions to solver-based

include/swift/IDE/CompletionLookup.h

@@ -113,6 +113,12 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   /// Expected types of the code completion expression.
   ExpectedTypeContext expectedTypeContext;
 
+  /// Variables whose type was determined while type checking the code
+  /// completion expression and that are thus not recorded in the AST.
+  /// This in particular applies to params of closures that contain the code
+  /// completion token.
+  llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+
   bool CanCurrDeclContextHandleAsync = false;
   bool HaveDot = false;
   bool IsUnwrappedOptional = false;
@@ -206,6 +212,11 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
                    const DeclContext *CurrDeclContext,
                    CodeCompletionContext *CompletionContext = nullptr);
 
+  void setSolutionSpecificVarTypes(
+      llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes) {
+    this->SolutionSpecificVarTypes = SolutionSpecificVarTypes;
+  }
+
   void setHaveDot(SourceLoc DotLoc) {
     HaveDot = true;
     this->DotLoc = DotLoc;
@@ -226,6 +237,10 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
 
   void setIdealExpectedType(Type Ty) { expectedTypeContext.setIdealType(Ty); }
 
+  void setCanCurrDeclContextHandleAsync(bool CanCurrDeclContextHandleAsync) {
+    this->CanCurrDeclContextHandleAsync = CanCurrDeclContextHandleAsync;
+  }
+
   const ExpectedTypeContext *getExpectedTypeContext() const {
     return &expectedTypeContext;
   }

include/swift/IDE/ExprCompletion.h

@@ -0,0 +1,66 @@
+//===--- ExprCompletion.h -------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 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_EXPRCOMPLETION_H
+#define SWIFT_IDE_EXPRCOMPLETION_H
+
+#include "swift/IDE/CodeCompletionConsumer.h"
+#include "swift/IDE/CodeCompletionContext.h"
+#include "swift/Sema/CodeCompletionTypeChecking.h"
+
+namespace swift {
+namespace ide {
+
+class ExprTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
+public:
+  struct Result {
+    /// The contextual type that the code completion expression should produce.
+    Type ExpectedType;
+
+    /// If the code completion expression is an implicit return in a
+    /// single-expression closure.
+    bool IsImplicitSingleExpressionReturn;
+
+    /// Whether the surrounding context is async and thus calling async
+    /// functions is supported.
+    bool IsInAsyncContext;
+
+    /// Types of variables that were determined in the solution that produced
+    /// this result. This in particular includes parameters of closures that
+    /// were type-checked with the code completion expression.
+    llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+  };
+
+private:
+  CodeCompletionExpr *CompletionExpr;
+  DeclContext *DC;
+
+  SmallVector<Result, 4> Results;
+
+public:
+  /// \param DC The decl context in which the \p CompletionExpr occurs.
+  ExprTypeCheckCompletionCallback(CodeCompletionExpr *CompletionExpr,
+                                  DeclContext *DC)
+      : CompletionExpr(CompletionExpr), DC(DC) {}
+
+  void sawSolution(const constraints::Solution &solution) override;
+
+  /// \param CCLoc The location of the code completion token.
+  void deliverResults(SourceLoc CCLoc,
+                      ide::CodeCompletionContext &CompletionCtx,
+                      CodeCompletionConsumer &Consumer);
+};
+
+} // end namespace ide
+} // end namespace swift
+
+#endif // SWIFT_IDE_EXPRCOMPLETION_H

include/swift/Sema/ConstraintSystem.h

@@ -978,6 +978,7 @@ class SolutionApplicationTargetsKey {
     tombstone,
     stmtCondElement,
     expr,
+    closure,
     stmt,
     pattern,
     patternBindingEntry,
@@ -1020,6 +1021,11 @@ class SolutionApplicationTargetsKey {
     storage.expr = expr;
   }
 
+  SolutionApplicationTargetsKey(const ClosureExpr *closure) {
+    kind = Kind::closure;
+    storage.expr = closure;
+  }
+
   SolutionApplicationTargetsKey(const Stmt *stmt) {
     kind = Kind::stmt;
     storage.stmt = stmt;
@@ -1056,6 +1062,7 @@ class SolutionApplicationTargetsKey {
       return lhs.storage.stmtCondElement == rhs.storage.stmtCondElement;
 
     case Kind::expr:
+    case Kind::closure:
       return lhs.storage.expr == rhs.storage.expr;
 
     case Kind::stmt:
@@ -1096,6 +1103,7 @@ class SolutionApplicationTargetsKey {
           DenseMapInfo<void *>::getHashValue(storage.stmtCondElement));
 
     case Kind::expr:
+    case Kind::closure:
       return hash_combine(
           DenseMapInfo<unsigned>::getHashValue(static_cast<unsigned>(kind)),
           DenseMapInfo<void *>::getHashValue(storage.expr));
@@ -1623,6 +1631,7 @@ class SolutionApplicationTarget {
 public:
   enum class Kind {
     expression,
+    closure,
     function,
     stmtCondition,
     caseLabelItem,
@@ -1684,6 +1693,14 @@ class SolutionApplicationTarget {
       };
     } expression;
 
+    struct {
+      /// The closure expression being type-checked.
+      ClosureExpr *closure;
+
+      /// The type to which the expression should be converted.
+      Type convertType;
+    } closure;
+
     struct {
       AnyFunctionRef function;
       BraceStmt *body;
@@ -1735,6 +1752,12 @@ class SolutionApplicationTarget {
     setPattern(pattern);
   }
 
+  SolutionApplicationTarget(ClosureExpr *closure, Type convertType) {
+    kind = Kind::closure;
+    this->closure.closure = closure;
+    this->closure.convertType = convertType;
+  }
+
   SolutionApplicationTarget(AnyFunctionRef fn)
       : SolutionApplicationTarget(fn, fn.getBody()) { }
 
@@ -1831,6 +1854,7 @@ class SolutionApplicationTarget {
     case Kind::expression:
       return expression.expression;
 
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -1846,6 +1870,9 @@ class SolutionApplicationTarget {
     case Kind::expression:
       return expression.dc;
 
+    case Kind::closure:
+      return closure.closure;
+
     case Kind::function:
       return function.function.getAsDeclContext();
 
@@ -1934,6 +1961,11 @@ class SolutionApplicationTarget {
     return cast<ExprPattern>(expression.pattern);
   }
 
+  Type getClosureContextualType() const {
+    assert(kind == Kind::closure);
+    return closure.convertType;
+  }
+
   /// For a pattern initialization target, retrieve the contextual pattern.
   ContextualPattern getContextualPattern() const;
 
@@ -2062,6 +2094,7 @@ class SolutionApplicationTarget {
   Optional<AnyFunctionRef> getAsFunction() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
     case Kind::patternBinding:
@@ -2077,6 +2110,7 @@ class SolutionApplicationTarget {
   Optional<StmtCondition> getAsStmtCondition() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::caseLabelItem:
     case Kind::patternBinding:
@@ -2092,6 +2126,7 @@ class SolutionApplicationTarget {
   Optional<CaseLabelItem *> getAsCaseLabelItem() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::patternBinding:
@@ -2107,6 +2142,7 @@ class SolutionApplicationTarget {
   PatternBindingDecl *getAsPatternBinding() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -2122,6 +2158,7 @@ class SolutionApplicationTarget {
   VarDecl *getAsUninitializedWrappedVar() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -2137,6 +2174,7 @@ class SolutionApplicationTarget {
   Pattern *getAsUninitializedVar() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -2152,6 +2190,7 @@ class SolutionApplicationTarget {
   Type getTypeOfUninitializedVar() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -2167,6 +2206,7 @@ class SolutionApplicationTarget {
   PatternBindingDecl *getPatternBindingOfUninitializedVar() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -2182,6 +2222,7 @@ class SolutionApplicationTarget {
   unsigned getIndexOfUninitializedVar() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::closure:
     case Kind::function:
     case Kind::stmtCondition:
     case Kind::caseLabelItem:
@@ -2210,6 +2251,9 @@ class SolutionApplicationTarget {
     case Kind::expression:
       return expression.expression->getSourceRange();
 
+    case Kind::closure:
+      return closure.closure->getSourceRange();
+
     case Kind::function:
       return function.body->getSourceRange();
 
@@ -2240,6 +2284,9 @@ class SolutionApplicationTarget {
     case Kind::expression:
       return expression.expression->getLoc();
 
+    case Kind::closure:
+      return closure.closure->getLoc();
+
     case Kind::function:
       return function.function.getLoc();
 

lib/IDE/ArgumentCompletion.cpp

@@ -106,7 +106,9 @@ void ArgumentTypeCheckCompletionCallback::sawSolution(const Solution &S) {
   }
 
   if (!ParentCall || ParentCall == CompletionExpr) {
-    assert(false && "no containing call?");
+    // We might not have a call that contains the code completion expression if
+    // we type-checked the fallback code completion expression that only
+    // contains the code completion token, but not the surrounding call.
     return;
   }
 

lib/IDE/CMakeLists.txt

@@ -20,6 +20,7 @@ add_swift_host_library(swiftIDE STATIC
   ConformingMethodList.cpp
   DependencyChecking.cpp
   DotExprCompletion.cpp
+  ExprCompletion.cpp
   ExprContextAnalysis.cpp
   Formatting.cpp
   FuzzyStringMatcher.cpp

lib/IDE/CodeCompletion.cpp

@@ -40,6 +40,7 @@
 #include "swift/IDE/CompletionLookup.h"
 #include "swift/IDE/CompletionOverrideLookup.h"
 #include "swift/IDE/DotExprCompletion.h"
+#include "swift/IDE/ExprCompletion.h"
 #include "swift/IDE/KeyPathCompletion.h"
 #include "swift/IDE/UnresolvedMemberCompletion.h"
 #include "swift/IDE/Utils.h"
@@ -1392,6 +1393,26 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
                           CurDeclContext, CompletionContext, Consumer);
     return true;
   }
+  case CompletionKind::StmtOrExpr: {
+    assert(CodeCompleteTokenExpr);
+    assert(CurDeclContext);
+
+    ExprTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr,
+                                           CurDeclContext);
+    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
+        Context.CompletionCallback, &Lookup);
+    typeCheckContextAt(CurDeclContext, CompletionLoc);
+
+    if (!Lookup.gotCallback()) {
+      Lookup.fallbackTypeCheck(CurDeclContext);
+    }
+
+    addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
+
+    SourceLoc CCLoc = P.Context.SourceMgr.getCodeCompletionLoc();
+    Lookup.deliverResults(CCLoc, CompletionContext, Consumer);
+    return true;
+  }
   default:
     return false;
   }
@@ -1514,10 +1535,10 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::UnresolvedMember:
   case CompletionKind::KeyPathExprSwift:
   case CompletionKind::CallArg:
+  case CompletionKind::StmtOrExpr:
     llvm_unreachable("should be already handled");
     return;
 
-  case CompletionKind::StmtOrExpr:
   case CompletionKind::ForEachSequence:
   case CompletionKind::PostfixExprBeginning: {
     ExprContextInfo ContextInfo(CurDeclContext, CodeCompleteTokenExpr);

lib/IDE/CompletionLookup.cpp

@@ -804,8 +804,13 @@ void CompletionLookup::addVarDeclRef(const VarDecl *VD,
   assert(!Name.empty() && "name should not be empty");
 
   Type VarType;
-  if (VD->hasInterfaceType())
+  auto SolutionSpecificType = SolutionSpecificVarTypes.find(VD);
+  if (SolutionSpecificType != SolutionSpecificVarTypes.end()) {
+    assert(!VarType && "Type recorded in the AST and is also solution-specific?");
+    VarType = SolutionSpecificType->second;
+  } else if (VD->hasInterfaceType()) {
     VarType = getTypeOfMember(VD, dynamicLookupInfo);
+  }
 
   Optional<ContextualNotRecommendedReason> NotRecommended;
   // "not recommended" in its own getter.