Sema: Record case label items in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -63,8 +63,9 @@ CHANGE(RecordedResultBuilderTransform)
 CHANGE(RecordedClosureType)
 CHANGE(RecordedContextualInfo)
 CHANGE(RecordedTarget)
+CHANGE(RecordedCaseLabelItemInfo)
 
-LAST_CHANGE(RecordedTarget)
+LAST_CHANGE(RecordedCaseLabelItemInfo)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE

include/swift/Sema/CSTrail.h

@@ -132,6 +132,7 @@ class SolverTrail {
       AnyFunctionRef TheRef;
       ClosureExpr *TheClosure;
       DeclContext *TheDeclContext;
+      CaseLabelItem *TheItem;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -220,6 +221,9 @@ class SolverTrail {
     /// Create a change that recorded a SyntacticElementTarget.
     static Change RecordedTarget(SyntacticElementTargetKey key);
 
+    /// Create a change that recorded a SyntacticElementTarget.
+    static Change RecordedCaseLabelItemInfo(CaseLabelItem *item);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1580,8 +1580,7 @@ class Solution {
 
   /// Maps case label items to information tracked about them as they are
   /// being solved.
-  llvm::MapVector<const CaseLabelItem *, CaseLabelItemInfo>
-      caseLabelItems;
+  llvm::DenseMap<const CaseLabelItem *, CaseLabelItemInfo> caseLabelItems;
 
   /// Maps catch nodes to the set of potential throw sites that will be caught
   /// at that location.
@@ -2281,8 +2280,7 @@ class ConstraintSystem {
   llvm::DenseMap<ASTNode, std::pair<ContextualTypeInfo, Type>> contextualTypes;
 
   /// Information about each case label item tracked by the constraint system.
-  llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
-      caseLabelItems;
+  llvm::SmallDenseMap<const CaseLabelItem *, CaseLabelItemInfo, 4> caseLabelItems;
 
   /// Keep track of all of the potential throw sites.
   /// FIXME: This data structure should be replaced with something that
@@ -2857,9 +2855,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c caseLabelItems.
-    unsigned numCaseLabelItems;
-
     /// The length of \c potentialThrowSites.
     unsigned numPotentialThrowSites;
 
@@ -3392,9 +3387,19 @@ class ConstraintSystem {
   }
 
   void setCaseLabelItemInfo(const CaseLabelItem *item, CaseLabelItemInfo info) {
-    assert(item != nullptr);
-    assert(caseLabelItems.count(item) == 0);
-    caseLabelItems[item] = info;
+    ASSERT(item);
+    bool inserted = caseLabelItems.insert({item, info}).second;
+    ASSERT(inserted);
+
+    if (solverState) {
+      recordChange(SolverTrail::Change::RecordedCaseLabelItemInfo(
+        const_cast<CaseLabelItem *>(item)));
+    }
+  }
+
+  void removeCaseLabelItemInfo(const CaseLabelItem *item) {
+    bool erased = caseLabelItems.erase(item);
+    ASSERT(erased);
   }
 
   /// Record a given ExprPattern as the parent of its sub-expression.

lib/Sema/CSSolver.cpp

@@ -693,7 +693,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numCaseLabelItems = cs.caseLabelItems.size();
   numPotentialThrowSites = cs.potentialThrowSites.size();
   numExprPatterns = cs.exprPatterns.size();
   numIsolatedParams = cs.isolatedParams.size();
@@ -737,9 +736,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any case label item infos.
-  truncate(cs.caseLabelItems, numCaseLabelItems);
-
   // Remove any potential throw sites.
   truncate(cs.potentialThrowSites, numPotentialThrowSites);
 

lib/Sema/CSTrail.cpp

@@ -303,6 +303,14 @@ SolverTrail::Change::RecordedTarget(SyntacticElementTargetKey key) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RecordedCaseLabelItemInfo(CaseLabelItem *item) {
+  Change result;
+  result.Kind = ChangeKind::RecordedCaseLabelItemInfo;
+  result.TheItem = item;
+  return result;
+}
+
 SyntacticElementTargetKey
 SolverTrail::Change::getSyntacticElementTargetKey() const {
   ASSERT(Kind == ChangeKind::RecordedTarget);
@@ -465,6 +473,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedTarget:
     cs.removeTargetFor(getSyntacticElementTargetKey());
     break;
+
+  case ChangeKind::RecordedCaseLabelItemInfo:
+    cs.removeCaseLabelItemInfo(TheItem);
+    break;
   }
 }
 
@@ -669,6 +681,11 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     getSyntacticElementTargetKey().dump(out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedCaseLabelItemInfo:
+    // FIXME: Print something here
+    out << "(RecordedCaseLabelItemInfo)\n";
+    break;
   }
 }