From 96e6952104f848ddf6963aad4c33fcae2cc471ea Mon Sep 17 00:00:00 2001
From: Pavel Yaskevich <pyaskevich@apple.com>
Date: Thu, 5 May 2022 13:28:55 -0700
Subject: [PATCH] [CSClosure] Teach syntactic element constraint generator
 about result builders

Closures and functions that have result builder transform applied have to be
handled specially by the constraint generator because transformed body is
associated with the context (function or closure) only during solution application.
---
 include/swift/Sema/ConstraintSystem.h |  8 +++++
 lib/Sema/CSClosure.cpp                | 51 +++++++++++++++++++--------
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
index f3096cbc7a4..1dc3fcd6e01 100644
--- a/include/swift/Sema/ConstraintSystem.h
+++ b/include/swift/Sema/ConstraintSystem.h
@@ -3734,6 +3734,14 @@ public:
     return known->second;
   }
 
+  Optional<AppliedBuilderTransform>
+  getAppliedResultBuilderTransform(AnyFunctionRef fn) const {
+    auto transformed = resultBuilderTransformed.find(fn);
+    if (transformed != resultBuilderTransformed.end())
+      return transformed->second;
+    return None;
+  }
+
   void setCaseLabelItemInfo(const CaseLabelItem *item, CaseLabelItemInfo info) {
     assert(item != nullptr);
     assert(caseLabelItems.count(item) == 0);
diff --git a/lib/Sema/CSClosure.cpp b/lib/Sema/CSClosure.cpp
index 21c92526f3a..2618f3115b7 100644
--- a/lib/Sema/CSClosure.cpp
+++ b/lib/Sema/CSClosure.cpp
@@ -874,30 +874,27 @@ private:
   }
 
   void visitReturnStmt(ReturnStmt *returnStmt) {
-    auto *closure =
-        dyn_cast_or_null<ClosureExpr>(context.getAbstractClosureExpr());
-
     // Single-expression closures are effectively a `return` statement,
     // so let's give them a special locator as to indicate that.
     // Return statements might not have a result if we have a closure whose
     // implicit returned value is coerced to Void.
-    if (closure && closure->hasSingleExpressionBody() &&
-        returnStmt->hasResult()) {
+    if (isInSingleExpressionClosure() && returnStmt->hasResult()) {
       auto *expr = returnStmt->getResult();
       assert(expr && "single expression closure without expression?");
 
-      expr = cs.generateConstraints(expr, closure,
+      expr = cs.generateConstraints(expr, context.getAsDeclContext(),
                                     /*isInputExpression=*/false);
       if (!expr) {
         hadError = true;
         return;
       }
 
-      cs.addConstraint(
-          ConstraintKind::Conversion, cs.getType(expr), resultType,
-          cs.getConstraintLocator(
-              closure, LocatorPathElt::ClosureBody(
-                           /*hasReturn=*/!returnStmt->isImplicit())));
+      cs.addConstraint(ConstraintKind::Conversion, cs.getType(expr),
+                       getContextualResultType(),
+                       cs.getConstraintLocator(
+                           context.getAbstractClosureExpr(),
+                           LocatorPathElt::ClosureBody(
+                               /*hasReturn=*/!returnStmt->isImplicit())));
       return;
     }
 
@@ -923,14 +920,40 @@ private:
       return;
     }
 
-    cs.setContextualType(target.getAsExpr(), TypeLoc::withoutLoc(resultType),
+    cs.setContextualType(target.getAsExpr(),
+                         TypeLoc::withoutLoc(getContextualResultType()),
                          CTP_ReturnStmt);
     cs.setSolutionApplicationTarget(returnStmt, target);
   }
 
-  bool isSupportedMultiStatementClosure() const {
+  bool isInSingleExpressionClosure() {
+    if (!isExpr<ClosureExpr>(context.getAbstractClosureExpr()))
+      return false;
+
+    // Result builder transformed bodies are never single-expression.
+    if (cs.getAppliedResultBuilderTransform(context))
+      return false;
+
+    return context.hasSingleExpressionBody();
+  }
+
+  Type getContextualResultType() const {
+    if (auto transform = cs.getAppliedResultBuilderTransform(context))
+      return transform->bodyResultType;
+
     if (auto *closure =
-            dyn_cast_or_null<ClosureExpr>(context.getAbstractClosureExpr())) {
+            getAsExpr<ClosureExpr>(context.getAbstractClosureExpr()))
+      return cs.getClosureType(closure)->getResult();
+
+    return context.getBodyResultType();
+  }
+
+  bool isSupportedMultiStatementClosure() const {
+    if (cs.getAppliedResultBuilderTransform(context))
+      return true;
+
+    if (auto *closure =
+            getAsExpr<ClosureExpr>(context.getAbstractClosureExpr())) {
       return !closure->hasSingleExpressionBody() &&
              cs.participatesInInference(closure);
     }