[CSBindings] Solver result builder transformed closures as soon as all contextual information becomes available

We determine that based on four criteria:

- Builder type doesn't have any unresolved generic parameters;
- Contextual type associated with the closure doesn't have any unresolved type variables (which means that all possible contextual information is provided to the body);
- All of the references to declarations from outer scope are resolved (i.e. a variable declaration from a parent closure);
- The contextual result type is either fully resolved or opaque type.

If all the criteria a met the conjunction that represents a closure is going to be picked and solved as the next solver step.

Resolves: rdar://104645543

(cherry picked from commit f27369d)
(cherry picked from commit a815de7)
(cherry picked from commit 38f8be1)
(cherry picked from commit 5bb9d75)
(cherry picked from commit fda7530)
(cherry picked from commit b84d70b)

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -486,6 +486,9 @@ class TypeVariableType::Implementation {
   /// expression.
   bool isCodeCompletionToken() const;
 
+  /// Determine whether this type variable represents an opened opaque type.
+  bool isOpaqueType() const;
+
   /// Retrieve the representative of the equivalence class to which this
   /// type variable belongs.
   ///
@@ -977,6 +980,11 @@ struct AppliedBuilderTransform {
   /// converted. Opaque types should be unopened.
   Type bodyResultType;
 
+  /// If transform is applied to a closure, this type represents
+  /// contextual type the closure is converted type (e.g. a parameter
+  /// type or or pattern type).
+  Type contextualType;
+
   /// The version of the original body with result builder applied
   /// as AST transformation.
   NullablePtr<BraceStmt> transformedBody;
@@ -5662,7 +5670,7 @@ class ConstraintSystem {
   Optional<TypeMatchResult>
   matchResultBuilder(AnyFunctionRef fn, Type builderType, Type bodyResultType,
                      ConstraintKind bodyResultConstraintKind,
-                     ConstraintLocatorBuilder locator);
+                     Type contextualType, ConstraintLocatorBuilder locator);
 
   /// Matches a wrapped or projected value parameter type to its backing
   /// property wrapper type by applying the property wrapper.

lib/Sema/BuilderTransform.cpp

@@ -2311,6 +2311,7 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
 
   if (auto result = cs.matchResultBuilder(
           func, builderType, resultContextType, resultConstraintKind,
+          /*contextualType=*/Type(),
           cs.getConstraintLocator(func->getBody()))) {
     if (result->isFailure())
       return nullptr;
@@ -2398,6 +2399,7 @@ Optional<ConstraintSystem::TypeMatchResult>
 ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
                                      Type bodyResultType,
                                      ConstraintKind bodyResultConstraintKind,
+                                     Type contextualType,
                                      ConstraintLocatorBuilder locator) {
   builderType = simplifyType(builderType);
   auto builder = builderType->getAnyNominal();
@@ -2522,6 +2524,7 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
 
     transformInfo.builderType = builderType;
     transformInfo.bodyResultType = bodyResultType;
+    transformInfo.contextualType = contextualType;
     transformInfo.transformedBody = transformedBody->second;
 
     // Record the transformation.

lib/Sema/CSBindings.cpp

@@ -1071,6 +1071,61 @@ bool BindingSet::favoredOverConjunction(Constraint *conjunction) const {
     if (forClosureResult() || forGenericParameter())
       return false;
   }
+
+  auto *locator = conjunction->getLocator();
+  if (locator->directlyAt<ClosureExpr>()) {
+    auto *closure = castToExpr<ClosureExpr>(locator->getAnchor());
+
+    if (auto transform = CS.getAppliedResultBuilderTransform(closure)) {
+      // Conjunctions that represent closures with result builder transformed
+      // bodies could be attempted right after their resolution if they meet
+      // all of the following criteria:
+      //
+      // - Builder type doesn't have any unresolved generic parameters;
+      // - Closure doesn't have any parameters;
+      // - The contextual result type is either concrete or opaque type.
+      auto contextualType = transform->contextualType;
+      if (!(contextualType && contextualType->is<FunctionType>()))
+        return true;
+
+      auto *contextualFnType =
+          CS.simplifyType(contextualType)->castTo<FunctionType>();
+      {
+        auto resultType = contextualFnType->getResult();
+        if (resultType->hasTypeVariable()) {
+          auto *typeVar = resultType->getAs<TypeVariableType>();
+          // If contextual result type is represented by an opaque type,
+          // it's a strong indication that body is self-contained, otherwise
+          // closure might rely on external types flowing into the body for
+          // disambiguation of `build{Partial}Block` or `buildFinalResult`
+          // calls.
+          if (!(typeVar && typeVar->getImpl().isOpaqueType()))
+            return true;
+        }
+      }
+
+      // If some of the closure parameters are unresolved, the conjunction
+      // has to be delayed to give them a chance to be inferred.
+      if (llvm::any_of(contextualFnType->getParams(), [](const auto &param) {
+            return param.getPlainType()->hasTypeVariable();
+          }))
+        return true;
+
+      // Check whether conjunction has any unresolved type variables
+      // besides the variable that represents the closure.
+      //
+      // Conjunction could refer to declarations from outer context
+      // (i.e. a variable declared in the outer closure) or generic
+      // parameters of the builder type), if any of such references
+      // are not yet inferred the conjunction has to be delayed.
+      auto *closureType = CS.getType(closure)->castTo<TypeVariableType>();
+      return llvm::any_of(
+          conjunction->getTypeVariables(), [&](TypeVariableType *typeVar) {
+            return !(typeVar == closureType || CS.getFixedType(typeVar));
+          });
+    }
+  }
+
   return true;
 }
 

lib/Sema/CSSimplify.cpp

@@ -10964,7 +10964,7 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
   if (resultBuilderType) {
     if (auto result = matchResultBuilder(
             closure, resultBuilderType, closureType->getResult(),
-            ConstraintKind::Conversion, locator)) {
+            ConstraintKind::Conversion, contextualType, locator)) {
       return result->isSuccess();
     }
   }

lib/Sema/TypeCheckConstraints.cpp

@@ -154,6 +154,22 @@ bool TypeVariableType::Implementation::isCodeCompletionToken() const {
   return locator && locator->directlyAt<CodeCompletionExpr>();
 }
 
+bool TypeVariableType::Implementation::isOpaqueType() const {
+  if (!locator)
+    return false;
+
+  auto GP = locator->getLastElementAs<LocatorPathElt::GenericParameter>();
+  if (!GP)
+    return false;
+
+  if (auto *GPT = GP->getType()->getAs<GenericTypeParamType>()) {
+    auto *decl = GPT->getDecl();
+    return decl && decl->isOpaqueType();
+  }
+
+  return false;
+}
+
 void *operator new(size_t bytes, ConstraintSystem& cs,
                    size_t alignment) {
   return cs.getAllocator().Allocate(bytes, alignment);

test/Constraints/result_builder_conjunction_selection.swift

@@ -0,0 +1,83 @@
+// RUN: %target-typecheck-verify-swift -debug-constraints -disable-availability-checking 2>%t.err
+// RUN: %FileCheck %s < %t.err
+
+protocol P<Output> {
+  associatedtype Output
+}
+
+struct S<Output> : P {
+  init(_: Output) {}
+}
+
+@resultBuilder
+struct Builder {
+  static func buildExpression<T>(_ e: T) -> T { e }
+
+  static func buildBlock<T1>(_ t1: T1) -> some P<T1> {
+    return S(t1)
+  }
+
+  static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> some P<(T1, T2)> {
+    return S((t1, t2))
+  }
+
+  static func buildBlock<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3) -> some P<(T1, T2, T3)> {
+    return S((t1, t2, t3))
+  }
+
+  static func buildOptional<T>(_ value: T?) -> T? { return value }
+}
+
+do {
+  func test<T, U>(_: T, @Builder _: () -> some P<U>) {}
+
+  // CHECK: ---Initial constraints for the given expression---
+  // CHECK: (integer_literal_expr type='[[LITERAL_VAR:\$T[0-9]+]]' {{.*}}
+  // CHECK: (attempting type variable [[CLOSURE:\$T[0-9]+]] := () -> {{.*}}
+  // CHECK-NOT: (attempting type variable [[LITERAL_VAR]] := {{.*}}
+  // CHECK: (attempting conjunction element pattern binding element @ 0
+  // CHECK: (applying conjunction result to outer context
+  // CHECK: (attempting type variable [[LITERAL_VAR]] := Int
+  test(42) {
+    1
+    ""
+  }
+}
+
+do {
+  func test<T, U>(_: T, @Builder _: (T) -> some P<U>) {}
+
+  // CHECK: ---Initial constraints for the given expression---
+  // CHECK: (integer_literal_expr type='[[LITERAL_VAR:\$T[0-9]+]]' {{.*}}
+  // CHECK: (attempting type variable [[LITERAL_VAR]] := Int
+  // CHECK: (attempting conjunction element pattern binding element @ 0
+  test(42) { v in
+    v
+    ""
+  }
+}
+
+do {
+  func test<T: BinaryInteger, U>(@Builder _: (Bool) -> some P<T>, transform: (T?) -> U) {}
+
+  // CHECK: ---Initial constraints for the given expression---
+  // CHECK: (attempting type variable {{.*}} := () -> {{.*}}
+  // CHECK: (attempting conjunction element pattern binding element @ 0 :
+  // CHECK-NEXT: (pattern_named 'x')
+  // CHECK: (attempting conjunction element syntactic element
+  // CHECK-NEXT: (call_expr {{.*}}
+  // CHECK: (attempting type variable {{.*}} := (Bool) -> {{.*}}
+  // CHECK: (attempting conjunction element pattern binding element @ 0
+  // CHECK: (pattern_named implicit '$__builder{{.*}}')
+  // CHECK: (applying conjunction result to outer context
+  // CHECK: (attempting type variable {{.*}} := (Int?) -> {{.*}}
+  // CHECK: (attempting disjunction choice {{.*}} bound to decl {{.*}}.Int.init(_:)
+  let _ = {
+    let x = 42
+    test { cond in
+      x
+    } transform: { v in
+      Int(v ?? 0)
+    }
+  }
+}