[CodeCompletion][Sema][Parse] Migrate unresolved member completion to the solver-based completion implementation

Following on from updating regular member completion, this hooks up unresolved
member completion (i.e. .<complete here>) to the typeCheckForCodeCompletion API
to generate completions from all solutions the constraint solver produces (even
those requiring fixes), rather than relying on a single solution being applied
to the AST (if any). This lets us produce unresolved member completions even
when the contextual type is ambiguous or involves errors.

Whenever typeCheckExpression is called on an expression containing a code
completion expression and a CompletionCallback has been set, each solution
formed is passed to the callback so the type of the completion expression can
be extracted and used to lookup up the members to return.

Author: Nathan Hawes

include/swift/Sema/CodeCompletionTypeChecking.h

@@ -18,6 +18,11 @@
 #ifndef SWIFT_SEMA_CODECOMPLETIONTYPECHECKING_H
 #define SWIFT_SEMA_CODECOMPLETIONTYPECHECKING_H
 
+#include "swift/Basic/LLVM.h"
+#include "swift/AST/Type.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
 namespace swift {
   class Decl;
   class DeclContext;
@@ -32,7 +37,7 @@ namespace swift {
   class TypeCheckCompletionCallback {
   public:
     /// Called for each solution produced while  type-checking an expression
-    /// containing a code completion expression.
+    /// that the code completion expression participates in.
     virtual void sawSolution(const constraints::Solution &solution) = 0;
     virtual ~TypeCheckCompletionCallback() {}
   };
@@ -78,6 +83,39 @@ namespace swift {
 
     void sawSolution(const constraints::Solution &solution) override;
   };
+
+  /// Used to collect and store information needed to perform unresolved member
+  /// completion (\c CompletionKind::UnresolvedMember ) from the solutions
+  /// formed during expression type-checking.
+  class UnresolvedMemberTypeCheckCompletionCallback: public TypeCheckCompletionCallback {
+  public:
+    struct Result {
+      Type ExpectedTy;
+      bool IsSingleExpressionBody;
+    };
+
+  private:
+    CodeCompletionExpr *CompletionExpr;
+    SmallVector<Result, 4> Results;
+    bool GotCallback = false;
+
+  public:
+    UnresolvedMemberTypeCheckCompletionCallback(CodeCompletionExpr *CompletionExpr)
+    : CompletionExpr(CompletionExpr) {}
+
+    ArrayRef<Result> getResults() const { return Results; }
+
+    /// True if at least one solution was passed via the \c sawSolution
+    /// callback.
+    bool gotCallback() const { return GotCallback; }
+
+    /// Typecheck the code completion expression in its outermost expression
+    /// context, calling \c sawSolution for each solution formed.
+    void fallbackTypeCheck(DeclContext *DC);
+
+    void sawSolution(const constraints::Solution &solution) override;
+  };
+
 }
 
 #endif

lib/IDE/CodeCompletion.cpp

@@ -4340,6 +4340,16 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     if (!T->mayHaveMembers())
       return;
 
+    if (auto objT = T->getOptionalObjectType()) {
+      // Add 'nil' keyword with erasing '.' instruction.
+      unsigned bytesToErase = 0;
+      auto &SM = CurrDeclContext->getASTContext().SourceMgr;
+      if (DotLoc.isValid())
+        bytesToErase = SM.getByteDistance(DotLoc, SM.getCodeCompletionLoc());
+      addKeyword("nil", T, SemanticContextKind::None,
+                 CodeCompletionKeywordKind::kw_nil, bytesToErase);
+    }
+
     // We can only say .foo where foo is a static member of the contextual
     // type and has the same type (or if the member is a function, then the
     // same result type) as the contextual type.
@@ -4378,14 +4388,6 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
         objT = objT->lookThroughAllOptionalTypes();
         if (seenTypes.insert(objT->getCanonicalType()).second)
           getUnresolvedMemberCompletions(objT);
-
-        // Add 'nil' keyword with erasing '.' instruction.
-        unsigned bytesToErase = 0;
-        auto &SM = CurrDeclContext->getASTContext().SourceMgr;
-        if (DotLoc.isValid())
-          bytesToErase = SM.getByteDistance(DotLoc, SM.getCodeCompletionLoc());
-        addKeyword("nil", T, SemanticContextKind::None,
-                   CodeCompletionKeywordKind::kw_nil, bytesToErase);
       }
       getUnresolvedMemberCompletions(T);
     }
@@ -6083,6 +6085,45 @@ static void deliverCompletionResults(CodeCompletionContext &CompletionContext,
                                    DCForModules);
 }
 
+void deliverUnresolvedMemberResults(
+    ArrayRef<UnresolvedMemberTypeCheckCompletionCallback::Result> Results,
+    DeclContext *DC, SourceLoc DotLoc,
+    ide::CodeCompletionContext &CompletionCtx,
+    CodeCompletionConsumer &Consumer) {
+  ASTContext &Ctx = DC->getASTContext();
+  CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
+                          &CompletionCtx);
+
+  assert(DotLoc.isValid());
+  Lookup.setHaveDot(DotLoc);
+  Lookup.shouldCheckForDuplicates(Results.size() > 1);
+
+  // Get the canonical versions of the top-level types
+  SmallPtrSet<CanType, 4> originalTypes;
+  for (auto &Result: Results)
+    originalTypes.insert(Result.ExpectedTy->getCanonicalType());
+
+  for (auto &Result: Results) {
+    Lookup.setExpectedTypes({Result.ExpectedTy},
+                            Result.IsSingleExpressionBody,
+                            /*expectsNonVoid*/true);
+    Lookup.setIdealExpectedType(Result.ExpectedTy);
+
+    // For optional types, also get members of the unwrapped type if it's not
+    // already equivalent to one of the top-level types. Handling it via the top
+    // level type and not here ensures we give the correct type relation
+    // (identical, rather than convertible).
+    if (Result.ExpectedTy->getOptionalObjectType()) {
+      Type Unwrapped = Result.ExpectedTy->lookThroughAllOptionalTypes();
+      if (originalTypes.insert(Unwrapped->getCanonicalType()).second)
+        Lookup.getUnresolvedMemberCompletions(Unwrapped);
+    }
+    Lookup.getUnresolvedMemberCompletions(Result.ExpectedTy);
+  }
+  SourceFile *SF = DC->getParentSourceFile();
+  deliverCompletionResults(CompletionCtx, Lookup, *SF, Consumer);
+}
+
 void deliverDotExprResults(
     ArrayRef<DotExprTypeCheckCompletionCallback::Result> Results,
     Expr *BaseExpr, DeclContext *DC, SourceLoc DotLoc, bool IsInSelector,
@@ -6158,6 +6199,23 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
                           Consumer);
     return true;
   }
+  case CompletionKind::UnresolvedMember: {
+    assert(CodeCompleteTokenExpr);
+    assert(CurDeclContext);
+
+    UnresolvedMemberTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr);
+    llvm::SaveAndRestore<TypeCheckCompletionCallback*>
+      CompletionCollector(Context.CompletionCallback, &Lookup);
+    typeCheckContextAt(CurDeclContext, CompletionLoc);
+
+    if (!Lookup.gotCallback())
+      Lookup.fallbackTypeCheck(CurDeclContext);
+
+    addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
+    deliverUnresolvedMemberResults(Lookup.getResults(), CurDeclContext, DotLoc,
+                                   CompletionContext, Consumer);
+    return true;
+  }
   default:
     return false;
   }
@@ -6277,6 +6335,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   switch (Kind) {
   case CompletionKind::None:
   case CompletionKind::DotExpr:
+  case CompletionKind::UnresolvedMember:
     llvm_unreachable("should be already handled");
     return;
 
@@ -6478,15 +6537,6 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       Lookup.addImportModuleNames();
     break;
   }
-  case CompletionKind::UnresolvedMember: {
-    Lookup.setHaveDot(DotLoc);
-    ExprContextInfo ContextInfo(CurDeclContext, CodeCompleteTokenExpr);
-    Lookup.setExpectedTypes(ContextInfo.getPossibleTypes(),
-                            ContextInfo.isSingleExpressionBody());
-    Lookup.setIdealExpectedType(CodeCompleteTokenExpr->getType());
-    Lookup.getUnresolvedMemberCompletions(ContextInfo.getPossibleTypes());
-    break;
-  }
   case CompletionKind::CallArg: {
     ExprContextInfo ContextInfo(CurDeclContext, CodeCompleteTokenExpr);
 

lib/Parse/ParseExpr.cpp

@@ -3581,26 +3581,22 @@ Parser::parseExprCollectionElement(Optional<bool> &isDictionary) {
     return Element;
 
   // Parse the ':'.
-  if (!consumeIf(tok::colon)) {
-    if (Element.hasCodeCompletion()) {
-      // Return the completion expression itself so we can analyze the type
-      // later.
-      return Element;
+  ParserResult<Expr> Value;
+  if (consumeIf(tok::colon)) {
+    // Parse the value.
+    Value = parseExpr(diag::expected_value_in_dictionary_literal);
+    if (Value.isNull()) {
+      if (!Element.hasCodeCompletion()) {
+        Value = makeParserResult(Value, new (Context) ErrorExpr(PreviousLoc));
+      } else {
+        Value = makeParserResult(Value,
+                                 new (Context) CodeCompletionExpr(PreviousLoc));
+      }
     }
+  } else {
     diagnose(Tok, diag::expected_colon_in_dictionary_literal);
-    return ParserStatus(Element) | makeParserError();
-  }
-
-  // Parse the value.
-  auto Value = parseExpr(diag::expected_value_in_dictionary_literal);
-
-  if (Value.isNull()) {
-    if (!Element.hasCodeCompletion()) {
-      Value = makeParserResult(Value, new (Context) ErrorExpr(PreviousLoc));
-    } else {
-      Value = makeParserResult(Value,
-                               new (Context) CodeCompletionExpr(PreviousLoc));
-    }
+    Value = makeParserResult(makeParserError(),
+                             new (Context) ErrorExpr(SourceRange()));
   }
 
   // Make a tuple of Key Value pair.

lib/Sema/CSBindings.cpp

@@ -1384,7 +1384,8 @@ TypeVariableBinding::fixForHole(ConstraintSystem &cs) const {
     // under-constrained due to e.g. lack of expressions on the
     // right-hand side of the token, which are required for a
     // regular type-check.
-    if (dstLocator->directlyAt<CodeCompletionExpr>())
+    if (dstLocator->directlyAt<CodeCompletionExpr>() ||
+        srcLocator->directlyAt<CodeCompletionExpr>())
       return None;
   }
 

lib/Sema/CSGen.cpp

@@ -1020,7 +1020,7 @@ namespace {
       // independent of the wider expression containing the ErrorExpr, so
       // there's no point attempting to produce a solution for it.
       SourceRange range = E->getSourceRange();
-      if (range.isInvalid() ||
+      if (range.isValid() &&
           CS.getASTContext().SourceMgr.rangeContainsCodeCompletionLoc(range))
         return nullptr;
 

lib/Sema/CSRanking.cpp

@@ -32,74 +32,83 @@ using namespace constraints;
 #define DEBUG_TYPE "Constraint solver overall"
 STATISTIC(NumDiscardedSolutions, "Number of solutions discarded");
 
-void ConstraintSystem::increaseScore(ScoreKind kind, unsigned value) {
-  unsigned index = static_cast<unsigned>(kind);
-  CurrentScore.Data[index] += value;
+static StringRef getScoreKindName(ScoreKind kind) {
+  switch (kind) {
+  case SK_Hole:
+    return "hole in the constraint system";
 
-  if (isDebugMode() && value > 0) {
-    if (solverState)
-      llvm::errs().indent(solverState->depth * 2);
-    llvm::errs() << "(increasing score due to ";
-    switch (kind) {
-    case SK_Hole:
-      llvm::errs() << "hole in the constraint system";
-      break;
+  case SK_Unavailable:
+    return "use of an unavailable declaration";
 
-    case SK_Unavailable:
-      llvm::errs() << "use of an unavailable declaration";
-      break;
+  case SK_AsyncSyncMismatch:
+    return "async/synchronous mismatch";
 
-    case SK_AsyncSyncMismatch:
-      llvm::errs() << "async/synchronous mismatch";
-      break;
+  case SK_ForwardTrailingClosure:
+    return "forward scan when matching a trailing closure";
 
-    case SK_ForwardTrailingClosure:
-      llvm::errs() << "forward scan when matching a trailing closure";
-      break;
+  case SK_Fix:
+    return "attempting to fix the source";
 
-    case SK_Fix:
-      llvm::errs() << "attempting to fix the source";
-      break;
+  case SK_DisfavoredOverload:
+    return "disfavored overload";
 
-    case SK_DisfavoredOverload:
-      llvm::errs() << "disfavored overload";
-      break;
+  case SK_ForceUnchecked:
+    return "force of an implicitly unwrapped optional";
 
-    case SK_ForceUnchecked:
-      llvm::errs() << "force of an implicitly unwrapped optional";
-      break;
+  case SK_UserConversion:
+    return "user conversion";
 
-    case SK_UserConversion:
-      llvm::errs() << "user conversion";
-      break;
+  case SK_FunctionConversion:
+    return "function conversion";
 
-    case SK_FunctionConversion:
-      llvm::errs() << "function conversion";
-      break;
+  case SK_NonDefaultLiteral:
+    return "non-default literal";
+
+  case SK_CollectionUpcastConversion:
+    return "collection upcast conversion";
+
+  case SK_ValueToOptional:
+    return "value to optional";
+
+  case SK_EmptyExistentialConversion:
+    return "empty-existential conversion";
+
+  case SK_KeyPathSubscript:
+    return "key path subscript";
 
+  case SK_ValueToPointerConversion:
+    return "value-to-pointer conversion";
+  }
+}
+
+void ConstraintSystem::increaseScore(ScoreKind kind, unsigned value) {
+  if (isForCodeCompletion()) {
+    switch (kind) {
     case SK_NonDefaultLiteral:
-      llvm::errs() << "non-default literal";
-      break;
-        
-    case SK_CollectionUpcastConversion:
-      llvm::errs() << "collection upcast conversion";
-      break;
-        
-    case SK_ValueToOptional:
-      llvm::errs() << "value to optional";
-      break;
-    case SK_EmptyExistentialConversion:
-      llvm::errs() << "empty-existential conversion";
-      break;
-    case SK_KeyPathSubscript:
-      llvm::errs() << "key path subscript";
-      break;
-    case SK_ValueToPointerConversion:
-      llvm::errs() << "value-to-pointer conversion";
+      // Don't increase score for non-default literals in expressions involving
+      // a code completion. In the below example, members of EnumA and EnumB
+      // should be ranked equally:
+      //   func overloaded(_ x: Float, _ y: EnumA) {}
+      //   func overloaded(_ x: Int, _ y: EnumB) {}
+      //   func overloaded(_ x: Float) -> EnumA {}
+      //   func overloaded(_ x: Int) -> EnumB {}
+      //
+      //   overloaded(1, .<complete>) {}
+      //   overloaded(1).<complete>
+      return;
+    default:
       break;
     }
-    llvm::errs() << ")\n";
   }
+
+  if (isDebugMode() && value > 0) {
+    if (solverState)
+      llvm::errs().indent(solverState->depth * 2);
+    llvm::errs() << "(increasing score due to " << getScoreKindName(kind) << ")\n";
+  }
+
+  unsigned index = static_cast<unsigned>(kind);
+  CurrentScore.Data[index] += value;
 }
 
 bool ConstraintSystem::worseThanBestSolution() const {