ConstraintAnalysis: Properly handle local changes

src/ir/abstract.h

@@ -331,6 +331,24 @@ inline Op negateRelational(Op op) {
   }
 }
 
+inline bool isRelationalSymmetric(Op op) { return op == Eq || op == Ne; }
+
+inline bool isRelationalAntisymmetric(Op op) {
+  switch (op) {
+    case LtS:
+    case LtU:
+    case LeS:
+    case LeU:
+    case GtS:
+    case GtU:
+    case GeS:
+    case GeU:
+      return true;
+    default:
+      return false;
+  }
+}
+
 } // namespace wasm::Abstract
 
 #endif // wasm_ir_abstract_h

src/ir/constraint.cpp

@@ -64,6 +64,11 @@ Result provesPair(const Constraint& a, const Constraint& b) {
     return True;
   }
 
+  // A thing always implies its negation is false.
+  if (a == b.negate()) {
+    return False;
+  }
+
   // Comparisons of two constants.
   auto* aConstant = std::get_if<Literal>(&a.term);
   auto* bConstant = std::get_if<Literal>(&b.term);
@@ -250,6 +255,36 @@ LocalConstraint::parseCondition(Expression* curr) {
   return parse(curr);
 };
 
+void LocalConstraint::flip() {
+  auto other = std::get<Index>(constraint.term);
+  constraint.term = Term{local};
+  local = other;
+  if (Abstract::isRelationalAntisymmetric(constraint.op)) {
+    constraint.op = Abstract::negateRelational(constraint.op);
+  } else {
+    // All we support for now are symmetric and antisymmetric operations.
+    assert(Abstract::isRelationalSymmetric(constraint.op));
+  }
+}
+
+void BasicBlockConstraintMap::set(Index index, const Constraint& c) {
+  assert(!unreachable);
+  eraseStaleRefs(index);
+  map[index].set(c);
+  noteRefs(index, c);
+
+  // If the constraint refers to another local, add it there too.
+  if (std::holds_alternative<Index>(c.term)) {
+    approximateAndInternal(index, c, true);
+  }
+}
+
+void BasicBlockConstraintMap::setProvesNothing(Index index) {
+  assert(!unreachable);
+  eraseStaleRefs(index);
+  map.erase(index);
+}
+
 void BasicBlockConstraintMap::approximateOr(
   const BasicBlockConstraintMap& other) {
   // If one is unreachable, it adds nothing to the other.
@@ -276,9 +311,19 @@ void BasicBlockConstraintMap::approximateOr(
   });
 }
 
-void BasicBlockConstraintMap::approximateAnd(Index index, const Constraint& c) {
+void BasicBlockConstraintMap::approximateAndInternal(Index index,
+                                                     const Constraint& c,
+                                                     bool flip) {
+  Constraint actual = c;
+  if (flip) {
+    LocalConstraint flipped{index, c};
+    flipped.flip();
+    index = flipped.local;
+    actual = flipped.constraint;
+  }
+
   auto combined = get(index);
-  combined.approximateAnd(c);
+  combined.approximateAnd(actual);
 
   if (combined.provesEverything()) {
     // We just proved we are in unreachable code.
@@ -293,6 +338,50 @@ void BasicBlockConstraintMap::approximateAnd(Index index, const Constraint& c) {
 
   // Otherwise, this is an interesting state; set it.
   map[index] = std::move(combined);
+
+  // Add a ref of what we are adding. Note that the approximation above may end
+  // up not actually adding this, or adding only part of this, but it is safe to
+  // always add a ref (at the cost of minor wasted work).
+  noteRefs(index, actual);
+
+  // If this is not the flipped version, and it refers to a local, add the
+  // flipped one too.
+  if (!flip && std::holds_alternative<Index>(actual.term)) {
+    approximateAndInternal(index, actual, true);
+  }
+}
+
+void BasicBlockConstraintMap::noteRefs(Index index, const Constraint& c) {
+  if (auto* i = std::get_if<Index>(&c.term)) {
+    refs[*i].insert(index);
+  }
+}
+
+void BasicBlockConstraintMap::eraseStaleRefs(Index index) {
+  auto iter = refs.find(index);
+  if (iter == refs.end()) {
+    return;
+  }
+
+  auto& refIndexes = iter->second;
+
+  for (auto refIndex : refIndexes) {
+    if (auto iter = map.find(refIndex); iter != map.end()) {
+      auto& refConstraints = iter->second;
+      std::erase_if(refConstraints, [&](const auto& c) {
+        if (auto* i = std::get_if<Index>(&c.term)) {
+          if (*i == index) {
+            return true;
+          }
+        }
+        return false;
+      });
+      if (refConstraints.empty()) {
+        // This became trivial.
+        map.erase(iter);
+      }
+    }
+  }
 }
 
 std::ostream& operator<<(std::ostream& o, const Constraint& c) {

src/ir/constraint.h

@@ -45,7 +45,7 @@ struct Constraint {
 
   bool operator==(const Constraint&) const = default;
 
-  Constraint negate() {
+  Constraint negate() const {
     return Constraint{Abstract::negateRelational(op), term};
   }
 };
@@ -186,6 +186,10 @@ struct LocalConstraint {
   // Parse in a condition context, i.e., where (local.get $x) is the same as
   // $x != 0 (e.g., in an if condition, or a br_on ref).
   static std::optional<LocalConstraint> parseCondition(Expression* curr);
+
+  // Reverse the constraint. The constraint's term must, of course, be another
+  // local.
+  void flip();
 };
 
 // A map of locals and their constraints, representing the state at a basic
@@ -205,6 +209,10 @@ struct LocalConstraint {
 //
 // As a result, the map only contains interesting things, where we can prove
 // something (but not everything).
+//
+// Cross-local constraints (like x == y) are duplicated, that is, they appear in
+// the constraints for both x and y. This makes things simple by having all
+// constraints related to a local in the same place.
 struct BasicBlockConstraintMap {
   // Blocks begin unreachable, like AndedConstraintSet.
   bool unreachable = true;
@@ -215,20 +223,23 @@ struct BasicBlockConstraintMap {
   }
 
   // Apply a constraint to a local.
-  void set(Index index, const Constraint& c) {
-    assert(!unreachable);
-    map[index].set(c);
-  }
+  void set(Index index, const Constraint& c);
 
-  void setProvesNothing(Index index) {
-    assert(!unreachable);
-    map.erase(index);
-  }
+  // Mark a local as unknown and able to prove nothing.
+  void setProvesNothing(Index index);
 
   // Get the constraints for a local.
   AndedConstraintSet get(Index index) const {
+    // We should not be called in unreachable code.
+    assert(!unreachable);
+
     if (auto iter = map.find(index); iter != map.end()) {
-      return iter->second;
+      auto& constraints = iter->second;
+      // If we can prove nothing, we should have removed it from the map.
+      assert(!constraints.provesNothing());
+      // If we can prove everything, we should be entirely unreachable.
+      assert(!constraints.provesEverything());
+      return constraints;
     }
     return AndedConstraintSet::makeProvesNothing();
   }
@@ -239,7 +250,13 @@ struct BasicBlockConstraintMap {
   void approximateOr(const BasicBlockConstraintMap& other);
 
   // Perform an AND as above, on a particular index.
-  void approximateAnd(Index index, const Constraint& c);
+  void approximateAnd(Index index, const Constraint& c) {
+    approximateAndInternal(index, c);
+  }
+
+  // TODO: Add proves() here, which could do things like: if asked x == y, we
+  // can answer False if we see x == c1, y == c2, and the constants c1, c2
+  // differ.
 
   bool operator!=(const BasicBlockConstraintMap& other) {
     return unreachable != other.unreachable || map != other.map;
@@ -250,6 +267,29 @@ struct BasicBlockConstraintMap {
 
 private:
   std::unordered_map<Index, AndedConstraintSet> map;
+
+  // Maps an index to the locals that have constraints referring to it. When a
+  // local is modified, we need to wipe all those constraints, which become
+  // stale.
+  //
+  // It is ok (but unoptimal in efficiency) if we have stale refs here, e.g. due
+  // to approximation removing a constraint. Whenever there is a reference,
+  // however, it must be noted here, so that when things get stale we can remove
+  // them.
+  std::unordered_map<Index, std::unordered_set<Index>> refs;
+
+  // Given a constraint on a local, note refs.
+  void noteRefs(Index index, const Constraint& c);
+
+  // Given an index, erase constraints referring to it.
+  void eraseStaleRefs(Index index);
+
+  // Internal version, with a flag to flip the constraint. Whenever we apply
+  // e.g. x == y, we also apply y == x to y, to maintain the invariant described
+  // above. When flip is true, we flip the constraint and apply it to the other
+  // index (y == x, in this example).
+  void
+  approximateAndInternal(Index index, const Constraint& c, bool flip = false);
 };
 
 std::ostream& operator<<(std::ostream& o, const Constraint& c);

src/passes/ConstraintAnalysis.cpp

@@ -302,9 +302,12 @@ struct ConstraintAnalysis
                           BasicBlockConstraintMap& constraints) {
     if (auto* set = curr->dynCast<LocalSet>()) {
       if (Properties::isSingleConstantExpression(set->value)) {
-        // We know this one constraint.
+        // Apply a constraint to this value.
         auto value = Properties::getLiteral(set->value);
         constraints.set(set->index, Constraint{Abstract::Eq, {value}});
+      } else if (auto* get = set->value->dynCast<LocalGet>()) {
+        // Apply a constraint to this local.
+        constraints.set(set->index, Constraint{Abstract::Eq, {get->index}});
       } else {
         // We know and can prove nothing.
         constraints.setProvesNothing(set->index);