Merge pull request #68079 from ahoppen/ahoppen/return-solver-based

[CodeCompletion] Migrate completions after `return` and `yield` to solver-based

Author: ahoppen

lib/IDE/CodeCompletion.cpp

@@ -890,6 +890,25 @@ static void addObserverKeywords(CodeCompletionResultSink &Sink) {
   addKeyword(Sink, "didSet", CodeCompletionKeywordKind::None);
 }
 
+static void addKeywordsAfterReturn(CodeCompletionResultSink &Sink, DeclContext *DC) {
+  // `return nil` is not actually represented as a `ReturnExpr` in the AST but
+  // gets translated to a `FailStmt`. We thus can't produce the 'nil' completion
+  // using the solver-based implementation. Add the result manually.
+  if (auto ctor = dyn_cast_or_null<ConstructorDecl>(DC->getAsDecl())) {
+    if (ctor->isFailable()) {
+      CodeCompletionResultBuilder Builder(Sink, CodeCompletionResultKind::Literal,
+                                          SemanticContextKind::None);
+      Builder.setLiteralKind(CodeCompletionLiteralKind::NilLiteral);
+      Builder.addKeyword("nil");
+      Builder.addTypeAnnotation(ctor->getResultInterfaceType(), {});
+      Builder.setResultTypes(ctor->getResultInterfaceType());
+      ExpectedTypeContext TypeContext;
+      TypeContext.setPossibleTypes({ctor->getResultInterfaceType()});
+      Builder.setTypeContext(TypeContext, DC);
+    }
+  }
+}
+
 void swift::ide::addExprKeywords(CodeCompletionResultSink &Sink,
                                  DeclContext *DC) {
   // Expression is invalid at top-level of non-script files.
@@ -1026,6 +1045,8 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
 
     LLVM_FALLTHROUGH;
   case CompletionKind::ReturnStmtExpr:
+    addKeywordsAfterReturn(Sink, CurDeclContext);
+    LLVM_FALLTHROUGH;
   case CompletionKind::YieldStmtExpr:
   case CompletionKind::ForEachSequence:
     addSuperKeyword(Sink, CurDeclContext);
@@ -1533,7 +1554,9 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
   case CompletionKind::CaseStmtBeginning:
   case CompletionKind::ForEachSequence:
   case CompletionKind::PostfixExprBeginning:
-  case CompletionKind::StmtOrExpr: {
+  case CompletionKind::StmtOrExpr: 
+  case CompletionKind::ReturnStmtExpr:
+  case CompletionKind::YieldStmtExpr: {
     assert(CurDeclContext);
 
     bool AddUnresolvedMemberCompletions =
@@ -1705,6 +1728,8 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
   case CompletionKind::CaseStmtBeginning:
   case CompletionKind::PostfixExprParen:
   case CompletionKind::PostfixExpr:
+  case CompletionKind::ReturnStmtExpr:
+  case CompletionKind::YieldStmtExpr:
     llvm_unreachable("should be already handled");
     return;
 
@@ -1953,29 +1978,6 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
     break;
   }
 
-  case CompletionKind::ReturnStmtExpr : {
-    SourceLoc Loc = P.Context.SourceMgr.getIDEInspectionTargetLoc();
-    SmallVector<Type, 2> possibleReturnTypes;
-    collectPossibleReturnTypesFromContext(CurDeclContext, possibleReturnTypes);
-    Lookup.setExpectedTypes(possibleReturnTypes,
-                            /*isImplicitSingleExpressionReturn*/ false);
-    Lookup.getValueCompletionsInDeclContext(Loc);
-    break;
-  }
-
-  case CompletionKind::YieldStmtExpr: {
-    SourceLoc Loc = P.Context.SourceMgr.getIDEInspectionTargetLoc();
-    if (auto FD = dyn_cast<AccessorDecl>(CurDeclContext)) {
-      if (FD->isCoroutine()) {
-        // TODO: handle multi-value yields.
-        Lookup.setExpectedTypes(FD->getStorage()->getValueInterfaceType(),
-                                /*isImplicitSingleExpressionReturn*/ false);
-      }
-    }
-    Lookup.getValueCompletionsInDeclContext(Loc);
-    break;
-  }
-
   case CompletionKind::AfterPoundDirective: {
     addPoundDirectives(CompletionContext.getResultSink());
 

lib/IDE/PostfixCompletion.cpp

@@ -178,8 +178,6 @@ void PostfixCompletionCallback::sawSolutionImpl(
   auto *Locator = CS.getConstraintLocator(SemanticExpr);
   Type ExpectedTy = getTypeForCompletion(S, CompletionExpr);
   Expr *ParentExpr = CS.getParentExpr(CompletionExpr);
-  if (!ParentExpr && !ExpectedTy)
-    ExpectedTy = CS.getContextualType(CompletionExpr, /*forConstraint=*/false);
 
   auto *CalleeLocator = S.getCalleeLocator(Locator);
   ValueDecl *ReferencedDecl = nullptr;

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -47,6 +47,10 @@ void TypeCheckCompletionCallback::fallbackTypeCheck(DeclContext *DC) {
 
 Type swift::ide::getTypeForCompletion(const constraints::Solution &S,
                                       ASTNode Node) {
+  if (auto ContextualType = S.getContextualType(Node)) {
+    return ContextualType;
+  }
+
   if (!S.hasType(Node)) {
     assert(false && "Expression wasn't type checked?");
     return nullptr;
@@ -60,9 +64,6 @@ Type swift::ide::getTypeForCompletion(const constraints::Solution &S,
     Result = S.getResolvedType(Node);
   }
 
-  if (!Result || Result->is<UnresolvedType>())
-    Result = S.getContextualType(Node);
-
   if (Result && Result->is<UnresolvedType>()) {
     Result = Type();
   }

lib/Parse/ParseStmt.cpp

@@ -880,9 +880,10 @@ ParserResult<Stmt> Parser::parseStmtYield(SourceLoc tryLoc) {
     }
 
     auto expr = parseExpr(diag::expected_expr_yield);
-    if (expr.hasCodeCompletion())
-      return makeParserCodeCompletionResult<Stmt>();
-    if (expr.isParseErrorOrHasCompletion()) {
+    if (expr.hasCodeCompletion() && expr.isNonNull()) {
+      status |= expr;
+      yields.push_back(expr.get());
+    } else if (expr.isParseErrorOrHasCompletion()) {
       auto endLoc = (Tok.getLoc() == beginLoc ? beginLoc : PreviousLoc);
       yields.push_back(
         new (Context) ErrorExpr(SourceRange(beginLoc, endLoc)));

test/IDE/complete_constructor.swift

@@ -312,8 +312,8 @@ class DependentTypeInClosure<Data: DataType> {
 func testDependentTypeInClosure() {
   let _: DependentTypeInClosure = .#^DEPENDENT_IN_CLOSURE_3^#
 // DEPENDENT_IN_CLOSURE_3: Begin completions, 2 items
-// DEPENDENT_IN_CLOSURE_3-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init({#(arg): DataType#}, {#fn: (Data.Content) -> Void##(Data.Content) -> Void#})[#DependentTypeInClosure<DataType>#];
-// DEPENDENT_IN_CLOSURE_3-DAG: Decl[Constructor]/CurrNominal/TypeRelation[Convertible]: init({#arg: DataType#}, {#fn: () -> Data.Content##() -> Data.Content#})[#DependentTypeInClosure<DataType>#];
+// DEPENDENT_IN_CLOSURE_3-DAG: Decl[Constructor]/CurrNominal: init({#(arg): DataType#}, {#fn: (DataType.Content) -> Void##(DataType.Content) -> Void#})[#DependentTypeInClosure<DataType>#];
+// DEPENDENT_IN_CLOSURE_3-DAG: Decl[Constructor]/CurrNominal: init({#arg: DataType#}, {#fn: () -> DataType.Content##() -> DataType.Content#})[#DependentTypeInClosure<DataType>#];
 
   let _ = DependentTypeInClosure(#^DEPENDENT_IN_CLOSURE_1^#)
 // DEPENDENT_IN_CLOSURE_1-DAG: Decl[Constructor]/CurrNominal/Flair[ArgLabels]:      ['(']{#(arg): DataType#}, {#fn: (_) -> Void##(_) -> Void#}[')'][#DependentTypeInClosure<DataType>#];

test/IDE/complete_yield.swift

@@ -0,0 +1,30 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-ide-test -batch-code-completion -source-filename %s -filecheck %raw-FileCheck -completion-output-dir %t
+
+struct YieldVariables {
+  var stringValue: String
+
+  var property: String {
+    _read {
+      yield #^IN_READ^#
+      // IN_READ-DAG: Decl[InstanceVar]/CurrNominal/TypeRelation[Convertible]: stringValue[#String#]; name=stringValue
+    }
+    _modify {
+      var localString = ""
+      yield #^IN_MODIFY^#
+      // IN_MODIFY-DAG: Decl[InstanceVar]/CurrNominal:      stringValue[#String#]; name=stringValue
+      // IN_MODIFY-DAG: Decl[LocalVar]/Local:               localString[#String#]; name=localString
+    }
+  }
+
+  var otherProperty: String {
+    _modify {
+      var localString = ""
+      yield &#^IN_MODIFY_WITH_REF^#
+      // FIXME: We should probably be returning a convertible type relation here.
+      // IN_MODIFY_WITH_REF: Decl[LocalVar]/Local:               localString[#String#]; name=localString
+      // IN_MODIFY_WITH_REF: Decl[InstanceVar]/CurrNominal:      stringValue[#String#]; name=stringValue
+    }
+  }
+}
+