[Completion] Complete ownership specifiers in parameters

Complete ownership specifiers such as `consuming`,
`borrowing`, and `inout` in parameter type
position. While here, also complete `isolated`.

rdar://127261573

Author: hamishknight

include/swift/IDE/CodeCompletionResult.h

@@ -197,6 +197,7 @@ enum class CompletionKind : uint8_t {
   PostfixExpr,
   KeyPathExprObjC,
   KeyPathExprSwift,
+  TypePossibleFunctionParamBeginning,
   TypeDeclResultBeginning,
   TypeBeginning,
   TypeSimpleOrComposition,

include/swift/Parse/IDEInspectionCallbacks.h

@@ -178,6 +178,12 @@ class CodeCompletionCallbacks {
   /// #keyPath argument have been parsed yet.
   virtual void completeExprKeyPath(KeyPathExpr *KPE, SourceLoc DotLoc) {};
 
+  /// Complete the beginning of the type for a parameter of a
+  /// func/subscript/closure, or the type for a parameter in a function type.
+  /// For the latter, we cannot know for sure whether the user is trying to
+  /// write a function type, so will complete for e.g `let x: (#^COMPLETE^#`.
+  virtual void completeTypePossibleFunctionParamBeginning() {}
+
   /// Complete the beginning of the type of result of func/var/let/subscript.
   virtual void completeTypeDeclResultBeginning() {};
 

lib/IDE/CodeCompletion.cpp

@@ -259,6 +259,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
   void completePostfixExpr(CodeCompletionExpr *E, bool hasSpace) override;
   void completeExprKeyPath(KeyPathExpr *KPE, SourceLoc DotLoc) override;
 
+  void completeTypePossibleFunctionParamBeginning() override;
   void completeTypeDeclResultBeginning() override;
   void completeTypeBeginning() override;
   void completeTypeSimpleOrComposition() override;
@@ -430,6 +431,11 @@ void CodeCompletionCallbacksImpl::completePoundAvailablePlatform() {
   CurDeclContext = P.CurDeclContext;
 }
 
+void CodeCompletionCallbacksImpl::completeTypePossibleFunctionParamBeginning() {
+  Kind = CompletionKind::TypePossibleFunctionParamBeginning;
+  CurDeclContext = P.CurDeclContext;
+}
+
 void CodeCompletionCallbacksImpl::completeTypeDeclResultBeginning() {
   Kind = CompletionKind::TypeDeclResultBeginning;
   CurDeclContext = P.CurDeclContext;
@@ -1058,6 +1064,12 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
     }
     break;
 
+  case CompletionKind::TypePossibleFunctionParamBeginning:
+    addKeyword(Sink, "inout", CodeCompletionKeywordKind::kw_inout);
+    addKeyword(Sink, "borrowing", CodeCompletionKeywordKind::None);
+    addKeyword(Sink, "consuming", CodeCompletionKeywordKind::None);
+    addKeyword(Sink, "isolated", CodeCompletionKeywordKind::None);
+    LLVM_FALLTHROUGH;
   case CompletionKind::TypeBeginning:
     addKeyword(Sink, "repeat", CodeCompletionKeywordKind::None);
     LLVM_FALLTHROUGH;
@@ -1262,6 +1274,7 @@ void swift::ide::postProcessCompletionResults(
     // names at non-type name position are "rare".
     if (result->getKind() == CodeCompletionResultKind::Declaration &&
         result->getAssociatedDeclKind() == CodeCompletionDeclKind::Protocol &&
+        Kind != CompletionKind::TypePossibleFunctionParamBeginning &&
         Kind != CompletionKind::TypeBeginning &&
         Kind != CompletionKind::TypeSimpleOrComposition &&
         Kind != CompletionKind::TypeSimpleBeginning &&
@@ -1755,6 +1768,7 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
     break;
   }
 
+  case CompletionKind::TypePossibleFunctionParamBeginning:
   case CompletionKind::TypeDeclResultBeginning:
   case CompletionKind::TypeBeginning:
   case CompletionKind::TypeSimpleOrComposition:

lib/Parse/ParsePattern.cpp

@@ -355,6 +355,22 @@ Parser::parseParameterClause(SourceLoc &leftParenLoc,
         .fixItReplace(Tok.getLoc(), "`" + Tok.getText().str() + "`");
     }
 
+    auto parseParamType = [&]() -> ParserResult<TypeRepr> {
+      // Currently none of the parameter type completions are relevant for
+      // enum cases, so don't include them. We'll complete for a regular type
+      // beginning instead.
+      if (Tok.is(tok::code_complete) &&
+          paramContext != ParameterContextKind::EnumElement) {
+        if (CodeCompletionCallbacks)
+          CodeCompletionCallbacks->completeTypePossibleFunctionParamBeginning();
+
+        auto CCLoc = consumeToken(tok::code_complete);
+        auto *ET = ErrorTypeRepr::create(Context, CCLoc);
+        return makeParserCodeCompletionResult<TypeRepr>(ET);
+      }
+      return parseType(diag::expected_parameter_type);
+    };
+
     if (startsParameterName(isClosure)) {
       // identifier-or-none for the first name
       param.FirstNameLoc = consumeArgumentLabel(param.FirstName,
@@ -402,7 +418,7 @@ Parser::parseParameterClause(SourceLoc &leftParenLoc,
       // (':' type)?
       if (consumeIf(tok::colon)) {
 
-        auto type = parseType(diag::expected_parameter_type);
+        auto type = parseParamType();
         status |= type;
         param.Type = type.getPtrOrNull();
 
@@ -425,7 +441,7 @@ Parser::parseParameterClause(SourceLoc &leftParenLoc,
       }
 
       if (isBareType && paramContext == ParameterContextKind::EnumElement) {
-        auto type = parseType(diag::expected_parameter_type);
+        auto type = parseParamType();
         status |= type;
         param.Type = type.getPtrOrNull();
         param.FirstName = Identifier();
@@ -440,7 +456,7 @@ Parser::parseParameterClause(SourceLoc &leftParenLoc,
         // the user is about to type the parameter label and we shouldn't
         // suggest types.
         SourceLoc typeStartLoc = Tok.getLoc();
-        auto type = parseType(diag::expected_parameter_type);
+        auto type = parseParamType();
         status |= type;
         param.Type = type.getPtrOrNull();
 

lib/Parse/ParseType.cpp

@@ -1202,6 +1202,16 @@ ParserResult<TypeRepr> Parser::parseTypeTupleBody() {
     }
     Backtracking.reset();
 
+    // If we have a code completion token, treat this as a possible parameter
+    // type since the user may be writing this as a function type.
+    if (Tok.is(tok::code_complete)) {
+      if (CodeCompletionCallbacks)
+        CodeCompletionCallbacks->completeTypePossibleFunctionParamBeginning();
+
+      consumeToken();
+      return makeParserCodeCompletionStatus();
+    }
+
     // Parse the type annotation.
     auto type = parseType(diag::expected_type);
     if (type.hasCodeCompletion())

test/IDE/complete_ownership_specifier.swift

@@ -0,0 +1,34 @@
+// RUN: %batch-code-completion
+
+func foo(x: #^FUNC_PARAM?check=SPECIFIER^#) {
+  let fn1 = { (x: #CLOSURE_PARAM?check=SPECIFIER#) in }
+  let fn2: (#^CLOSURE_PARAM_TY_COMPLETE?check=SPECIFIER^#) -> Void
+  let fn3: (#^CLOSURE_PARAM_TY_INCOMPLETE?check=SPECIFIER^#
+}
+
+struct S {
+  init(x: #^INIT_PARAM?check=SPECIFIER^#) {}
+  subscript(x: #SUBSCRIPT_PARAM?check=SPECIFIER#) {}
+}
+
+// Don't complete for enum cases.
+enum E {
+  case e(#^ENUM_CASE_TY?check=SPECIFIER_NOT^#)
+  case f(x: #^ENUM_CASE_LABELED_TY?check=SPECIFIER_NOT^#)
+}
+
+// Don't complete for a regular variable type.
+let x: #^VAR_TY?check=SPECIFIER_NOT^#
+
+// Or for a return type.
+func bar() -> #^RETURN_TY?check=SPECIFIER_NOT^# {}
+
+// SPECIFIER-DAG: Keyword[inout]/None: inout; name=inout
+// SPECIFIER-DAG: Keyword/None: consuming; name=consuming
+// SPECIFIER-DAG: Keyword/None: borrowing; name=borrowing
+// SPECIFIER-DAG: Keyword/None: isolated; name=isolated
+
+// SPECIFIER_NOT-NOT: Keyword[inout]/None: inout
+// SPECIFIER_NOT-NOT: Keyword/None: consuming
+// SPECIFIER_NOT-NOT: Keyword/None: borrowing
+// SPECIFIER_NOT-NOT: Keyword/None: isolated

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp

@@ -175,6 +175,7 @@ bool SourceKit::CodeCompletion::addCustomCompletions(
         addCompletion(custom);
       }
       break;
+    case CompletionKind::TypePossibleFunctionParamBeginning:
     case CompletionKind::TypeDeclResultBeginning:
     case CompletionKind::TypeBeginning:
     case CompletionKind::TypeSimpleOrComposition:
@@ -452,6 +453,7 @@ void CodeCompletionOrganizer::Impl::addCompletionsWithFilter(
   if (filterText.empty()) {
     bool hideLowPriority =
         options.hideLowPriority &&
+        completionKind != CompletionKind::TypePossibleFunctionParamBeginning &&
         completionKind != CompletionKind::TypeDeclResultBeginning &&
         completionKind != CompletionKind::TypeBeginning &&
         completionKind != CompletionKind::TypeSimpleOrComposition &&