[Parse][CodeCompletion] Don't special case code completion when forming single-expression closures/function bodies (NFC)

Code completion used to avoid forming single expression closures/function
bodies when the single expression contained the code completion expression
because a contextual type mismatch could result in types not being applied
to the AST, giving no completions.

Completions that have been migrated to the new solver-based completion
mechanism don't need this behavior, however. Rather than trying to guess
whether the type of completion we're going to end up performing is one of
the ones that haven't been migrated to the solver yet when parsing, instead
just always form single-expression closures/function bodies (like we do for
regular compilation) and undo the transformation if and when we know we're
going to perform a completion kind we haven't migrated yet.

Once all completion kinds are migrated, the undo-ing code can be removed.

Author: Nathan Hawes

include/swift/Parse/Parser.h

@@ -709,15 +709,6 @@ class Parser {
       Context.LangOpts.ParseForSyntaxTreeOnly;
   }
 
-  /// If a function or closure body consists of a single expression, determine
-  /// whether we should turn wrap it in a return statement or not.
-  ///
-  /// We don't do this transformation for non-solver-based code completion
-  /// positions, as the source may be incomplete and the type mismatch in the
-  /// return statement will just confuse the type checker.
-  bool shouldSuppressSingleExpressionBodyTransform(
-      ParserStatus Status, MutableArrayRef<ASTNode> BodyElems);
-
 public:
   InFlightDiagnostic diagnose(SourceLoc Loc, Diagnostic Diag) {
     if (Diags.isDiagnosticPointsToFirstBadToken(Diag.getID()) &&

lib/IDE/CodeCompletion.cpp

@@ -1672,7 +1672,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
 
 private:
   void addKeywords(CodeCompletionResultSink &Sink, bool MaybeFuncBody);
-  bool trySolverCompletion();
+  bool trySolverCompletion(bool MaybeFuncBody);
 };
 } // end anonymous namespace
 
@@ -6054,25 +6054,7 @@ void deliverDotExprResults(
   deliverCompletionResults(CompletionCtx, Lookup, *SF, Consumer);
 }
 
-bool CodeCompletionCallbacksImpl::trySolverCompletion() {
-  CompletionContext.CodeCompletionKind = Kind;
-
-  if (Kind == CompletionKind::None)
-    return true;
-
-  bool MaybeFuncBody = true;
-  if (CurDeclContext) {
-    auto *CD = CurDeclContext->getLocalContext();
-    if (!CD || CD->getContextKind() == DeclContextKind::Initializer ||
-        CD->getContextKind() == DeclContextKind::TopLevelCodeDecl)
-      MaybeFuncBody = false;
-  }
-
-  if (auto *DC = dyn_cast_or_null<DeclContext>(ParsedDecl)) {
-    if (DC->isChildContextOf(CurDeclContext))
-      CurDeclContext = DC;
-  }
-
+bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
   assert(ParsedExpr || CurDeclContext);
 
   SourceLoc CompletionLoc = ParsedExpr
@@ -6111,12 +6093,42 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion() {
   }
 }
 
+// Undoes the single-expression closure/function body transformation on the
+// given DeclContext and its parent contexts if they have a single expression
+// body that contains the code completion location.
+//
+// FIXME: Remove this once all expression position completions are migrated
+// to work via TypeCheckCompletionCallback.
+static void undoSingleExpressionReturn(DeclContext *DC) {
+  auto updateBody = [](BraceStmt *BS, ASTContext &Ctx) -> bool {
+    ASTNode FirstElem = BS->getFirstElement();
+    auto *RS = dyn_cast_or_null<ReturnStmt>(FirstElem.dyn_cast<Stmt*>());
+
+    if (!RS || !RS->isImplicit())
+      return false;
+
+    BS->setFirstElement(RS->getResult());
+    return true;
+  };
+
+  while (ClosureExpr *CE = dyn_cast_or_null<ClosureExpr>(DC)) {
+    if (CE->hasSingleExpressionBody()) {
+      if (updateBody(CE->getBody(), CE->getASTContext()))
+        CE->setBody(CE->getBody(), false);
+    }
+    DC = DC->getParent();
+  }
+  if (FuncDecl *FD = dyn_cast_or_null<FuncDecl>(DC)) {
+    if (FD->hasSingleExpressionBody()) {
+      if (updateBody(FD->getBody(), FD->getASTContext()))
+        FD->setHasSingleExpressionBody(false);
+    }
+  }
+}
+
 void CodeCompletionCallbacksImpl::doneParsing() {
   CompletionContext.CodeCompletionKind = Kind;
 
-  if (trySolverCompletion())
-    return;
-
   if (Kind == CompletionKind::None) {
     return;
   }
@@ -6134,6 +6146,10 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       CurDeclContext = DC;
   }
 
+  if (trySolverCompletion(MaybeFuncBody))
+    return;
+
+  undoSingleExpressionReturn(CurDeclContext);
   typeCheckContextAt(
       CurDeclContext,
       ParsedExpr

lib/Parse/ParseDecl.cpp

@@ -6588,11 +6588,7 @@ BraceStmt *Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
 
   // If the body consists of a single expression, turn it into a return
   // statement.
-  //
-  // But don't do this transformation when performing certain kinds of code
-  // completion, as the source may be incomplete and the type mismatch in return
-  // statement will just confuse the type checker.
-  if (shouldSuppressSingleExpressionBodyTransform(Body, BS->getElements()))
+  if (BS->getNumElements() != 1)
     return BS;
 
   auto Element = BS->getFirstElement();

lib/Parse/ParseExpr.cpp

@@ -2841,13 +2841,8 @@ ParserResult<Expr> Parser::parseExprClosure() {
 
   // If the body consists of a single expression, turn it into a return
   // statement.
-  //
-  // But don't do this transformation when performing certain kinds of code
-  // completion, as the source may be incomplete and the type mismatch in return
-  // statement will just confuse the type checker.
   bool hasSingleExpressionBody = false;
-  if (!missingRBrace &&
-      !shouldSuppressSingleExpressionBodyTransform(Status, bodyElements)) {
+  if (!missingRBrace && bodyElements.size() == 1) {
     // If the closure's only body element is a single return statement,
     // use that instead of creating a new wrapping return expression.
     Expr *returnExpr = nullptr;

lib/Parse/Parser.cpp

@@ -1185,34 +1185,6 @@ Parser::getStringLiteralIfNotInterpolated(SourceLoc Loc,
                                                Segments.front().Length));
 }
 
-bool Parser::
-shouldSuppressSingleExpressionBodyTransform(ParserStatus Status,
-                                            MutableArrayRef<ASTNode> BodyElems) {
-  if (BodyElems.size() != 1)
-    return true;
-
-  if (!Status.hasCodeCompletion())
-    return false;
-
-  struct HasMemberCompletion: public ASTWalker {
-    bool Value = false;
-    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      if (auto *CCE = dyn_cast<CodeCompletionExpr>(E)) {
-        // If it has a base expression this is member completion, which is
-        // performed using the new solver-based mechanism, so it's ok to go
-        // ahead with the transform (and necessary to pick up the correct
-        // expected type).
-        Value = CCE->getBase();
-        return {false, nullptr};
-      }
-      return {true, E};
-    }
-  };
-  HasMemberCompletion Check;
-  BodyElems.front().walk(Check);
-  return !Check.Value;
-}
-
 struct ParserUnit::Implementation {
   std::shared_ptr<SyntaxParseActions> SPActions;
   LangOptions LangOpts;