Sema: Record applied property wrappers in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -91,6 +91,8 @@ class SolverTrail {
     FavoredConstraint,
     /// Recorded a result builder transform.
     RecordedResultBuilderTransform,
+    /// Applied a property wrapper.
+    AppliedPropertyWrapper,
   };
 
   /// A change made to the constraint system.
@@ -168,6 +170,7 @@ class SolverTrail {
       ConstraintLocator *Locator;
       PackExpansionType *ExpansionTy;
       PackElementExpr *ElementExpr;
+      Expr *TheExpr;
       AnyFunctionRef AFR;
     };
 
@@ -262,6 +265,9 @@ class SolverTrail {
     /// Create a change that recorded a result builder transform.
     static Change recordedResultBuilderTransform(AnyFunctionRef fn);
 
+    /// Create a change that recorded a result builder transform.
+    static Change appliedPropertyWrapper(Expr *anchor);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1604,7 +1604,7 @@ class Solution {
       resultBuilderTransformed;
 
   /// A map from argument expressions to their applied property wrapper expressions.
-  llvm::MapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>> appliedPropertyWrappers;
+  llvm::DenseMap<ASTNode, SmallVector<AppliedPropertyWrapper, 2>> appliedPropertyWrappers;
 
   /// A mapping from the constraint locators for references to various
   /// names (e.g., member references, normal name references, possible
@@ -2403,7 +2403,7 @@ class ConstraintSystem {
 
 public:
   /// A map from argument expressions to their applied property wrapper expressions.
-  llvm::SmallMapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>, 4>
+  llvm::SmallDenseMap<ASTNode, SmallVector<AppliedPropertyWrapper, 2>, 4>
       appliedPropertyWrappers;
 
   /// The locators of \c Defaultable constraints whose defaults were used.
@@ -2857,9 +2857,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c appliedPropertyWrappers
-    unsigned numAppliedPropertyWrappers;
-
     /// The length of \c ResolvedOverloads.
     unsigned numResolvedOverloads;
 
@@ -5329,6 +5326,13 @@ class ConstraintSystem {
       Type wrapperType, Type paramType, ParamDecl *param, Identifier argLabel,
       ConstraintKind matchKind, ConstraintLocatorBuilder locator);
 
+  /// Used by applyPropertyWrapperToParameter() to update appliedPropertyWrappers
+  /// and record a change in the trail.
+  void applyPropertyWrapper(Expr *anchor, AppliedPropertyWrapper applied);
+
+  /// Undo the above change.
+  void removePropertyWrapper(Expr *anchor);
+
   /// Determine whether given type variable with its set of bindings is viable
   /// to be attempted on the next step of the solver.
   std::optional<BindingSet> determineBestBindings(

lib/Sema/CSGen.cpp

@@ -5046,6 +5046,22 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
   return false;
 }
 
+void ConstraintSystem::applyPropertyWrapper(
+    Expr *anchor, AppliedPropertyWrapper applied) {
+  appliedPropertyWrappers[anchor].push_back(applied);
+
+  if (solverState)
+    recordChange(SolverTrail::Change::appliedPropertyWrapper(anchor));
+}
+
+void ConstraintSystem::removePropertyWrapper(Expr *anchor) {
+  auto found = appliedPropertyWrappers.find(anchor);
+  ASSERT(found != appliedPropertyWrappers.end());
+  auto &wrappers = found->second;
+  ASSERT(!wrappers.empty());
+  wrappers.pop_back();
+}
+
 ConstraintSystem::TypeMatchResult
 ConstraintSystem::applyPropertyWrapperToParameter(
     Type wrapperType, Type paramType, ParamDecl *param, Identifier argLabel,
@@ -5079,13 +5095,13 @@ ConstraintSystem::applyPropertyWrapperToParameter(
       setType(param->getPropertyWrapperProjectionVar(), projectionType);
     }
 
-    appliedPropertyWrappers[anchor].push_back({ wrapperType, PropertyWrapperInitKind::ProjectedValue });
+    applyPropertyWrapper(anchor, { wrapperType, PropertyWrapperInitKind::ProjectedValue });
   } else if (param->hasExternalPropertyWrapper()) {
     Type wrappedValueType = computeWrappedValueType(param, wrapperType);
     addConstraint(matchKind, paramType, wrappedValueType, locator);
     setType(param->getPropertyWrapperWrappedValueVar(), wrappedValueType);
 
-    appliedPropertyWrappers[anchor].push_back({ wrapperType, PropertyWrapperInitKind::WrappedValue });
+    applyPropertyWrapper(anchor, { wrapperType, PropertyWrapperInitKind::WrappedValue });
   } else {
     return getTypeMatchFailure(locator);
   }

lib/Sema/CSSolver.cpp

@@ -415,7 +415,14 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   for (const auto &appliedWrapper : solution.appliedPropertyWrappers) {
-    appliedPropertyWrappers.insert(appliedWrapper);
+    auto found = appliedPropertyWrappers.find(appliedWrapper.first);
+    if (found == appliedPropertyWrappers.end()) {
+      appliedPropertyWrappers.insert(appliedWrapper);
+    } else {
+      auto &existing = found->second;
+      ASSERT(existing.size() <= appliedWrapper.second.size());
+      existing = appliedWrapper.second;
+    }
   }
 
   for (auto &valueConversion : solution.ImplicitValueConversions) {
@@ -681,7 +688,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numAppliedPropertyWrappers = cs.appliedPropertyWrappers.size();
   numResolvedOverloads = cs.ResolvedOverloads.size();
   numInferredClosureTypes = cs.ClosureTypes.size();
   numImpliedResults = cs.ImpliedResults.size();
@@ -733,9 +739,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any applied property wrappers.
-  truncate(cs.appliedPropertyWrappers, numAppliedPropertyWrappers);
-
   // Remove any inferred closure types (e.g. used in result builder body).
   truncate(cs.ClosureTypes, numInferredClosureTypes);
 

lib/Sema/CSTrail.cpp

@@ -270,6 +270,14 @@ SolverTrail::Change::recordedResultBuilderTransform(AnyFunctionRef fn) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::appliedPropertyWrapper(Expr *expr) {
+  Change result;
+  result.Kind = ChangeKind::AppliedPropertyWrapper;
+  result.TheExpr = expr;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -380,6 +388,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedResultBuilderTransform:
     cs.removeResultBuilderTransform(AFR);
     break;
+
+  case ChangeKind::AppliedPropertyWrapper:
+    cs.removePropertyWrapper(TheExpr);
+    break;
   }
 }
 
@@ -588,6 +600,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     simple_display(out, AFR);
     out << ")\n";
     break;
+
+  case ChangeKind::AppliedPropertyWrapper:
+    out << "(applied property wrapper to ";
+    simple_display(out, TheExpr);
+    out << ")\n";
+    break;
   }
 }