[CS] Generalize implied result handling

Track the implied result exprs in the constraint
system, and allow arbitrary propagation of
implied results down if/switch expression
branches. This is required for allowing implied
results in non-single-expression closures.

Author: hamishknight

include/swift/Sema/ConstraintLocator.h

@@ -102,18 +102,6 @@ namespace constraints {
 class ConstraintSystem;
 enum class ConversionRestrictionKind;
 
-/// The kind of SingleValueStmtExpr branch the locator identifies.
-enum class SingleValueStmtBranchKind {
-  /// Explicitly written 'then <expr>'.
-  Explicit,
-
-  /// Implicitly written '<expr>'.
-  Implicit,
-
-  /// Implicitly written '<expr>' in a single expr closure body.
-  ImplicitInSingleExprClosure
-};
-
 /// Locates a given constraint within the expression being
 /// type-checked, which may refer down into subexpressions and parts of
 /// the types of those subexpressions.
@@ -301,6 +289,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// Determine whether this locator points to the contextual type.
   bool isForContextualType() const;
 
+  /// Determine whether this locator points to the contextual type for a given
+  /// purpose.
+  bool isForContextualType(ContextualTypePurpose ctp) const;
+
   /// Determine whether this locator points to the assignment expression.
   bool isForAssignment() const;
 
@@ -325,8 +317,8 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   bool isForSingleValueStmtConjunctionOrBrace() const;
 
   /// Whether this locator identifies a conversion for a SingleValueStmtExpr
-  /// branch, and if so, the kind of branch.
-  llvm::Optional<SingleValueStmtBranchKind> isForSingleValueStmtBranch() const;
+  /// branch.
+  bool isForSingleValueStmtBranch() const;
 
   /// If the locator in question is for a pattern match, returns the pattern,
   /// otherwise \c nullptr.
@@ -818,21 +810,6 @@ class LocatorPathElt::TypeParameterRequirement final
   }
 };
 
-class LocatorPathElt::ClosureBody final : public StoredIntegerElement<1> {
-  public:
-  ClosureBody(bool hasImpliedReturn)
-    : StoredIntegerElement(ConstraintLocator::ClosureBody, hasImpliedReturn) {}
-
-  /// Indicates whether body of the closure has an implied `return` statement,
-  /// this is the case for single expression bodies where the `return` was not
-  /// written explicitly.
-  bool hasImpliedReturn() const { return bool(getValue()); }
-
-  static bool classof(const LocatorPathElt *elt) {
-    return elt->getKind() == ConstraintLocator::ClosureBody;
-  }
-};
-
 class LocatorPathElt::Witness final : public StoredPointerElement<ValueDecl> {
 public:
   Witness(ValueDecl *witness)

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -61,7 +61,7 @@ SIMPLE_LOCATOR_PATH_ELT(ClosureThrownError)
 /// closure returns.
 ///
 /// The type of a closure body.
-CUSTOM_LOCATOR_PATH_ELT(ClosureBody)
+SIMPLE_LOCATOR_PATH_ELT(ClosureBody)
 
 /// The lookup for a constructor member.
 SIMPLE_LOCATOR_PATH_ELT(ConstructorMember)

include/swift/Sema/ConstraintSystem.h

@@ -1468,6 +1468,15 @@ struct PotentialThrowSite {
   ConstraintLocator *locator;
 };
 
+enum class ImpliedResultKind {
+  /// A regular implied result, this applies to e.g implied 'then' statements
+  /// outside of closures.
+  Regular,
+
+  /// An implied result for a closure, e.g a single expression body.
+  ForClosure
+};
+
 /// A complete solution to a constraint system.
 ///
 /// A solution to a constraint system consists of type variable bindings to
@@ -1515,6 +1524,11 @@ class Solution {
   /// to make the solution work.
   std::vector<ConstraintFix *> Fixes;
 
+  /// Maps expressions for implied results (e.g implicit 'then' statements,
+  /// implicit 'return' statements in single expression body closures) to their
+  /// result kind.
+  llvm::MapVector<const Expr *, ImpliedResultKind> ImpliedResults;
+
   /// For locators associated with call expressions, the trailing closure
   /// matching rule and parameter bindings that were applied.
   llvm::MapVector<ConstraintLocator *, MatchCallArgumentResult>
@@ -1671,6 +1685,16 @@ class Solution {
     return DisjunctionChoices.find(locator)->second;
   }
 
+  /// Whether the given expression is the implied result for either a ReturnStmt
+  /// or ThenStmt, and if so, the kind of implied result.
+  llvm::Optional<ImpliedResultKind> isImpliedResult(const Expr *E) const {
+    auto result = ImpliedResults.find(E);
+    if (result == ImpliedResults.end())
+      return llvm::None;
+
+    return result->second;
+  }
+
   /// Retrieve the fixed score of this solution
   const Score &getFixedScore() const { return FixedScore; }
 
@@ -2196,6 +2220,11 @@ class ConstraintSystem {
   /// from declared parameters/result and body.
   llvm::MapVector<const ClosureExpr *, FunctionType *> ClosureTypes;
 
+  /// Maps expressions for implied results (e.g implicit 'then' statements,
+  /// implicit 'return' statements in single expression body closures) to their
+  /// result kind.
+  llvm::MapVector<const Expr *, ImpliedResultKind> ImpliedResults;
+
   /// This is a *global* list of all result builder bodies that have
   /// been determined to be incorrect by failing constraint generation.
   ///
@@ -2869,6 +2898,9 @@ class ConstraintSystem {
     /// The length of \c ClosureTypes.
     unsigned numInferredClosureTypes;
 
+    /// The length of \c ImpliedResults.
+    unsigned numImpliedResults;
+
     /// The length of \c contextualTypes.
     unsigned numContextualTypes;
 
@@ -3077,6 +3109,24 @@ class ConstraintSystem {
     return !IgnoredArguments.empty();
   }
 
+  /// Record an implied result for a ReturnStmt or ThenStmt.
+  void recordImpliedResult(const Expr *E, ImpliedResultKind kind) {
+    assert(E);
+    auto inserted = ImpliedResults.insert({E, kind}).second;
+    assert(inserted && "Duplicate implied result?");
+    (void)inserted;
+  }
+
+  /// Whether the given expression is the implied result for either a ReturnStmt
+  /// or ThenStmt, and if so, the kind of implied result.
+  llvm::Optional<ImpliedResultKind> isImpliedResult(const Expr *E) const {
+    auto result = ImpliedResults.find(E);
+    if (result == ImpliedResults.end())
+      return llvm::None;
+
+    return result->second;
+  }
+
   void setClosureType(const ClosureExpr *closure, FunctionType *type) {
     assert(closure);
     assert(type && "Expected non-null type");

lib/Sema/CSApply.cpp

@@ -9638,6 +9638,7 @@ ExprWalker::rewriteTarget(SyntacticElementTarget target) {
       if (solution.simplifyType(convertType)->isVoid()) {
         auto contextPurpose = cs.getContextualTypePurpose(target.getAsExpr());
         if (contextPurpose == CTP_ImpliedReturnStmt ||
+            contextPurpose == CTP_ClosureResult ||
             contextPurpose == CTP_SingleValueStmtBranch) {
           return false;
         }

lib/Sema/CSSimplify.cpp

@@ -7854,22 +7854,11 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
   }
 
   // Allow '() -> T' to '() -> ()' and '() -> Never' to '() -> T' for closure
-  // literals and expressions representing an implicit return type of the single
-  // expression functions.
+  // literals and expressions representing an implied result of closures and
+  // if/switch expressions.
   if (auto elt = locator.last()) {
     if (kind >= ConstraintKind::Subtype &&
         (type1->isUninhabited() || type2->isVoid())) {
-      // A conversion from closure body type to its signature result type.
-      if (auto resultElt = elt->getAs<LocatorPathElt::ClosureBody>()) {
-        // If a single statement closure has explicit `return` let's
-        // forbid conversion to `Void` and report an error instead to
-        // honor user's intent.
-        if (type1->isUninhabited() || resultElt->hasImpliedReturn()) {
-          increaseScore(SK_FunctionConversion, locator);
-          return getTypeMatchSuccess();
-        }
-      }
-
       // Function with an implied `return`, e.g a single expression body.
       if (auto contextualType = elt->getAs<LocatorPathElt::ContextualType>()) {
         if (contextualType->isFor(CTP_ImpliedReturnStmt)) {
@@ -7878,26 +7867,34 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
         }
       }
 
-      // We also need to propagate this conversion into the branches for single
-      // value statements.
+      // Implied results can occur for closure bodies, returns, and if/switch
+      // expression branches.
       //
-      // As with the previous checks, we only allow the Void conversion in
-      // an implicit single-expression closure. In the more general case for
-      // implicit branches, we only allow the Never conversion. For explicit
-      // branches, no conversions are allowed.
+      // We only allow the Void conversion for implied results in a closure
+      // context. In the more general case, we only allow the Never conversion.
+      // For explicit branches, no conversions are allowed, unless this is for
+      // a single expression body closure, in which case we still allow the
+      // Never conversion.
       auto *loc = getConstraintLocator(locator);
-      if (auto branchKind = loc->isForSingleValueStmtBranch()) {
+      if (elt->is<LocatorPathElt::ClosureBody>() || 
+          loc->isForContextualType(CTP_ClosureResult) ||
+          loc->isForSingleValueStmtBranch()) {
         bool allowConversion = false;
-        switch (*branchKind) {
-        case SingleValueStmtBranchKind::Explicit:
-          allowConversion = false;
-          break;
-        case SingleValueStmtBranchKind::Implicit:
-          allowConversion = type1->isUninhabited();
-          break;
-        case SingleValueStmtBranchKind::ImplicitInSingleExprClosure:
-          allowConversion = true;
-          break;
+        if (auto *E = getAsExpr(simplifyLocatorToAnchor(loc))) {
+          if (auto kind = isImpliedResult(E)) {
+            switch (*kind) {
+            case ImpliedResultKind::Regular:
+              allowConversion = type1->isUninhabited();
+              break;
+            case ImpliedResultKind::ForClosure:
+              allowConversion = true;
+              break;
+            }
+          } else if (elt->is<LocatorPathElt::ClosureBody>()) {
+            // Even if explicit, we always allow uninhabited types in single
+            // expression closures.
+            allowConversion = type1->isUninhabited();
+          }
         }
         if (allowConversion) {
           increaseScore(SK_FunctionConversion, locator);

lib/Sema/CSSolver.cpp

@@ -148,6 +148,10 @@ Solution ConstraintSystem::finalize() {
     (void)inserted;
   }
 
+  // Remember implied results.
+  for (auto impliedResult : ImpliedResults)
+    solution.ImpliedResults.insert(impliedResult);
+
   // Remember the opened types.
   for (const auto &opened : OpenedTypes) {
     // We shouldn't ever register opened types multiple times,
@@ -294,6 +298,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     argumentMatchingChoices.insert(argumentMatch);
   }
 
+  // Remember implied results.
+  for (auto impliedResult : solution.ImpliedResults)
+    ImpliedResults.insert(impliedResult);
+
   // Register the solution's opened types.
   for (const auto &opened : solution.OpenedTypes) {
     OpenedTypes.insert(opened);
@@ -667,6 +675,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numAppliedPropertyWrappers = cs.appliedPropertyWrappers.size();
   numResolvedOverloads = cs.ResolvedOverloads.size();
   numInferredClosureTypes = cs.ClosureTypes.size();
+  numImpliedResults = cs.ImpliedResults.size();
   numContextualTypes = cs.contextualTypes.size();
   numTargets = cs.targets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
@@ -789,6 +798,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any inferred closure types (e.g. used in result builder body).
   truncate(cs.ClosureTypes, numInferredClosureTypes);
 
+  // Remove any implied results.
+  truncate(cs.ImpliedResults, numImpliedResults);
+
   // Remove any contextual types.
   truncate(cs.contextualTypes, numContextualTypes);
 

lib/Sema/CSSyntacticElement.cpp

@@ -1242,6 +1242,14 @@ class SyntacticElementConstraintGenerator
   }
 
   void visitReturnStmt(ReturnStmt *returnStmt) {
+    // Record an implied result if we have one.
+    if (returnStmt->isImplied()) {
+      auto kind = context.getAsClosureExpr() ? ImpliedResultKind::ForClosure
+                                             : ImpliedResultKind::Regular;
+      auto *result = returnStmt->getResult();
+      cs.recordImpliedResult(result, kind);
+    }
+
     Expr *resultExpr;
     if (returnStmt->hasResult()) {
       resultExpr = returnStmt->getResult();
@@ -1316,11 +1324,9 @@ class SyntacticElementConstraintGenerator
       // on the closure itself. Otherwise we use the default contextual type
       // locator, which will be created for us.
       ConstraintLocator *loc = nullptr;
-      if (context.isSingleExpressionClosure(cs) && returnStmt->hasResult()) {
-        loc = cs.getConstraintLocator(
-            closure, {LocatorPathElt::ClosureBody(
-                         /*hasImpliedReturn=*/returnStmt->isImplied())});
-      }
+      if (context.isSingleExpressionClosure(cs) && returnStmt->hasResult())
+        loc = cs.getConstraintLocator(closure, {LocatorPathElt::ClosureBody()});
+
       return {cs.getClosureType(closure)->getResult(), CTP_ClosureResult, loc};
     }
 
@@ -1461,18 +1467,30 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
   Type resultTy = createTypeVariable(loc, /*options*/ 0);
   setType(E, resultTy);
 
+  // Propagate the implied result kind from the if/switch expression itself
+  // into the branches.
+  auto impliedResultKind =
+      isImpliedResult(E).value_or(ImpliedResultKind::Regular);
+
   // Assign contextual types for each of the result exprs.
-  SmallVector<Expr *, 4> scratch;
-  auto branches = E->getResultExprs(scratch);
+  SmallVector<ThenStmt *, 4> scratch;
+  auto branches = E->getThenStmts(scratch);
   for (auto idx : indices(branches)) {
-    auto *branch = branches[idx];
+    auto *thenStmt = branches[idx];
+    auto *result = thenStmt->getResult();
+
+    // If we have an implicit 'then' statement, record it as an implied result.
+    // TODO: Should we track 'implied' as a separate bit on ThenStmt? Currently
+    // it's the same as being implicit, but may not always be.
+    if (thenStmt->isImplicit())
+      recordImpliedResult(result, impliedResultKind);
 
     auto ctpElt = LocatorPathElt::ContextualType(CTP_SingleValueStmtBranch);
     auto *loc = getConstraintLocator(
         E, {LocatorPathElt::SingleValueStmtResult(idx), ctpElt});
 
     ContextualTypeInfo info(resultTy, CTP_SingleValueStmtBranch, loc);
-    setContextualInfo(branch, info);
+    setContextualInfo(result, info);
   }
 
   TypeJoinExpr *join = nullptr;
@@ -1488,9 +1506,9 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
         ctx, resultTy, E, AllocationArena::ConstraintSolver);
   }
 
-  // If this is the single expression body of a closure, we need to account
-  // for the fact that the result type may be bound to Void. This is necessary
-  // to correctly handle the following case:
+  // If this is an implied return in a closure, we need to account for the fact
+  // that the result type may be bound to Void. This is necessary to correctly
+  // handle the following case:
   //
   // func foo<T>(_ fn: () -> T) {}
   // foo {
@@ -1514,12 +1532,11 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
   // need to do this with 'return if'. We also don't need to do it for function
   // decls, as we proactively avoid transforming the if/switch into an
   // expression if the result is known to be Void.
-  if (auto *CE = dyn_cast<ClosureExpr>(E->getDeclContext())) {
-    if (CE->hasSingleExpressionBody() && !hasExplicitResult(CE) &&
-        CE->getSingleExpressionBody()->getSemanticsProvidingExpr() == E) {
-      assert(!getAppliedResultBuilderTransform(CE) &&
-             "Should have applied the builder with statement semantics");
-
+  if (impliedResultKind == ImpliedResultKind::ForClosure) {
+    auto *CE = cast<ClosureExpr>(E->getDeclContext());
+    assert(!getAppliedResultBuilderTransform(CE) &&
+           "Should have applied the builder with statement semantics");
+    if (getParentExpr(E) == CE) {
       // We may not have a closure type if we're solving a sub-expression
       // independently for e.g code completion.
       // TODO: This won't be necessary once we stop doing the fallback