Merge pull request #60449 from xedin/dont-erase-result-during-transform

[ResultBuilder] AST transform: don't try type erasure of result expressions

Author: xedin

lib/Sema/BuilderTransform.cpp

@@ -919,8 +919,6 @@ class ResultBuilderTransform
 
   using UnsupportedElt = SkipUnhandledConstructInResultBuilder::UnhandledNode;
 
-  /// The constraint system this transform is associated with.
-  ConstraintSystem &CS;
   /// The result type of this result builder body.
   Type ResultType;
 
@@ -930,8 +928,7 @@ class ResultBuilderTransform
 public:
   ResultBuilderTransform(ConstraintSystem &cs, DeclContext *dc,
                          Type builderType, Type resultTy)
-      : BuilderTransformerBase(&cs, dc, builderType), CS(cs),
-        ResultType(resultTy) {}
+      : BuilderTransformerBase(&cs, dc, builderType), ResultType(resultTy) {}
 
   UnsupportedElt getUnsupportedElement() const { return FirstUnsupported; }
 
@@ -1180,10 +1177,6 @@ class ResultBuilderTransform
                               {buildBlockResult}, {Identifier()});
       }
 
-      // Type erase return if the result type requires it.
-      buildBlockResult = CS.buildTypeErasedExpr(buildBlockResult, dc,
-                                                ResultType, CTP_ReturnStmt);
-
       elements.push_back(new (ctx) ReturnStmt(resultLoc, buildBlockResult,
                                               /*Implicit=*/true));
     }

lib/Sema/CSClosure.cpp

@@ -1607,6 +1607,8 @@ class SyntacticElementSolutionApplication
   }
 
   ASTNode visitReturnStmt(ReturnStmt *returnStmt) {
+    auto &cs = solution.getConstraintSystem();
+
     if (!returnStmt->hasResult()) {
       // If contextual is not optional, there is nothing to do here.
       if (resultType->isVoid())
@@ -1621,8 +1623,6 @@ class SyntacticElementSolutionApplication
       assert(resultType->getOptionalObjectType() &&
              resultType->lookThroughAllOptionalTypes()->isVoid());
 
-      auto &cs = solution.getConstraintSystem();
-
       auto target = *cs.getSolutionApplicationTarget(returnStmt);
       returnStmt->setResult(target.getAsExpr());
     }
@@ -1652,13 +1652,24 @@ class SyntacticElementSolutionApplication
       mode = convertToResult;
     }
 
-    SolutionApplicationTarget resultTarget(
-        resultExpr, context.getAsDeclContext(),
-        mode == convertToResult ? CTP_ReturnStmt : CTP_Unused,
-        mode == convertToResult ? resultType : Type(),
-        /*isDiscarded=*/false);
-    if (auto newResultTarget = rewriteTarget(resultTarget))
+    Optional<SolutionApplicationTarget> resultTarget;
+    if (auto target = cs.getSolutionApplicationTarget(returnStmt)) {
+      resultTarget = *target;
+    } else {
+      // Single-expression closures have to handle returns in a special
+      // way so the target has to be created for them during solution
+      // application based on the resolved type.
+      assert(isSingleExpression);
+      resultTarget = SolutionApplicationTarget(
+          resultExpr, context.getAsDeclContext(),
+          mode == convertToResult ? CTP_ReturnStmt : CTP_Unused,
+          mode == convertToResult ? resultType : Type(),
+          /*isDiscarded=*/false);
+    }
+
+    if (auto newResultTarget = rewriteTarget(*resultTarget)) {
       resultExpr = newResultTarget->getAsExpr();
+    }
 
     switch (mode) {
     case convertToResult: