Merge pull request #63677 from xedin/dont-delay-result-builder-conjunctions

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

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -487,6 +487,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.
   ///
@@ -985,6 +988,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;
@@ -5727,7 +5735,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

@@ -2341,6 +2341,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;
@@ -2428,6 +2429,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();
@@ -2552,6 +2554,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

@@ -11036,7 +11036,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)
+    }
+  }
+}