Sema: Record result builder transforms in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -17,6 +17,9 @@
 #ifndef SWIFT_SEMA_CSTRAIL_H
 #define SWIFT_SEMA_CSTRAIL_H
 
+#include "swift/AST/AnyFunctionRef.h"
+#include "swift/AST/Type.h"
+#include "swift/AST/Types.h"
 #include <vector>
 
 namespace llvm {
@@ -86,6 +89,8 @@ class SolverTrail {
     DisabledConstraint,
     /// Favored a constraint.
     FavoredConstraint,
+    /// Recorded a result builder transform.
+    RecordedResultBuilderTransform,
   };
 
   /// A change made to the constraint system.
@@ -163,6 +168,7 @@ class SolverTrail {
       ConstraintLocator *Locator;
       PackExpansionType *ExpansionTy;
       PackElementExpr *ElementExpr;
+      AnyFunctionRef AFR;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -253,6 +259,9 @@ class SolverTrail {
     /// Create a change that favored a constraint.
     static Change favoredConstraint(Constraint *constraint);
 
+    /// Create a change that recorded a result builder transform.
+    static Change recordedResultBuilderTransform(AnyFunctionRef fn);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2857,8 +2857,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    unsigned numResultBuilderTransformed;
-
     /// The length of \c appliedPropertyWrappers
     unsigned numAppliedPropertyWrappers;
 
@@ -5317,6 +5315,14 @@ class ConstraintSystem {
                      ConstraintKind bodyResultConstraintKind,
                      Type contextualType, ConstraintLocatorBuilder locator);
 
+  /// Used by matchResultBuilder() to update resultBuilderTransformed and record
+  /// a change in the trail.
+  void recordResultBuilderTransform(AnyFunctionRef fn,
+                                    AppliedBuilderTransform transformInfo);
+
+  /// Undo the above change.
+  void removeResultBuilderTransform(AnyFunctionRef fn);
+
   /// Matches a wrapped or projected value parameter type to its backing
   /// property wrapper type by applying the property wrapper.
   TypeMatchResult applyPropertyWrapperToParameter(

lib/Sema/BuilderTransform.cpp

@@ -1211,22 +1211,30 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
   transformInfo.transformedBody = transformedBody->second;
 
   // Record the transformation.
-  assert(
-      std::find_if(
-          resultBuilderTransformed.begin(), resultBuilderTransformed.end(),
-          [&](const std::pair<AnyFunctionRef, AppliedBuilderTransform> &elt) {
-            return elt.first == fn;
-          }) == resultBuilderTransformed.end() &&
-      "already transformed this body along this path!?!");
-  resultBuilderTransformed.insert(
-      std::make_pair(fn, std::move(transformInfo)));
+  recordResultBuilderTransform(fn, std::move(transformInfo));
 
   if (generateConstraints(fn, transformInfo.transformedBody))
     return getTypeMatchFailure(locator);
 
   return getTypeMatchSuccess();
 }
 
+void ConstraintSystem::recordResultBuilderTransform(AnyFunctionRef fn,
+                                          AppliedBuilderTransform transformInfo) {
+  bool inserted = resultBuilderTransformed.insert(
+      std::make_pair(fn, std::move(transformInfo))).second;
+  ASSERT(inserted);
+
+  if (solverState)
+    recordChange(SolverTrail::Change::recordedResultBuilderTransform(fn));
+}
+
+/// Undo the above change.
+void ConstraintSystem::removeResultBuilderTransform(AnyFunctionRef fn) {
+  bool erased = resultBuilderTransformed.erase(fn);
+  ASSERT(erased);
+}
+
 namespace {
 class ReturnStmtFinder : public ASTWalker {
   std::vector<ReturnStmt *> ReturnStmts;

lib/Sema/CSSolver.cpp

@@ -410,7 +410,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   for (const auto &transformed : solution.resultBuilderTransformed) {
-    resultBuilderTransformed.insert(transformed);
+    if (resultBuilderTransformed.count(transformed.first) == 0)
+      recordResultBuilderTransform(transformed.first, transformed.second);
   }
 
   for (const auto &appliedWrapper : solution.appliedPropertyWrappers) {
@@ -680,7 +681,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numResultBuilderTransformed = cs.resultBuilderTransformed.size();
   numAppliedPropertyWrappers = cs.appliedPropertyWrappers.size();
   numResolvedOverloads = cs.ResolvedOverloads.size();
   numInferredClosureTypes = cs.ClosureTypes.size();
@@ -733,9 +733,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  /// Remove any builder transformed closures.
-  truncate(cs.resultBuilderTransformed, numResultBuilderTransformed);
-
   // Remove any applied property wrappers.
   truncate(cs.appliedPropertyWrappers, numAppliedPropertyWrappers);
 

lib/Sema/CSTrail.cpp

@@ -262,6 +262,14 @@ SolverTrail::Change::favoredConstraint(Constraint *constraint) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedResultBuilderTransform(AnyFunctionRef fn) {
+  Change result;
+  result.Kind = ChangeKind::RecordedResultBuilderTransform;
+  result.AFR = fn;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -368,6 +376,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     if (TheConstraint.Constraint->isFavored())
       TheConstraint.Constraint->setFavored(false);
     break;
+
+  case ChangeKind::RecordedResultBuilderTransform:
+    cs.removeResultBuilderTransform(AFR);
+    break;
   }
 }
 
@@ -570,6 +582,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
                                     indent + 2);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedResultBuilderTransform:
+    out << "(recorded result builder transform ";
+    simple_display(out, AFR);
+    out << ")\n";
+    break;
   }
 }