ConstraintAnalysis: Handle cases with extra constraints

src/ir/abstract.h

@@ -304,6 +304,7 @@ inline BinaryOp getBinary(Type type, Op op) {
   WASM_UNREACHABLE("invalid type");
 }
 
+// Logical negation, e.g. !(x < 10)  == x >= 10
 inline Op negateRelational(Op op) {
   switch (op) {
     case Eq:
@@ -331,6 +332,35 @@ inline Op negateRelational(Op op) {
   }
 }
 
+// Side flipping, e.g.  x < 10 flips to 10 > x (while still saying the same
+// thing, not negated).
+inline Op flipRelational(Op op) {
+  switch (op) {
+    case Eq:
+      return Eq;
+    case Ne:
+      return Ne;
+    case LtS:
+      return GtS;
+    case LtU:
+      return GtU;
+    case LeS:
+      return GeS;
+    case LeU:
+      return GeU;
+    case GtS:
+      return LtS;
+    case GtU:
+      return LtU;
+    case GeS:
+      return LeS;
+    case GeU:
+      return LeU;
+    default:
+      WASM_UNREACHABLE("invalid relational");
+  }
+}
+
 inline bool isRelationalSymmetric(Op op) { return op == Eq || op == Ne; }
 
 inline bool isRelationalAntisymmetric(Op op) {

src/ir/constraint.cpp

@@ -498,7 +498,7 @@ void LocalConstraint::flip() {
   constraint.term = Term{local};
   local = other;
   if (Abstract::isRelationalAntisymmetric(constraint.op)) {
-    constraint.op = Abstract::negateRelational(constraint.op);
+    constraint.op = Abstract::flipRelational(constraint.op);
   } else {
     // All we support for now are symmetric and antisymmetric operations.
     assert(Abstract::isRelationalSymmetric(constraint.op));
@@ -552,60 +552,62 @@ void BasicBlockConstraintMap::set(Index index, Expression* value) {
     auto old = get(y);
 
     // Iterate over the old constraints and increment each one.
-    auto success = true;
-    for (auto& c : old) {
+    for (auto iter = old.begin(); iter != old.end();) {
+      auto& c = *iter;
       auto* N = std::get_if<Literal>(&c.term);
       if (!N) {
-        // A non-constant term, which we don't know how to increment.
-        success = false;
-        break;
+        // A non-constant term, which we don't know how to increment. Simply
+        // remove it: we are losing proving power here, but doing so is never
+        // invalid.
+        iter = old.erase(iter);
+        continue;
       }
 
       switch (c.op) {
         // x == N, x++  =>  x == N+1.
         case Eq:
           *N = N->add(Literal::makeFromInt32(1, N->type));
-          continue;
+          break;
         // x >= N, x++  =>  x > N
         case GeS:
           c.op = GtS;
-          continue;
+          break;
         case GeU:
           c.op = GtU;
-          continue;
+          break;
         // x < N, x++  =>  x <= N
         case LtS:
           c.op = LeS;
-          continue;
+          break;
         case LtU:
           c.op = LeU;
-          continue;
+          break;
         // x <= N, x++ => x <= N+1 if no overflow
         case LeS:
           if (N->isSignedMax()) {
-            success = false;
-            break;
+            iter = old.erase(iter);
+            continue;
           }
           *N = N->add(Literal::makeFromInt32(1, N->type));
-          continue;
+          break;
         case LeU:
           if (N->isUnsignedMax()) {
-            success = false;
-            break;
+            iter = old.erase(iter);
+            continue;
           }
           *N = N->add(Literal::makeFromInt32(1, N->type));
-          continue;
+          break;
         default:
           // Something we don't recognize.
-          success = false;
-          break;
+          iter = old.erase(iter);
+          continue;
       }
-    }
 
-    if (success) {
-      set(index, old);
-      return;
+      ++iter;
     }
+
+    set(index, old);
+    return;
   }
 
   // We know and can prove nothing.
@@ -666,12 +668,20 @@ void BasicBlockConstraintMap::approximateAndInternal(Index index,
     actual = flipped.constraint;
   }
 
-  // Never add constraints to ourselves (x == x, etc., which can happen due to
-  // copying/flipping).
   if (auto* other = std::get_if<Index>(&actual.term)) {
+    // Never add constraints to ourselves (x == x, etc., which can happen due to
+    // copying/flipping).
     if (*other == index) {
       return;
     }
+
+    // If we are applying a constraint to another local, and we know that
+    // local's value, propagate it. That is, if x == 42, then if we try to apply
+    // y < x we instead apply y < 42, which is better.
+    auto otherConstraints = get(*other);
+    if (auto lit = otherConstraints.getLiteral()) {
+      actual.term = Term(*lit);
+    }
   }
 
   // Refer to the constraints for this index. If this is the first access of

src/ir/constraint.h

@@ -194,6 +194,18 @@ struct AndedConstraintSet : inplace_vector<Constraint, MaxConstraints> {
     setProvesNothing();
     push_back(c);
   }
+
+  // If the set of constraints shows us as equal to a literal, return it.
+  std::optional<Literal> getLiteral() const {
+    for (auto& c : *this) {
+      if (c.op == Abstract::Eq) {
+        if (auto* cc = std::get_if<Literal>(&c.term)) {
+          return *cc;
+        }
+      }
+    }
+    return {};
+  }
 };
 
 // A local plus a constraint on it.

src/passes/ConstraintAnalysis.cpp

@@ -592,14 +592,9 @@ struct ConstraintAnalysis
       return false;
     }
 
-    auto localConstraints = constraints.get(branch.local);
     // Handle the case of simple equality of the local to a constant.
-    // TODO: Handle more constraints here as well, and non-constant ones.
-    if (localConstraints.size() != 1 ||
-        localConstraints[0].op != Abstract::Eq) {
-      return false;
-    }
-    auto* N = std::get_if<Literal>(&localConstraints[0].term);
+    // TODO: Handle non-constant ones.
+    auto N = constraints.get(branch.local).getLiteral();
     if (!N) {
       return false;
     }

test/gtest/constraint.cpp

@@ -631,4 +631,32 @@ TEST(ConstraintTest, TestIncrement) {
   EXPECT_EQ(map.get(0),
             (AndedConstraintSet{{GtS, {Literal(int32_t(10))}},
                                 {LeS, {Literal(int32_t(20))}}}));
+
+  // $0 >= 10 && $0 <= max_signed, $0++  =>  $0 > 10 (overflowing constraint
+  // removed)
+  map.set(0, {GeS, {Literal(int32_t(10))}});
+  map.approximateAnd(0, {LeS, {Literal::makeSignedMax(Type::i32)}});
+  map.set(0, &add);
+  EXPECT_EQ(map.get(0), (AndedConstraintSet{{GtS, {Literal(int32_t(10))}}}));
+
+  // $0 >= 10 && $0 == $2, $0++  =>  $0 > 10 (non-constant term removed)
+  map.set(0, {GeS, {Literal(int32_t(10))}});
+  map.approximateAnd(0, {Eq, {Index(2)}});
+  map.set(0, &add);
+  EXPECT_EQ(map.get(0), (AndedConstraintSet{{GtS, {Literal(int32_t(10))}}}));
+}
+
+TEST(ConstraintTest, TestEqConstraints) {
+  BasicBlockConstraintMap map;
+  map.setReachable();
+
+  // $0 == 42
+  map.set(0, {Eq, {Literal(int32_t(42))}});
+
+  // $0 < $1
+  map.approximateAnd(0, {LtS, {Index(int32_t(1))}});
+
+  // $1 has $1 > 42: we constant-propagated the value of $0. This is better than
+  // having $1 > $0 and needing to look $0 up.
+  check(map.get(1), {GtS, {Literal(int32_t(42))}});
 }

test/lit/passes/constraint-analysis-loops.wast

@@ -682,4 +682,102 @@
       )
     )
   )
+
+  ;; CHECK:      (func $extra-constraint (type $0)
+  ;; CHECK-NEXT:  (local $x i32)
+  ;; CHECK-NEXT:  (block $out
+  ;; CHECK-NEXT:   (loop $loop
+  ;; CHECK-NEXT:    (local.set $x
+  ;; CHECK-NEXT:     (i32.add
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (i32.gt_u
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:      (i32.const 100)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (br $out)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (i32.eq
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:      (i32.const 42)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (unreachable)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (br $loop)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $extra-constraint
+    ;; As in the last testcase, but with extra code below.
+    (local $x i32)
+    (block $out
+      (loop $loop
+        (local.set $x
+          (i32.add
+            (local.get $x)
+            (i32.const 1)
+          )
+        )
+        (if
+          (i32.gt_u
+            (local.get $x)
+            (i32.const 100)
+          )
+          (then
+            (br $out)
+          )
+        )
+
+        ;; An extra constraint gets added to the local here. We should still be
+        ;; able to optimize the things below.
+        (if
+          (i32.eq
+            (local.get $x)
+            (i32.const 42)
+          )
+          (then
+            (unreachable)
+          )
+        )
+
+        (drop
+          (i32.gt_u
+            (local.get $x)
+            (i32.const 0)
+          )
+        )
+        (drop
+          (i32.le_u
+            (local.get $x)
+            (i32.const 100)
+          )
+        )
+        ;; And we also optimize that extra constraint: this is true.
+        (drop
+          (i32.ne
+            (local.get $x)
+            (i32.const 42)
+          )
+        )
+        (br $loop)
+      )
+    )
+  )
 )