[CodeCompletion] Migrate labeled trailing closure completions to solver-based

rdar://113472967

Author: ahoppen

include/swift/IDE/ArgumentCompletion.h

@@ -25,31 +25,48 @@ 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.
     ValueDecl *FuncD;
+
     /// The type of the function being called.
     AnyFunctionType *FuncTy;
+
     /// The index of the argument containing the completion location
     unsigned ArgIdx;
+
     /// The index of the parameter corresponding to the completion argument.
     llvm::Optional<unsigned> ParamIdx;
+
     /// The indices of all params that were bound to non-synthesized
     /// arguments. Used so we don't suggest them even when the args are out of
     /// order.
     std::set<unsigned> ClaimedParamIndices;
+
     /// True if the completion is a noninitial term in a variadic argument.
     bool IsNoninitialVariadic;
+
     /// The base type of the call/subscript (null for free functions).
     Type BaseType;
+
     /// True if an argument label precedes the completion location.
     bool HasLabel;
+
+    /// The argument index of the first trailing closure.
+    ///
+    /// \c None if the call doesn't have a trailing closure.
+    llvm::Optional<unsigned> FirstTrailingClosureIndex;
+
     /// Whether the surrounding context is async and thus calling async
     /// functions is supported.
     bool IsInAsyncContext;
+
     /// A bitfield to mark whether the parameter at a given index is optional.
     /// Parameters can be optional if they have a default argument or belong to
     /// a parameter pack.
@@ -93,7 +110,11 @@ class ArgumentTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
   /// function signature instead of suggesting individual labels. Used when
   /// completing after the opening '(' of a function call \param Loc The
   /// location of the code completion token
-  void collectResults(bool IncludeSignature, SourceLoc Loc, DeclContext *DC,
+  /// \param IsLabeledTrailingClosure Whether we are completing the label of a
+  /// labeled trailing closure, ie. if the code completion location is outside
+  /// the call after the first trailing closure of the call.
+  void collectResults(bool IncludeSignature, bool IsLabeledTrailingClosure,
+                      SourceLoc Loc, DeclContext *DC,
                       CodeCompletionContext &CompletionCtx);
 };
 

lib/IDE/ArgumentCompletion.cpp

@@ -50,6 +50,24 @@ bool ArgumentTypeCheckCompletionCallback::addPossibleParams(
     const AnyFunctionType::Param *TypeParam = &ParamsToPass[Idx];
     bool Required = !Res.DeclParamIsOptional[Idx];
 
+    if (Res.FirstTrailingClosureIndex && Idx > *Res.FirstTrailingClosureIndex &&
+        !TypeParam->getPlainType()
+             ->lookThroughAllOptionalTypes()
+             ->is<AnyFunctionType>()) {
+      // We are completing an argument after the first trailing closure, i.e.
+      // a multitple trailing closure label but the parameter is not a function
+      // type. Since we only allow labeled trailing closures after the first
+      // trailing closure, we cannot pass an argument for this parameter.
+      // If the parameter is required, stop here since we cannot pass an argument
+      // for the parameter. If it's optional, keep looking for more trailing
+      // closures that can be passed.
+      if (Required) {
+        break;
+      } else {
+        continue;
+      }
+    }
+
     if (TypeParam->hasLabel() && !(IsCompletion && Res.IsNoninitialVariadic)) {
       // Suggest parameter label if parameter has label, we are completing in it
       // and it is not a variadic parameter that already has arguments
@@ -198,9 +216,11 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   }
 
   bool HasLabel = false;
+  llvm::Optional<unsigned> FirstTrailingClosureIndex = llvm::None;
   if (auto PE = CS.getParentExpr(CompletionExpr)) {
     if (auto Args = PE->getArgs()) {
       HasLabel = !Args->getLabel(ArgIdx).empty();
+      FirstTrailingClosureIndex = Args->getFirstTrailingClosureIndex();
     }
   }
 
@@ -258,11 +278,11 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
     }
   }
 
-  Results.push_back({ExpectedTy, ExpectedCallType,
-                     isa<SubscriptExpr>(ParentCall), Info.getValue(), FuncTy,
-                     ArgIdx, ParamIdx, std::move(ClaimedParams),
-                     IsNoninitialVariadic, Info.BaseTy, HasLabel, IsAsync,
-                     DeclParamIsOptional, SolutionSpecificVarTypes});
+  Results.push_back(
+      {ExpectedTy, ExpectedCallType, isa<SubscriptExpr>(ParentCall),
+       Info.getValue(), FuncTy, ArgIdx, ParamIdx, std::move(ClaimedParams),
+       IsNoninitialVariadic, Info.BaseTy, HasLabel, FirstTrailingClosureIndex,
+       IsAsync, DeclParamIsOptional, SolutionSpecificVarTypes});
 }
 
 void ArgumentTypeCheckCompletionCallback::computeShadowedDecls(
@@ -296,8 +316,8 @@ void ArgumentTypeCheckCompletionCallback::computeShadowedDecls(
 }
 
 void ArgumentTypeCheckCompletionCallback::collectResults(
-    bool IncludeSignature, SourceLoc Loc, DeclContext *DC,
-    ide::CodeCompletionContext &CompletionCtx) {
+    bool IncludeSignature, bool IsLabeledTrailingClosure, SourceLoc Loc,
+    DeclContext *DC, ide::CodeCompletionContext &CompletionCtx) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
@@ -374,7 +394,7 @@ void ArgumentTypeCheckCompletionCallback::collectResults(
       shouldPerformGlobalCompletion |=
           addPossibleParams(Ret, Params, ExpectedTypes);
     }
-    Lookup.addCallArgumentCompletionResults(Params);
+    Lookup.addCallArgumentCompletionResults(Params, IsLabeledTrailingClosure);
   }
 
   if (shouldPerformGlobalCompletion) {

lib/IDE/CodeCompletion.cpp

@@ -133,7 +133,6 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
   bool ShouldCompleteCallPatternAfterParen = true;
   bool PreferFunctionReferencesToCalls = false;
   bool AttTargetIsIndependent = false;
-  bool IsAtStartOfLine = false;
   llvm::Optional<DeclKind> AttTargetDK;
   llvm::Optional<StmtKind> ParentStmtKind;
 
@@ -583,6 +582,7 @@ void CodeCompletionCallbacksImpl::completeLabeledTrailingClosure(
   CodeCompleteTokenExpr = E;
   Kind = CompletionKind::LabeledTrailingClosure;
   IsAtStartOfLine = isAtStartOfLine;
+  ShouldCompleteCallPatternAfterParen = false;
 }
 
 void CodeCompletionCallbacksImpl::completeReturnStmt(CodeCompletionExpr *E) {
@@ -1491,6 +1491,43 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
 
     Lookup.collectResults(DotLoc, isInsideObjCSelector(), IncludeOperators,
                           HasSpace, CompletionContext);
+
+    // Check if we are completing after a call that already has a trailing
+    // closure. In that case, also suggest labels for additional trailing
+    // closures.
+    if (auto AE = dyn_cast<ApplyExpr>(ParsedExpr)) {
+      if (AE->getArgs()->hasAnyTrailingClosures()) {
+        ASTContext &Ctx = CurDeclContext->getASTContext();
+
+        // Modify the call that has the code completion expression as an
+        // additional argument, restore the original arguments afterwards.
+        auto OriginalArgs = AE->getArgs();
+        llvm::SmallVector<Argument> ArgsWithCC(OriginalArgs->begin(),
+                                               OriginalArgs->end());
+        auto CC = new (Ctx) CodeCompletionExpr(CodeCompleteTokenExpr->getLoc());
+        ArgsWithCC.emplace_back(SourceLoc(), Identifier(), CC);
+        auto ArgList =
+            ArgumentList::create(Ctx, OriginalArgs->getLParenLoc(), ArgsWithCC,
+                                 OriginalArgs->getRParenLoc(),
+                                 OriginalArgs->getFirstTrailingClosureIndex(),
+                                 OriginalArgs->isImplicit());
+        AE->setArgs(ArgList);
+        SWIFT_DEFER { AE->setArgs(OriginalArgs); };
+
+        // Perform argument label completions on the newly created call.
+        ArgumentTypeCheckCompletionCallback Lookup(CC, CurDeclContext);
+
+        llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
+            Context.CompletionCallback, &Lookup);
+        typeCheckContextAt(
+            TypeCheckASTNodeAtLocContext::node(CurDeclContext, AE),
+            CompletionLoc);
+        Lookup.collectResults(/*IncludeSignature=*/false,
+                              /*IsLabeledTrailingClosure=*/true, CompletionLoc,
+                              CurDeclContext, CompletionContext);
+      }
+    }
+
     Consumer.handleResults(CompletionContext);
     return true;
   }
@@ -1521,14 +1558,19 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     return true;
   }
   case CompletionKind::PostfixExprParen:
-  case CompletionKind::CallArg: {
+  case CompletionKind::CallArg:
+  case CompletionKind::LabeledTrailingClosure: {
+    // FIXME: Delete LabeledTrailingClosure
     assert(CodeCompleteTokenExpr);
     assert(CurDeclContext);
     ArgumentTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr,
                                                CurDeclContext);
     typeCheckWithLookup(Lookup);
 
-    Lookup.collectResults(ShouldCompleteCallPatternAfterParen, CompletionLoc,
+    bool IsLabeledTrailingClosure =
+        (Kind == CompletionKind::LabeledTrailingClosure);
+    Lookup.collectResults(ShouldCompleteCallPatternAfterParen,
+                          IsLabeledTrailingClosure, CompletionLoc,
                           CurDeclContext, CompletionContext);
     Consumer.handleResults(CompletionContext);
     return true;
@@ -1691,12 +1733,6 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
   if (PreferFunctionReferencesToCalls)
     Lookup.setPreferFunctionReferencesToCalls();
 
-  auto DoPostfixExprBeginning = [&] (){
-    SourceLoc Loc = P.Context.SourceMgr.getIDEInspectionTargetLoc();
-    Lookup.getValueCompletionsInDeclContext(Loc);
-    Lookup.getSelfTypeCompletionInDeclContext(Loc, /*isForDeclResult=*/false);
-  };
-
   switch (Kind) {
   case CompletionKind::None:
   case CompletionKind::DotExpr:

lib/IDE/PostfixCompletion.cpp

@@ -433,12 +433,20 @@ void PostfixCompletionCallback::collectResults(
   UnifiedTypeContext.setPreferNonVoid(true);
   bool UnifiedCanHandleAsync = false;
 
+  // The base types of the result for which we already returned results.
+  // Used so we only return keyword and operator completions once for each base
+  // type.
+  llvm::SmallPtrSet<Type, 2> ProcessedBaseTypes;
+
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
+
   for (auto &Result : Results) {
     Lookup.setCanCurrDeclContextHandleAsync(Result.IsInAsyncContext);
     Lookup.setClosureActorIsolations(Result.ClosureActorIsolations);
     Lookup.setIsStaticMetatype(Result.BaseIsStaticMetaType);
-    Lookup.getPostfixKeywordCompletions(Result.BaseTy, BaseExpr);
+    if (!ProcessedBaseTypes.contains(Result.BaseTy)) {
+      Lookup.getPostfixKeywordCompletions(Result.BaseTy, BaseExpr);
+    }
     Lookup.setExpectedTypes(Result.ExpectedTypes,
                             Result.IsImplicitSingleExpressionReturn,
                             Result.ExpectsNonVoid);
@@ -447,12 +455,19 @@ void PostfixCompletionCallback::collectResults(
     Lookup.getValueExprCompletions(Result.BaseTy, Result.BaseDecl,
                                    Result.IsBaseDeclUnapplied);
 
-    if (IncludeOperators && !Result.BaseIsStaticMetaType) {
+    // `==`, `<=` etc can be used on `Void` because `Void` is just an empty
+    // tuple. But that doesn’t really make sense so we shouldn't be suggesting
+    // any operators based on `Void`.
+    if (IncludeOperators && !Result.BaseIsStaticMetaType &&
+        !Result.BaseTy->isVoid() &&
+        !ProcessedBaseTypes.contains(Result.BaseTy)) {
       addOperatorResults(Result.BaseTy, Operators, DC, Lookup);
     }
 
     UnifiedTypeContext.merge(*Lookup.getExpectedTypeContext());
     UnifiedCanHandleAsync |= Result.IsInAsyncContext;
+
+    ProcessedBaseTypes.insert(Result.BaseTy);
   }
 
   collectCompletionResults(CompletionCtx, Lookup, DC, UnifiedTypeContext,

lib/Parse/ParseExpr.cpp

@@ -3336,38 +3336,8 @@ Parser::parseTrailingClosures(bool isExprBasic, SourceRange calleeRange,
   // Parse labeled trailing closures.
   while (true) {
     if (!isStartOfLabelledTrailingClosure(*this)) {
-      if (!Tok.is(tok::code_complete))
-        break;
-
-      // FIXME: Additional trailing closure completion on newline positions.
-      //   let foo = SomeThing {
-      //     ...
-      //   }
-      //   <HERE>
-      // This was previously enabled, but it failed to suggest 'foo' because
-      // the token was considered a part of the initializer.
-      if (Tok.isAtStartOfLine())
-        break;
-
-      // If the current completion mode doesn't support trailing closure
-      // completion, leave the token here and let "postfix completion" to
-      // handle it.
-      if (CodeCompletionCallbacks &&
-          !CodeCompletionCallbacks->canPerformCompleteLabeledTrailingClosure())
-        break;
-
-      // foo() {} <token>
-      auto CCExpr = new (Context) CodeCompletionExpr(Tok.getLoc());
-      if (CodeCompletionCallbacks) {
-        CodeCompletionCallbacks->completeLabeledTrailingClosure(
-            CCExpr, Tok.isAtStartOfLine());
-      }
-      consumeToken(tok::code_complete);
-      result.setHasCodeCompletionAndIsError();
-      closures.emplace_back(SourceLoc(), Identifier(), CCExpr);
-      continue;
+      break;
     }
-
     Identifier label;
     auto labelLoc = consumeArgumentLabel(label, /*diagnoseDollarPrefix=*/false);
     consumeToken(tok::colon);