Merge pull request #32283 from rintaro/ide-completion-postfixfallback-rdar64176730

[CodeCompletion] Typecheck without CodeCompletionExpr

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -6172,15 +6172,33 @@ void CodeCompletionCallbacksImpl::doneParsing() {
         if (analyzedExpr->getEndLoc() != CodeCompleteTokenExpr->getLoc())
           break;
 
-        // If the call expression doesn't have a type, infer it from the
-        // possible callee info.
         Type resultTy = analyzedExpr->getType();
-        if (!resultTy) {
-          if (ContextInfo.getPossibleCallees().empty())
-            break;
-          auto calleeInfo = ContextInfo.getPossibleCallees()[0];
-          resultTy = calleeInfo.Type->getResult();
-          analyzedExpr->setType(resultTy);
+        // If the call expression doesn't have a type, fallback to:
+        if (!resultTy || resultTy->is<ErrorType>()) {
+          // 1) Try to type check removing CodeCompletionExpr from the call.
+          Expr *removedExpr = analyzedExpr;
+          removeCodeCompletionExpr(CurDeclContext->getASTContext(),
+                                   removedExpr);
+          ConcreteDeclRef referencedDecl;
+          auto optT = getTypeOfCompletionContextExpr(
+              CurDeclContext->getASTContext(), CurDeclContext,
+              CompletionTypeCheckKind::Normal, removedExpr, referencedDecl);
+          if (optT) {
+            resultTy = *optT;
+            analyzedExpr->setType(resultTy);
+          }
+        }
+        if (!resultTy || resultTy->is<ErrorType>()) {
+          // 2) Infer it from the possible callee info.
+          if (!ContextInfo.getPossibleCallees().empty()) {
+            auto calleeInfo = ContextInfo.getPossibleCallees()[0];
+            resultTy = calleeInfo.Type->getResult();
+            analyzedExpr->setType(resultTy);
+          }
+        }
+        if (!resultTy || resultTy->is<ErrorType>()) {
+          // 3) Give up providing postfix completions.
+          break;
         }
 
         auto &SM = CurDeclContext->getASTContext().SourceMgr;

lib/IDE/ExprContextAnalysis.cpp

@@ -12,6 +12,7 @@
 
 #include "ExprContextAnalysis.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DeclContext.h"
@@ -181,6 +182,81 @@ Expr *swift::ide::findParsedExpr(const DeclContext *DC,
   return finder.get();
 }
 
+//===----------------------------------------------------------------------===//
+// removeCodeCompletionExpr(ASTContext, Expr)
+//===----------------------------------------------------------------------===//
+
+namespace {
+// TODO: Implement other expressions?
+class CCExprRemover: public ASTWalker, public ExprVisitor<CCExprRemover, Expr *> {
+  ASTContext &Ctx;
+
+public:
+  bool Removed = false;
+
+  CCExprRemover(ASTContext &Ctx) : Ctx(Ctx) {}
+
+  Expr *visitCallExpr(CallExpr *E) {
+    SourceLoc lParenLoc, rParenLoc;
+    SmallVector<Identifier, 2> argLabels;
+    SmallVector<SourceLoc, 2> argLabelLocs;
+    SmallVector<Expr *, 2> args;
+    SmallVector<TrailingClosure, 2> trailingClosures;
+    bool removing = false;
+
+    if (auto paren = dyn_cast<ParenExpr>(E->getArg())) {
+      if (isa<CodeCompletionExpr>(paren->getSubExpr())) {
+        lParenLoc = paren->getLParenLoc();
+        rParenLoc = paren->getRParenLoc();
+        removing = true;
+      }
+    } else if (auto tuple = dyn_cast<TupleExpr>(E->getArg())) {
+      lParenLoc = tuple->getLParenLoc();
+      rParenLoc = tuple->getRParenLoc();
+      for (unsigned i = 0, e = tuple->getNumElements(); i != e; ++i) {
+        if (isa<CodeCompletionExpr>(tuple->getElement(i))) {
+          removing = true;
+          continue;
+        }
+
+        if (i < E->getUnlabeledTrailingClosureIndex()) {
+          // Normal arguments.
+          argLabels.push_back(E->getArgumentLabels()[i]);
+          argLabelLocs.push_back(E->getArgumentLabelLocs()[i]);
+          args.push_back(tuple->getElement(i));
+        } else {
+          // Trailing closure arguments.
+          trailingClosures.emplace_back(E->getArgumentLabels()[i],
+                                        E->getArgumentLabelLocs()[i],
+                                        tuple->getElement(i));
+        }
+      }
+    }
+    if (removing) {
+      Removed = true;
+      return CallExpr::create(Ctx, E->getFn(), lParenLoc, args, argLabels,
+                              argLabelLocs, rParenLoc, trailingClosures,
+                              E->isImplicit());
+    }
+    return E;
+  }
+
+  Expr *visitExpr(Expr *E) {
+    return E;
+  }
+
+  std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+    return {true, visit(E)};
+  }
+};
+}
+
+bool swift::ide::removeCodeCompletionExpr(ASTContext &Ctx, Expr *&expr) {
+  CCExprRemover remover(Ctx);
+  expr = expr->walk(remover);
+  return remover.Removed;
+}
+
 //===----------------------------------------------------------------------===//
 // getReturnTypeFromContext(DeclContext)
 //===----------------------------------------------------------------------===//

lib/IDE/ExprContextAnalysis.h

@@ -34,6 +34,13 @@ void typeCheckContextAt(DeclContext *DC, SourceLoc Loc);
 /// exact the same as \p TargetRange. Returns \c nullptr if not found.
 Expr *findParsedExpr(const DeclContext *DC, SourceRange TargetRange);
 
+/// Remove \c CodeCompletionExpr from \p expr . Returns \c true if it actually
+/// mutated the expression.
+///
+/// NOTE: Currently, this only removes CodeCompletionExpr at call argument
+///       position.
+bool removeCodeCompletionExpr(ASTContext &Ctx, Expr *&expr);
+
 /// Returns expected return type of the given decl context.
 /// \p DC should be an \c AbstractFunctionDecl or an \c AbstractClosureExpr.
 Type getReturnTypeFromContext(const DeclContext *DC);

test/IDE/complete_multiple_trailingclosure.swift

@@ -13,6 +13,8 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=INIT_REQUIRED_NEWLINE_2 | %FileCheck %s -check-prefix=INIT_REQUIRED_NEWLINE_2
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=INIT_REQUIRED_SAMELINE_3 | %FileCheck %s -check-prefix=INIT_REQUIRED_SAMELINE_3
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=INIT_REQUIRED_NEWLINE_3 | %FileCheck %s -check-prefix=INIT_REQUIRED_NEWLINE_3
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=INIT_FALLBACK_1 | %FileCheck %s -check-prefix=INIT_FALLBACK
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=INIT_FALLBACK_2 | %FileCheck %s -check-prefix=INIT_FALLBACK
 
 func globalFunc1(fn1: () -> Int, fn2: () -> String) {}
 func testGlobalFunc() {
@@ -184,3 +186,24 @@ func testOptionalInit() {
 // INIT_REQUIRED_NEWLINE_3-NOT: name=fn3
 // INIT_REQUIRED_NEWLINE_3: End completions
 }
+
+struct MyStruct4<T> {
+  init(arg1: Int = 0, arg2: () -> T) {}
+  init(name: String, arg2: () -> String, arg3: () -> T) {}
+
+  func testStructMethod() {}
+}
+func testFallbackPostfix() {
+  let _ = MyStruct4 {
+    1
+  } #^INIT_FALLBACK_1^#
+// INIT_FALLBACK: Begin completions, 2 items
+// INIT_FALLBACK-DAG: Decl[InstanceMethod]/CurrNominal:   .testStructMethod()[#Void#];
+// INIT_FALLBACK-DAG: Keyword[self]/CurrNominal:          .self[#MyStruct4<Int>#];
+// INIT_FALLBACK: End completions
+  let _ = MyStruct4(name: "test") {
+    ""
+  } arg3: {
+    1
+  } #^INIT_FALLBACK_2^#
+}