[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.

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 @@ public:
   /// 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 @@ public:
     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 @@ private:
   /// 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 @@ public:
     /// 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 @@ public:
     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");