Sema: Record SyntacticElementTargetKeys in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

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

include/swift/Sema/CSTrail.h

@@ -34,6 +34,7 @@ class TypeVariableType;
 namespace constraints {
 
 class Constraint;
+struct SyntacticElementTargetKey;
 
 class SolverTrail {
 public:
@@ -102,7 +103,6 @@ class SolverTrail {
         Type DstType;
       } Restriction;
 
-      ConstraintFix *Fix;
 
       struct {
         GenericTypeParamType *GP;
@@ -119,12 +119,19 @@ class SolverTrail {
         Type OldType;
       } KeyPath;
 
-      ConstraintLocator *Locator;
-      PackExpansionType *ExpansionTy;
-      PackElementExpr *ElementExpr;
+      ConstraintFix *TheFix;
+      ConstraintLocator *TheLocator;
+      PackExpansionType *TheExpansion;
+      PackElementExpr *TheElement;
       Expr *TheExpr;
-      AnyFunctionRef AFR;
-      const ClosureExpr *Closure;
+      Stmt *TheStmt;
+      StmtConditionElement *TheCondElt;
+      Pattern *ThePattern;
+      PatternBindingDecl *ThePatternBinding;
+      VarDecl *TheVar;
+      AnyFunctionRef TheRef;
+      ClosureExpr *TheClosure;
+      DeclContext *TheDeclContext;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -205,11 +212,14 @@ class SolverTrail {
     static Change RecordedResultBuilderTransform(AnyFunctionRef fn);
 
     /// Create a change that recorded a closure type.
-    static Change RecordedClosureType(const ClosureExpr *closure);
+    static Change RecordedClosureType(ClosureExpr *closure);
 
     /// Create a change that recorded the contextual type of an AST node.
     static Change RecordedContextualInfo(ASTNode node);
 
+    /// Create a change that recorded a SyntacticElementTarget.
+    static Change RecordedTarget(SyntacticElementTargetKey key);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///
@@ -218,6 +228,9 @@ class SolverTrail {
 
     void dump(llvm::raw_ostream &out, ConstraintSystem &cs,
               unsigned indent = 0) const;
+
+  private:
+    SyntacticElementTargetKey getSyntacticElementTargetKey() const;
   };
 
   SolverTrail(ConstraintSystem &cs);

include/swift/Sema/ConstraintSystem.h

@@ -1203,8 +1203,7 @@ struct CaseLabelItemInfo {
 
 /// Key to the constraint solver's mapping from AST nodes to their corresponding
 /// target.
-class SyntacticElementTargetKey {
-public:
+struct SyntacticElementTargetKey {
   enum class Kind {
     empty,
     tombstone,
@@ -1218,72 +1217,75 @@ class SyntacticElementTargetKey {
     functionRef,
   };
 
-private:
   Kind kind;
 
   union {
-    const StmtConditionElement *stmtCondElement;
+    StmtConditionElement *stmtCondElement;
 
-    const Expr *expr;
+    Expr *expr;
 
-    const Stmt *stmt;
+    Stmt *stmt;
 
-    const Pattern *pattern;
+    Pattern *pattern;
 
     struct PatternBindingEntry {
-      const PatternBindingDecl *patternBinding;
+      PatternBindingDecl *patternBinding;
       unsigned index;
     } patternBindingEntry;
 
-    const VarDecl *varDecl;
+    VarDecl *varDecl;
 
-    const DeclContext *functionRef;
+    DeclContext *functionRef;
   } storage;
 
-public:
   SyntacticElementTargetKey(Kind kind) {
     assert(kind == Kind::empty || kind == Kind::tombstone);
     this->kind = kind;
   }
 
-  SyntacticElementTargetKey(const StmtConditionElement *stmtCondElement) {
+  SyntacticElementTargetKey(StmtConditionElement *stmtCondElement) {
     kind = Kind::stmtCondElement;
     storage.stmtCondElement = stmtCondElement;
   }
 
-  SyntacticElementTargetKey(const Expr *expr) {
+  SyntacticElementTargetKey(Expr *expr) {
     kind = Kind::expr;
     storage.expr = expr;
   }
 
-  SyntacticElementTargetKey(const ClosureExpr *closure) {
+  SyntacticElementTargetKey(ClosureExpr *closure) {
     kind = Kind::closure;
     storage.expr = closure;
   }
 
-  SyntacticElementTargetKey(const Stmt *stmt) {
+  SyntacticElementTargetKey(Stmt *stmt) {
     kind = Kind::stmt;
     storage.stmt = stmt;
   }
 
-  SyntacticElementTargetKey(const Pattern *pattern) {
+  SyntacticElementTargetKey(Pattern *pattern) {
     kind = Kind::pattern;
     storage.pattern = pattern;
   }
 
-  SyntacticElementTargetKey(const PatternBindingDecl *patternBinding,
+  SyntacticElementTargetKey(PatternBindingDecl *patternBinding,
                             unsigned index) {
     kind = Kind::patternBindingEntry;
     storage.patternBindingEntry.patternBinding = patternBinding;
     storage.patternBindingEntry.index = index;
   }
 
-  SyntacticElementTargetKey(const VarDecl *varDecl) {
+  SyntacticElementTargetKey(VarDecl *varDecl) {
     kind = Kind::varDecl;
     storage.varDecl = varDecl;
   }
 
-  SyntacticElementTargetKey(const AnyFunctionRef functionRef) {
+  SyntacticElementTargetKey(DeclContext *dc) {
+    kind = Kind::functionRef;
+    storage.functionRef = dc;
+  }
+
+  SyntacticElementTargetKey(AnyFunctionRef functionRef) {
     kind = Kind::functionRef;
     storage.functionRef = functionRef.getAsDeclContext();
   }
@@ -1574,7 +1576,7 @@ class Solution {
   std::vector<std::pair<ASTNode, ContextualTypeInfo>> contextualTypes;
 
   /// Maps AST nodes to their target.
-  llvm::MapVector<SyntacticElementTargetKey, SyntacticElementTarget> targets;
+  llvm::DenseMap<SyntacticElementTargetKey, SyntacticElementTarget> targets;
 
   /// Maps case label items to information tracked about them as they are
   /// being solved.
@@ -1727,12 +1729,7 @@ class Solution {
   }
 
   std::optional<SyntacticElementTarget>
-  getTargetFor(SyntacticElementTargetKey key) const {
-    auto known = targets.find(key);
-    if (known == targets.end())
-      return std::nullopt;
-    return known->second;
-  }
+  getTargetFor(SyntacticElementTargetKey key) const;
 
   ConstraintLocator *getCalleeLocator(ConstraintLocator *locator,
                                       bool lookThroughApply = true) const;
@@ -2275,7 +2272,7 @@ class ConstraintSystem {
       KeyPaths;
 
   /// Maps AST entries to their targets.
-  llvm::MapVector<SyntacticElementTargetKey, SyntacticElementTarget> targets;
+  llvm::DenseMap<SyntacticElementTargetKey, SyntacticElementTarget> targets;
 
   /// Contextual type information for expressions that are part of this
   /// constraint system. The second type, if valid, contains the type as it
@@ -2860,9 +2857,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c targets.
-    unsigned numTargets;
-
     /// The length of \c caseLabelItems.
     unsigned numCaseLabelItems;
 
@@ -3090,8 +3084,10 @@ class ConstraintSystem {
     bool inserted = ClosureTypes.insert({closure, type}).second;
     ASSERT(inserted);
 
-    if (solverState)
-      recordChange(SolverTrail::Change::RecordedClosureType(closure));
+    if (solverState) {
+      recordChange(SolverTrail::Change::RecordedClosureType(
+        const_cast<ClosureExpr *>(closure)));
+    }
   }
 
   void removeClosureType(const ClosureExpr *closure) {
@@ -3359,18 +3355,12 @@ class ConstraintSystem {
   }
 
   void setTargetFor(SyntacticElementTargetKey key,
-                    SyntacticElementTarget target) {
-    assert(targets.count(key) == 0 && "Already set this target");
-    targets.insert({key, target});
-  }
+                    SyntacticElementTarget target);
+
+  void removeTargetFor(SyntacticElementTargetKey key);
 
   std::optional<SyntacticElementTarget>
-  getTargetFor(SyntacticElementTargetKey key) const {
-    auto known = targets.find(key);
-    if (known == targets.end())
-      return std::nullopt;
-    return known->second;
-  }
+  getTargetFor(SyntacticElementTargetKey key) const;
 
   std::optional<AppliedBuilderTransform>
   getAppliedResultBuilderTransform(AnyFunctionRef fn) const {

lib/Sema/CSGen.cpp

@@ -4995,7 +4995,7 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
   }
 
   Type boolTy = boolDecl->getDeclaredInterfaceType();
-  for (const auto &condElement : condition) {
+  for (auto &condElement : condition) {
     switch (condElement.getKind()) {
     case StmtConditionElement::CK_Availability:
       // Nothing to do here.

lib/Sema/CSSolver.cpp

@@ -218,7 +218,8 @@ Solution ConstraintSystem::finalize() {
     solution.contextualTypes.push_back({entry.first, entry.second.first});
   }
 
-  solution.targets = targets;
+  for (auto &target : targets)
+    solution.targets.insert(target);
 
   for (const auto &item : caseLabelItems)
     solution.caseLabelItems.insert(item);
@@ -692,7 +693,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numTargets = cs.targets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
   numPotentialThrowSites = cs.potentialThrowSites.size();
   numExprPatterns = cs.exprPatterns.size();
@@ -737,9 +737,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any targets.
-  truncate(cs.targets, numTargets);
-
   // Remove any case label item infos.
   truncate(cs.caseLabelItems, numCaseLabelItems);
 

lib/Sema/CSSyntacticElement.cpp

@@ -26,6 +26,36 @@
 using namespace swift;
 using namespace swift::constraints;
 
+void ConstraintSystem::setTargetFor(SyntacticElementTargetKey key,
+                                    SyntacticElementTarget target) {
+  bool inserted = targets.insert({key, target}).second;
+  ASSERT(inserted);
+
+  if (solverState)
+    recordChange(SolverTrail::Change::RecordedTarget(key));
+}
+
+void ConstraintSystem::removeTargetFor(SyntacticElementTargetKey key) {
+  bool erased = targets.erase(key);
+  ASSERT(erased);
+}
+
+std::optional<SyntacticElementTarget>
+ConstraintSystem::getTargetFor(SyntacticElementTargetKey key) const {
+  auto known = targets.find(key);
+  if (known == targets.end())
+    return std::nullopt;
+  return known->second;
+}
+
+std::optional<SyntacticElementTarget>
+Solution::getTargetFor(SyntacticElementTargetKey key) const {
+  auto known = targets.find(key);
+  if (known == targets.end())
+    return std::nullopt;
+  return known->second;
+}
+
 namespace {
 
 // Produce an implicit empty tuple expression.