[Constraint system] Generalize record of transformed function builders.

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);

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;

lib/Sema/ConstraintSystem.h

@@ -823,9 +823,9 @@ class Solution {
   std::vector<std::pair<ConstraintLocator *, ProtocolConformanceRef>>
       Conformances;
 
-  /// The set of closures that have been transformed by a function builder.
-  llvm::MapVector<ClosureExpr *, AppliedBuilderTransform>
-      builderTransformedClosures;
+  /// The set of functions that have been transformed by a function builder.
+  llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
+      functionBuilderTransformed;
 
   /// Simplify the given type by substituting all occurrences of
   /// type variables for their fixed types.
@@ -1341,9 +1341,9 @@ class ConstraintSystem {
   std::vector<std::pair<ConstraintLocator *, ProtocolConformanceRef>>
       CheckedConformances;
 
-  /// The set of closures that have been transformed by a function builder.
-  std::vector<std::pair<ClosureExpr *, AppliedBuilderTransform>>
-      builderTransformedClosures;
+  /// The set of functions that have been transformed by a function builder.
+  std::vector<std::pair<AnyFunctionRef, AppliedBuilderTransform>>
+      functionBuilderTransformed;
 
 public:
   /// The locators of \c Defaultable constraints whose defaults were used.
@@ -1840,7 +1840,7 @@ class ConstraintSystem {
 
     unsigned numFavoredConstraints;
 
-    unsigned numBuilderTransformedClosures;
+    unsigned numFunctionBuilderTransformed;
 
     /// The length of \c ResolvedOverloads.
     unsigned numResolvedOverloads;