Sema: Record implied results in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -22,6 +22,10 @@
 #define LOCATOR_CHANGE(Name) CHANGE(Name)
 #endif
 
+#ifndef EXPR_CHANGE
+#define EXPR_CHANGE(Name) CHANGE(Name)
+#endif
+
 #ifndef LAST_CHANGE
 #define LAST_CHANGE(Name)
 #endif
@@ -34,6 +38,9 @@ LOCATOR_CHANGE(RecordedPackExpansionEnvironment)
 LOCATOR_CHANGE(RecordedDefaultedConstraint)
 LOCATOR_CHANGE(ResolvedOverload)
 
+EXPR_CHANGE(AppliedPropertyWrapper)
+EXPR_CHANGE(RecordedImpliedResult)
+
 CHANGE(AddedTypeVariable)
 CHANGE(AddedConstraint)
 CHANGE(RemovedConstraint)
@@ -53,11 +60,11 @@ CHANGE(RecordedKeyPathComponentType)
 CHANGE(DisabledConstraint)
 CHANGE(FavoredConstraint)
 CHANGE(RecordedResultBuilderTransform)
-CHANGE(AppliedPropertyWrapper)
 CHANGE(RecordedClosureType)
 
 LAST_CHANGE(RecordedClosureType)
 
 #undef LOCATOR_CHANGE
+#undef EXPR_CHANGE
 #undef LAST_CHANGE
 #undef CHANGE

include/swift/Sema/CSTrail.h

@@ -130,6 +130,7 @@ class SolverTrail {
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
 
 #define LOCATOR_CHANGE(Name) static Change Name(ConstraintLocator *locator);
+#define EXPR_CHANGE(Name) static Change Name(Expr *expr);
 #include "swift/Sema/CSTrail.def"
 
     /// Create a change that added a type variable.
@@ -203,9 +204,6 @@ 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);
-
     /// Create a change that recorded a closure type.
     static Change RecordedClosureType(const ClosureExpr *closure);
 

include/swift/Sema/ConstraintSystem.h

@@ -1515,7 +1515,7 @@ class Solution {
   /// 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;
+  llvm::DenseMap<Expr *, ImpliedResultKind> ImpliedResults;
 
   /// For locators associated with call expressions, the trailing closure
   /// matching rule and parameter bindings that were applied.
@@ -2233,7 +2233,7 @@ class ConstraintSystem {
   /// 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;
+  llvm::DenseMap<Expr *, ImpliedResultKind> ImpliedResults;
 
   /// This is a *global* list of all result builder bodies that have
   /// been determined to be incorrect by failing constraint generation.
@@ -2860,9 +2860,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c ImpliedResults.
-    unsigned numImpliedResults;
-
     /// The length of \c contextualTypes.
     unsigned numContextualTypes;
 
@@ -3065,16 +3062,24 @@ class ConstraintSystem {
   }
 
   /// Record an implied result for a ReturnStmt or ThenStmt.
-  void recordImpliedResult(const Expr *E, ImpliedResultKind kind) {
-    assert(E);
+  void recordImpliedResult(Expr *E, ImpliedResultKind kind) {
+    ASSERT(E);
     auto inserted = ImpliedResults.insert({E, kind}).second;
-    assert(inserted && "Duplicate implied result?");
-    (void)inserted;
+    ASSERT(inserted && "Duplicate implied result?");
+
+    if (solverState)
+      recordChange(SolverTrail::Change::RecordedImpliedResult(E));
+  }
+
+  /// Undo the above change.
+  void removeImpliedResult(Expr *E) {
+    bool erased = ImpliedResults.erase(E);
+    ASSERT(erased);
   }
 
   /// Whether the given expression is the implied result for either a ReturnStmt
   /// or ThenStmt, and if so, the kind of implied result.
-  std::optional<ImpliedResultKind> isImpliedResult(const Expr *E) const {
+  std::optional<ImpliedResultKind> isImpliedResult(Expr *E) const {
     auto result = ImpliedResults.find(E);
     if (result == ImpliedResults.end())
       return std::nullopt;

lib/Sema/CSSolver.cpp

@@ -323,8 +323,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   // Remember implied results.
-  for (auto impliedResult : solution.ImpliedResults)
-    ImpliedResults.insert(impliedResult);
+  for (auto impliedResult : solution.ImpliedResults) {
+    if (ImpliedResults.count(impliedResult.first) == 0)
+      recordImpliedResult(impliedResult.first, impliedResult.second);
+  }
 
   // Register the solution's opened types.
   for (const auto &opened : solution.OpenedTypes) {
@@ -690,7 +692,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numImpliedResults = cs.ImpliedResults.size();
   numContextualTypes = cs.contextualTypes.size();
   numTargets = cs.targets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
@@ -737,9 +738,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any implied results.
-  truncate(cs.ImpliedResults, numImpliedResults);
-
   // Remove any contextual types.
   truncate(cs.contextualTypes, numContextualTypes);
 

lib/Sema/CSTrail.cpp

@@ -56,6 +56,14 @@ SolverTrail::~SolverTrail() {
     result.Locator = locator; \
     return result; \
   }
+#define EXPR_CHANGE(Name) \
+  SolverTrail::Change \
+  SolverTrail::Change::Name(Expr *expr) { \
+    Change result; \
+    result.Kind = ChangeKind::Name; \
+    result.TheExpr = expr; \
+    return result; \
+  }
 #include "swift/Sema/CSTrail.def"
 
 SolverTrail::Change
@@ -239,14 +247,6 @@ 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;
-}
-
 SolverTrail::Change
 SolverTrail::Change::RecordedClosureType(const ClosureExpr *closure) {
   Change result;
@@ -376,6 +376,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedClosureType:
     cs.removeClosureType(Closure);
     break;
+
+  case ChangeKind::RecordedImpliedResult:
+    cs.removeImpliedResult(TheExpr);
+    break;
   }
 }
 
@@ -395,6 +399,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out); \
     out << ")\n"; \
     break;
+#define EXPR_CHANGE(Name) \
+  case ChangeKind::Name: \
+    out << "(" << #Name << " "; \
+    simple_display(out, TheExpr); \
+    out << ")\n"; \
+    break;
 #include "swift/Sema/CSTrail.def"
 
   case ChangeKind::AddedTypeVariable:
@@ -558,12 +568,6 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     out << ")\n";
     break;
 
-  case ChangeKind::AppliedPropertyWrapper:
-    out << "(AppliedPropertyWrapper ";
-    simple_display(out, TheExpr);
-    out << ")\n";
-    break;
-
   case ChangeKind::RecordedClosureType:
     out << "(RecordedClosureType ";
     simple_display(out, Closure);