Merge pull request swiftlang#34253 from nathawes/nfc-always-form-single-expression-return-and-undo-if-needed

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

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;