Merge pull request #28837 from xedin/rdar-56340587

[ConstraintSystem] Delay constraint generation for single-statement closure body

Author: xedin

include/swift/AST/Types.h

@@ -3388,7 +3388,6 @@ END_CAN_TYPE_WRAPPER(FunctionType, AnyFunctionType)
 /// has a default argument.
 struct ParameterListInfo {
   SmallBitVector defaultArguments;
-  std::vector<Type> functionBuilderTypes;
 
 public:
   ParameterListInfo() { }
@@ -3406,9 +3405,6 @@ struct ParameterListInfo {
 
   /// Retrieve the number of parameters for which we have information.
   unsigned size() const { return defaultArguments.size(); }
-
-  /// Retrieve the function builder type for the given parameter.
-  Type getFunctionBuilderType(unsigned paramIdx) const;
 };
 
 /// Turn a param list into a symbolic and printable representation that does not

lib/AST/Type.cpp

@@ -818,7 +818,6 @@ ParameterListInfo::ParameterListInfo(
     const ValueDecl *paramOwner,
     bool skipCurriedSelf) {
   defaultArguments.resize(params.size());
-  functionBuilderTypes.resize(params.size());
 
   // No parameter owner means no parameter list means no default arguments
   // - hand back the zeroed bitvector.
@@ -863,10 +862,6 @@ ParameterListInfo::ParameterListInfo(
     if (param->isDefaultArgument()) {
       defaultArguments.set(i);
     }
-
-    if (Type functionBuilderType = param->getFunctionBuilderType()) {
-      functionBuilderTypes[i] = functionBuilderType;
-    }
   }
 }
 
@@ -875,12 +870,6 @@ bool ParameterListInfo::hasDefaultArgument(unsigned paramIdx) const {
       : false;
 }
 
-Type ParameterListInfo::getFunctionBuilderType(unsigned paramIdx) const {
-  return paramIdx < functionBuilderTypes.size()
-      ? functionBuilderTypes[paramIdx]
-      : Type();
-}
-
 /// Turn a param list into a symbolic and printable representation that does not
 /// include the types, something like (_:, b:, c:)
 std::string swift::getParamListAsString(ArrayRef<AnyFunctionType::Param> params) {

lib/Sema/BuilderTransform.cpp

@@ -615,13 +615,6 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::matchFunctionBuilder(
   assert(builder && "Bad function builder type");
   assert(builder->getAttrs().hasAttribute<FunctionBuilderAttr>());
 
-  // FIXME: Right now, single-expression closures suppress the function
-  // builder translation.
-  if (auto closure = fn.getAbstractClosureExpr()) {
-    if (closure->hasSingleExpressionBody())
-      return getTypeMatchSuccess();
-  }
-
   // Pre-check the body: pre-check any expressions in it and look
   // for return statements.
   auto request = PreCheckFunctionBuilderRequest{fn};
@@ -716,6 +709,12 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::matchFunctionBuilder(
         fn,
         AppliedBuilderTransform{builderType, singleExpr, bodyResultType}));
 
+  // If builder is applied to the closure expression then
+  // `closure body` to `closure result` matching should
+  // use special locator.
+  if (auto *closure = fn.getAbstractClosureExpr())
+    locator = getConstraintLocator(closure, ConstraintLocator::ClosureResult);
+
   // Bind the body result type to the type of the transformed expression.
   addConstraint(bodyResultConstraintKind, transformedType, bodyResultType,
                 locator);

lib/Sema/CSApply.cpp

@@ -152,6 +152,18 @@ static bool shouldAccessStorageDirectly(Expr *base, VarDecl *member,
   return true;
 }
 
+ConstraintLocator *Solution::getCalleeLocator(ConstraintLocator *locator,
+                                              bool lookThroughApply) const {
+  auto &cs = getConstraintSystem();
+  return cs.getCalleeLocator(
+      locator, lookThroughApply,
+      [&](const Expr *expr) -> Type { return getType(expr); },
+      [&](Type type) -> Type { return simplifyType(type)->getRValueType(); },
+      [&](ConstraintLocator *locator) -> Optional<SelectedOverload> {
+        return getOverloadChoiceIfAvailable(locator);
+      });
+}
+
 /// Return the implicit access kind for a MemberRefExpr with the
 /// specified base and member in the specified DeclContext.
 static AccessSemantics
@@ -7412,6 +7424,21 @@ ProtocolConformanceRef Solution::resolveConformance(
   return ProtocolConformanceRef::forInvalid();
 }
 
+Type Solution::getType(const Expr *expr) const {
+  auto result = llvm::find_if(
+      addedNodeTypes, [&](const std::pair<TypedNode, Type> &node) -> bool {
+        if (auto *e = node.first.dyn_cast<const Expr *>())
+          return expr == e;
+        return false;
+      });
+
+  if (result != addedNodeTypes.end())
+    return result->second;
+
+  auto &cs = getConstraintSystem();
+  return cs.getType(expr);
+}
+
 void Solution::setExprTypes(Expr *expr) const {
   if (!expr)
     return;

lib/Sema/CSBindings.cpp

@@ -228,6 +228,22 @@ bool ConstraintSystem::PotentialBindings::isViable(
   return true;
 }
 
+bool ConstraintSystem::PotentialBindings::favoredOverDisjunction() const {
+  if (IsHole || FullyBound)
+    return false;
+
+  // If this bindings are for a closure and there are no holes,
+  // it shouldn't matter whether it there are any type variables
+  // or not because e.g. parameter type can have type variables,
+  // but we still want to resolve closure body early (instead of
+  // attempting any disjunction) to gain additional contextual
+  // information.
+  if (TypeVar->getImpl().isClosureType())
+    return true;
+
+  return !InvolvesTypeVariables;
+}
+
 static bool hasNilLiteralConstraint(TypeVariableType *typeVar,
                                     const ConstraintSystem &CS) {
   // Look for a literal-conformance constraint on the type variable.
@@ -353,6 +369,11 @@ ConstraintSystem::getPotentialBindingForRelationalConstraint(
 
     result.InvolvesTypeVariables = true;
 
+    if (constraint->getKind() == ConstraintKind::Subtype &&
+        kind == AllowedBindingKind::Subtypes) {
+      result.SubtypeOf.insert(bindingTypeVar);
+    }
+
     // If we've already set addOptionalSupertypeBindings, or we aren't
     // allowing supertype bindings, we're done.
     if (addOptionalSupertypeBindings || kind != AllowedBindingKind::Supertypes)
@@ -525,6 +546,7 @@ ConstraintSystem::getPotentialBindings(TypeVariableType *typeVar) const {
     }
 
     case ConstraintKind::Defaultable:
+    case ConstraintKind::DefaultClosureType:
       // Do these in a separate pass.
       if (getFixedTypeRecursive(constraint->getFirstType(), true)
               ->getAs<TypeVariableType>() == typeVar) {
@@ -749,6 +771,15 @@ ConstraintSystem::getPotentialBindings(TypeVariableType *typeVar) const {
     if (!exactTypes.insert(type->getCanonicalType()).second)
       continue;
 
+    if (constraint->getKind() == ConstraintKind::DefaultClosureType) {
+      // If there are no other possible bindings for this closure
+      // let's default it to the type inferred from its parameters/body,
+      // otherwise we should only attempt contextual types as a
+      // top-level closure type.
+      if (!result.Bindings.empty())
+        continue;
+    }
+
     result.addPotentialBinding({type, AllowedBindingKind::Exact, constraint});
   }
 

lib/Sema/CSDiagnostics.cpp

@@ -639,6 +639,18 @@ bool GenericArgumentsMismatchFailure::diagnoseAsError() {
     case ConstraintLocator::ContextualType: {
       auto purpose = getContextualTypePurpose();
       assert(!(purpose == CTP_Unused && purpose == CTP_CannotFail));
+
+      // If this is call to a closure e.g. `let _: A = { B() }()`
+      // let's point diagnostic to its result.
+      if (auto *call = dyn_cast<CallExpr>(anchor)) {
+        auto *fnExpr = call->getFn();
+        if (auto *closure = dyn_cast<ClosureExpr>(fnExpr)) {
+          purpose = CTP_ClosureResult;
+          if (closure->hasSingleExpressionBody())
+            anchor = closure->getSingleExpressionBody();
+        }
+      }
+
       diagnostic = getDiagnosticFor(purpose);
       break;
     }
@@ -1958,6 +1970,11 @@ bool ContextualFailure::diagnoseAsError() {
       return true;
     }
 
+    if (auto *call = dyn_cast<CallExpr>(anchor)) {
+      if (isa<ClosureExpr>(call->getFn()))
+        CTP = CTP_ClosureResult;
+    }
+
     if (auto msg = getDiagnosticFor(CTP, toType->isExistentialType())) {
       diagnostic = *msg;
       break;
@@ -2621,7 +2638,15 @@ bool ContextualFailure::trySequenceSubsequenceFixIts(
       if (auto *CE = dyn_cast<CoerceExpr>(anchor)) {
         anchor = CE->getSubExpr();
       }
-      
+
+      if (auto *call = dyn_cast<CallExpr>(anchor)) {
+        auto *fnExpr = call->getFn();
+        if (auto *closure = dyn_cast<ClosureExpr>(fnExpr)) {
+          if (closure->hasSingleExpressionBody())
+            anchor = closure->getSingleExpressionBody();
+        }
+      }
+
       auto range = anchor->getSourceRange();
       diagnostic.fixItInsert(range.Start, "String(");
       diagnostic.fixItInsertAfter(range.End, ")");