Merge pull request swiftlang#29026 from DougGregor/generalize-solution-apply

Generalize solution application

Author: DougGregor

include/swift/AST/AnyFunctionRef.h

@@ -208,6 +208,14 @@ class AnyFunctionRef {
     llvm_unreachable("unexpected AnyFunctionRef representation");
   }
 
+  friend bool operator==(AnyFunctionRef lhs, AnyFunctionRef rhs) {
+     return lhs.TheFunction == rhs.TheFunction;
+   }
+
+   friend bool operator!=(AnyFunctionRef lhs, AnyFunctionRef rhs) {
+     return lhs.TheFunction != rhs.TheFunction;
+   }
+
 private:
   ArrayRef<AnyFunctionType::Yield>
   getYieldResultsImpl(SmallVectorImpl<AnyFunctionType::Yield> &buffer,

lib/Sema/BuilderTransform.cpp

@@ -604,13 +604,13 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
 
   // Record the transformation.
   assert(std::find_if(
-      builderTransformedClosures.begin(),
-      builderTransformedClosures.end(),
-      [&](const std::pair<ClosureExpr *, AppliedBuilderTransform> &elt) {
+      functionBuilderTransformed.begin(),
+      functionBuilderTransformed.end(),
+      [&](const std::pair<AnyFunctionRef, AppliedBuilderTransform> &elt) {
         return elt.first == closure;
-      }) == builderTransformedClosures.end() &&
+      }) == functionBuilderTransformed.end() &&
          "already transformed this closure along this path!?!");
-  builderTransformedClosures.push_back(
+  functionBuilderTransformed.push_back(
       std::make_pair(closure,
                      AppliedBuilderTransform{builderType, singleExpr}));
 

lib/Sema/CSApply.cpp

@@ -7041,9 +7041,8 @@ namespace {
         // If this closure had a function builder applied, rewrite it to a
         // closure with a single expression body containing the builder
         // invocations.
-        auto builder =
-            Rewriter.solution.builderTransformedClosures.find(closure);
-        if (builder != Rewriter.solution.builderTransformedClosures.end()) {
+        auto builder = Rewriter.solution.functionBuilderTransformed.find(closure);
+        if (builder != Rewriter.solution.functionBuilderTransformed.end()) {
           auto singleExpr = builder->second.singleExpr;
           auto returnStmt = new (ctx) ReturnStmt(
              singleExpr->getStartLoc(), singleExpr, /*implicit=*/true);
@@ -7206,10 +7205,9 @@ bool ConstraintSystem::applySolutionFixes(const Solution &solution) {
 
 /// Apply a given solution to the expression, producing a fully
 /// type-checked expression.
-Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
-                                      Type convertType,
-                                      bool discardedExpr,
-                                      bool performingDiagnostics) {
+llvm::PointerUnion<Expr *, Stmt *> ConstraintSystem::applySolutionImpl(
+    Solution &solution, SolutionApplicationTarget target, Type convertType,
+    bool discardedExpr, bool performingDiagnostics) {
   // If any fixes needed to be applied to arrive at this solution, resolve
   // them to specific expressions.
   if (!solution.Fixes.empty()) {
@@ -7228,18 +7226,32 @@ Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
 
       // If we didn't manage to diagnose anything well, so fall back to
       // diagnosing mining the system to construct a reasonable error message.
-      diagnoseFailureForExpr(expr);
+      diagnoseFailureFor(target);
       return nullptr;
     }
   }
 
   ExprRewriter rewriter(*this, solution, shouldSuppressDiagnostics());
   ExprWalker walker(rewriter);
 
-  // Apply the solution to the expression.
-  auto result = expr->walk(walker);
-  if (!result)
-    return nullptr;
+  // Apply the solution to the target.
+  llvm::PointerUnion<Expr *, Stmt *> result;
+  if (auto expr = target.getAsExpr()) {
+    result = expr->walk(walker);
+  } else {
+    // FIXME: Implement this!
+    llvm_unreachable("Not yet implemented");
+
+#if false
+    auto fn = *target.getAsFunction();
+    auto transform = rewriter.getAppliedBuilderTransform(fn);
+    assert(transform && "Not applying builder transform?");
+    result = walker.applyFunctionBuilderBodyTransform(fn, transform);
+#endif
+  }
+
+  if (result.isNull())
+    return result;
 
   // If we're re-typechecking an expression for diagnostics, don't
   // visit closures that have non-single expression bodies.
@@ -7261,28 +7273,36 @@ Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
       return nullptr;
   }
 
-  // We are supposed to use contextual type only if it is present and
-  // this expression doesn't represent the implicit return of the single
-  // expression function which got deduced to be `Never`.
-  auto shouldCoerceToContextualType = [&]() {
-    return convertType && !(getType(result)->isUninhabited() &&
-                            getContextualTypePurpose() == CTP_ReturnSingleExpr);
-  };
+  if (auto resultExpr = result.dyn_cast<Expr *>()) {
+    Expr *expr = target.getAsExpr();
+    assert(expr && "Can't have expression result without expression target");
+    // We are supposed to use contextual type only if it is present and
+    // this expression doesn't represent the implicit return of the single
+    // expression function which got deduced to be `Never`.
+    auto shouldCoerceToContextualType = [&]() {
+      return convertType &&
+          !(getType(resultExpr)->isUninhabited() &&
+            getContextualTypePurpose() == CTP_ReturnSingleExpr);
+    };
 
-  // If we're supposed to convert the expression to some particular type,
-  // do so now.
-  if (shouldCoerceToContextualType()) {
-    result = rewriter.coerceToType(result, convertType,
-                                   getConstraintLocator(expr));
-    if (!result)
-      return nullptr;
-  } else if (getType(result)->hasLValueType() && !discardedExpr) {
-    // We referenced an lvalue. Load it.
-    result = rewriter.coerceToType(result, getType(result)->getRValueType(),
-                                   getConstraintLocator(expr));
+    // If we're supposed to convert the expression to some particular type,
+    // do so now.
+    if (shouldCoerceToContextualType()) {
+      result = rewriter.coerceToType(resultExpr, convertType,
+                                     getConstraintLocator(expr));
+      if (!result)
+        return nullptr;
+    } else if (getType(resultExpr)->hasLValueType() && !discardedExpr) {
+      // We referenced an lvalue. Load it.
+      result = rewriter.coerceToType(resultExpr,
+                                     getType(resultExpr)->getRValueType(),
+                                     getConstraintLocator(expr));
+    }
+
+    if (resultExpr)
+      solution.setExprTypes(resultExpr);
   }
 
-  solution.setExprTypes(result);
   rewriter.finalize();
 
   return result;
@@ -7431,3 +7451,14 @@ ArrayRef<Solution> SolutionResult::getAmbiguousSolutions() const {
   assert(getKind() == Ambiguous);
   return makeArrayRef(solutions, numSolutions);
 }
+
+llvm::PointerUnion<Expr *, Stmt *> SolutionApplicationTarget::walk(
+    ASTWalker &walker) {
+  switch (kind) {
+  case Kind::expression:
+    return getAsExpr()->walk(walker);
+
+  case Kind::function:
+    return getAsFunction()->getBody()->walk(walker);
+  }
+}

lib/Sema/CSDiag.cpp

@@ -2420,30 +2420,32 @@ bool FailureDiagnosis::diagnoseExprFailure() {
 ///
 /// This is guaranteed to always emit an error message.
 ///
-void ConstraintSystem::diagnoseFailureForExpr(Expr *expr) {
+void ConstraintSystem::diagnoseFailureFor(SolutionApplicationTarget target) {
   setPhase(ConstraintSystemPhase::Diagnostics);
 
   SWIFT_DEFER { setPhase(ConstraintSystemPhase::Finalization); };
 
-  // Look through RebindSelfInConstructorExpr to avoid weird Sema issues.
-  if (auto *RB = dyn_cast<RebindSelfInConstructorExpr>(expr))
-    expr = RB->getSubExpr();
+  if (auto expr = target.getAsExpr()) {
+    // Look through RebindSelfInConstructorExpr to avoid weird Sema issues.
+    if (auto *RB = dyn_cast<RebindSelfInConstructorExpr>(expr))
+      expr = RB->getSubExpr();
 
-  FailureDiagnosis diagnosis(expr, *this);
+    FailureDiagnosis diagnosis(expr, *this);
 
-  // Now, attempt to diagnose the failure from the info we've collected.
-  if (diagnosis.diagnoseExprFailure())
-    return;
+    // Now, attempt to diagnose the failure from the info we've collected.
+    if (diagnosis.diagnoseExprFailure())
+      return;
 
-  // If this is a contextual conversion problem, dig out some information.
-  if (diagnosis.diagnoseContextualConversionError(expr, getContextualType(),
-                                                  getContextualTypePurpose()))
-    return;
+    // If this is a contextual conversion problem, dig out some information.
+    if (diagnosis.diagnoseContextualConversionError(expr, getContextualType(),
+                                                    getContextualTypePurpose()))
+      return;
 
-  // If no one could find a problem with this expression or constraint system,
-  // then it must be well-formed... but is ambiguous.  Handle this by diagnostic
-  // various cases that come up.
-  diagnosis.diagnoseAmbiguity(expr);
+    // If no one could find a problem with this expression or constraint system,
+    // then it must be well-formed... but is ambiguous.  Handle this by diagnostic
+    // various cases that come up.
+    diagnosis.diagnoseAmbiguity(expr);
+  }
 }
 
 std::pair<Type, ContextualTypePurpose>

lib/Sema/CSSolver.cpp

@@ -171,13 +171,13 @@ Solution ConstraintSystem::finalize() {
   for (auto &e : CheckedConformances)
     solution.Conformances.push_back({e.first, e.second});
 
-  for (const auto &transformed : builderTransformedClosures) {
+  for (const auto &transformed : functionBuilderTransformed) {
     auto known =
-        solution.builderTransformedClosures.find(transformed.first);
-    if (known != solution.builderTransformedClosures.end()) {
+        solution.functionBuilderTransformed.find(transformed.first);
+    if (known != solution.functionBuilderTransformed.end()) {
       assert(known->second.singleExpr == transformed.second.singleExpr);
     }
-    solution.builderTransformedClosures.insert(transformed);
+    solution.functionBuilderTransformed.insert(transformed);
   }
 
   return solution;
@@ -241,8 +241,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   for (auto &conformance : solution.Conformances)
     CheckedConformances.push_back(conformance);
 
-  for (const auto &transformed : solution.builderTransformedClosures) {
-    builderTransformedClosures.push_back(transformed);
+  for (const auto &transformed : solution.functionBuilderTransformed) {
+    functionBuilderTransformed.push_back(transformed);
   }
 
   // Register any fixes produced along this path.
@@ -436,7 +436,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numCheckedConformances = cs.CheckedConformances.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
   numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
-  numBuilderTransformedClosures = cs.builderTransformedClosures.size();
+  numFunctionBuilderTransformed = cs.functionBuilderTransformed.size();
   numResolvedOverloads = cs.ResolvedOverloads.size();
 
   PreviousScore = cs.CurrentScore;
@@ -503,7 +503,7 @@ ConstraintSystem::SolverScope::~SolverScope() {
   truncate(cs.CheckedConformances, numCheckedConformances);
 
   /// Remove any builder transformed closures.
-  truncate(cs.builderTransformedClosures, numBuilderTransformedClosures);
+  truncate(cs.functionBuilderTransformed, numFunctionBuilderTransformed);
 
   // Reset the previous score.
   cs.CurrentScore = PreviousScore;
@@ -1127,7 +1127,7 @@ bool ConstraintSystem::solve(Expr *&expr,
       return true;
 
     case SolutionResult::Kind::UndiagnosedError:
-      diagnoseFailureForExpr(expr);
+      diagnoseFailureFor(expr);
       salvagedResult.markAsDiagnosed();
       return true;