ConstraintAnalysis: Parse AND operations

src/ir/constraint.cpp

@@ -734,6 +734,8 @@ bool AndedConstraintSet::approximateOr(const AndedConstraintSet& other) {
 }
 
 std::optional<LocalConstraint> LocalConstraint::parse(Expression* curr) {
+  using namespace Match;
+
   auto parseEqZArgument =
     [&](Expression* value) -> std::optional<LocalConstraint> {
     if (auto* get = value->dynCast<LocalGet>()) {
@@ -745,10 +747,18 @@ std::optional<LocalConstraint> LocalConstraint::parse(Expression* curr) {
     return {};
   };
 
-  if (auto* unary = curr->dynCast<Unary>()) {
-    if (Abstract::getUnary(unary->value->type, Abstract::EqZ) == unary->op) {
-      return parseEqZArgument(unary->value);
+  if (auto* u = curr->dynCast<Unary>()) {
+    if (Abstract::getUnary(u->value->type, Abstract::EqZ) == u->op) {
+      // EqZ of EqZ means a check that the value is *not* zero.
+      LocalGet* get;
+      if (matches(u->value, unary(Abstract::EqZ, Match::local(&get)))) {
+        auto value = Literal::makeZero(get->type);
+        return LocalConstraint{get->index, Constraint{Abstract::Ne, {value}}};
+      }
+
+      return parseEqZArgument(u->value);
     }
+
     return {};
   }
 
@@ -807,17 +817,107 @@ std::optional<LocalConstraint> LocalConstraint::parse(Expression* curr) {
   return {};
 }
 
-std::optional<LocalConstraint>
-LocalConstraint::parseCondition(Expression* curr) {
+ParsedAndedConstraints ParsedAndedConstraints::parse(Expression* curr) {
+  using namespace Match;
+
+  // The final return value.
+  ParsedAndedConstraints ret;
+
+  // Starting from |curr|, parse and recurse into sub-trees: when we see an AND,
+  // we push both children as further work.
+  SmallVector<Expression*, 4> work;
+  work.push_back(curr);
+  while (!work.empty()) {
+    auto* curr = work.back();
+    work.pop_back();
+
+    auto parsed = LocalConstraint::parse(curr);
+    if (parsed) {
+      ret.push_back(*parsed);
+      continue;
+    }
+
+    Binary* b;
+    if (matches(curr, binary(&b, Abstract::And, any(), any()))) {
+      // An AND can be recursively processed: both sides must be true.
+      work.push_back(b->left);
+      work.push_back(b->right);
+      continue;
+    }
+    // TODO: support OR
+
+    // We failed to parse this.
+    ret.hasUnknown = true;
+  }
+
+  return ret;
+}
+
+ParsedAndedConstraints
+ParsedAndedConstraints::parseCondition(Expression* curr) {
   // A get by itself is a check for not being null.
   if (auto* get = curr->dynCast<LocalGet>()) {
     auto value = Literal::makeZero(get->type);
-    return LocalConstraint{get->index, Constraint{Abstract::Ne, {value}}};
+    return {LocalConstraint{get->index, Constraint{Abstract::Ne, {value}}}};
   }
 
   // Otherwise, parse normally.
   return parse(curr);
-};
+}
+
+void ParsedAndedConstraints::negate() {
+  if (empty()) {
+    return;
+  }
+
+  if (hasUnknown) {
+    // This includes things we don't know about, and don't know how to negate.
+    clear();
+    return;
+  }
+
+  // The input is a list of constraints all applying at once, A & B & C. The
+  // negation is !A | !B | !C, but we cannot express a general OR like that,
+  // except in the simple case where they all talk about the same local: then
+  // we can at least approximateOr them all into one constraint.
+  auto& self = *this;
+  for (Index i = 1; i < size(); i++) {
+    if (self[i].local != self[0].local) {
+      // They refer to different locals. Give up.
+      clear();
+      return;
+    }
+  }
+
+  // Negate them before the OR.
+  for (auto& pair : self) {
+    pair.constraint = pair.constraint.negate();
+  }
+
+  if (size() == 1) {
+    // The simple case of 1 doesn't need any more work.
+    return;
+  }
+
+  // Do the OR.
+  AndedConstraintSet anded;
+  anded.set(self[0].constraint);
+  for (Index i = 1; i < size(); i++) {
+    anded.approximateOr({self[i].constraint});
+    if (anded.provesNothing()) {
+      // We have nothing useful here.
+      clear();
+      return;
+    }
+  }
+
+  // Return only the OR'ed result.
+  auto local = self[0].local;
+  clear();
+  for (auto& c : anded) {
+    emplace_back(local, c);
+  }
+}
 
 void LocalConstraint::flip() {
   auto other = std::get<Index>(constraint.term);
@@ -1152,6 +1252,11 @@ std::ostream& operator<<(std::ostream& o, const Constraint& c) {
   return o;
 }
 
+std::ostream& operator<<(std::ostream& o, const LocalConstraint& c) {
+  o << "LocalConstraint{$" << c.local << ", " << c.constraint << '}';
+  return o;
+}
+
 std::ostream& operator<<(std::ostream& o, const AndedConstraintSet& set) {
   if (set.provesEverything()) {
     o << "AndedConstraintSet(contradiction)";

src/ir/constraint.h

@@ -32,6 +32,7 @@
 
 #include "ir/abstract.h"
 #include "support/inplace_vector.h"
+#include "support/small_vector.h"
 #include "support/span.h"
 #include "support/utilities.h"
 #include "wasm.h"
@@ -235,6 +236,12 @@ struct LocalConstraint {
   Index local;
   Constraint constraint;
 
+  LocalConstraint() = default;
+  LocalConstraint(Index local, Constraint constraint)
+    : local(local), constraint(std::move(constraint)) {}
+
+  bool operator==(const LocalConstraint&) const = default;
+
   // Try to parse BinaryenIR into a local to which a constraint is applied. For
   // example
   //
@@ -246,15 +253,43 @@ struct LocalConstraint {
   //
   static std::optional<LocalConstraint> parse(Expression* curr);
 
-  // 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 utility to parse BinaryenIR into locals and constraints on them. This is
+// similar to LocalConstraint::parse, but that parses a single constraint, while
+// this can handle a list of ANDed ones:
+//
+//   (i32.and (..A..) (..B..))
+//
+// parses into [ A, B ].
+//
+// We also set a field |hasUnknown| if we saw things we could not parse. E.g.
+//
+//   (i32.and (call $unknown) (i32.eqz (local.get $x)))
+//
+// This parses into [ $x == 0 ] and sets hasUnknown=true. Even if there are
+// unknown things, we do know that definitely $x == 0 at least, which is useful
+// in some cases.
+struct ParsedAndedConstraints : public SmallVector<LocalConstraint, 1> {
+  using SmallVector<LocalConstraint, 1>::SmallVector;
+
+  bool hasUnknown = false;
+
+  static ParsedAndedConstraints parse(Expression* curr);
+
+  // 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 ParsedAndedConstraints parseCondition(Expression* curr);
+
+  // Negate the entire list of constraints. If we fail to generate something
+  // that can be represented as a list of ANDed constraints, the list will be
+  // empty (i.e., we can prove nothing).
+  void negate();
+};
+
 // A map of locals and their constraints, representing the state at a basic
 // block. We use the following representation:
 //
@@ -366,6 +401,7 @@ struct BasicBlockConstraintMap {
 };
 
 std::ostream& operator<<(std::ostream& o, const Constraint& c);
+std::ostream& operator<<(std::ostream& o, const LocalConstraint& c);
 std::ostream& operator<<(std::ostream& o, const AndedConstraintSet& set);
 
 } // namespace wasm::constraint

src/ir/match.h

@@ -896,6 +896,9 @@ inline decltype(auto) local() {
 inline decltype(auto) local(Index* binder) {
   return Internal::LocalGetMatcher(nullptr, Internal::Any(binder));
 }
+inline decltype(auto) local(LocalGet** binder) {
+  return Internal::LocalGetMatcher(binder, Internal::Any<Index>(nullptr));
+}
 
 } // namespace wasm::Match
 

src/passes/ConstraintAnalysis.cpp

@@ -66,6 +66,8 @@
 // function analysis).
 //
 
+#include <algorithm>
+
 #include "cfg/cfg-traversal.h"
 #include "ir/constraint.h"
 #include "ir/drop.h"
@@ -146,11 +148,12 @@ struct ConstraintAnalysis
 
   void maybeMarkRelevant(Expression* curr) {
     // If this parses into a constraint on a local, that local is relevant.
-    if (auto parsed = LocalConstraint::parseCondition(curr);
-        parsed && isRelevantType(getFunction()->getLocalType(parsed->local))) {
-      relevantLocals[parsed->local] = true;
-      if (auto* other = std::get_if<Index>(&parsed->constraint.term)) {
-        relevantLocals[*other] = true;
+    for (auto& pair : ParsedAndedConstraints::parseCondition(curr)) {
+      if (isRelevantType(getFunction()->getLocalType(pair.local))) {
+        relevantLocals[pair.local] = true;
+        if (auto* other = std::get_if<Index>(&pair.constraint.term)) {
+          relevantLocals[*other] = true;
+        }
       }
     }
   }
@@ -371,9 +374,9 @@ struct ConstraintAnalysis
         // Find the constraints sent to this specific successor, if there is a
         // branch, and use them.
         if (auto branch = getBranchConstraints(block, out);
-            branch && checkRelevancy(*branch)) {
+            filterRelevant(branch), !branch.empty()) {
           auto sentConstraints = constraints;
-          applyBranchConstraints(*branch, sentConstraints);
+          applyBranchConstraints(branch, sentConstraints);
 #if CONSTRAINT_DEBUG
           std::cout << block << " sending branch to " << out
                     << " with sent constraints: " << sentConstraints << '\n';
@@ -441,6 +444,9 @@ struct ConstraintAnalysis
   void optimizeExpression(Expression** currp,
                           const BasicBlockConstraintMap& constraints) {
     auto* curr = *currp;
+    // Note that we don't need to try to parse a series of constraints with
+    // ParsedAndedConstraints: if there is a tree of ANDed things, we will
+    // simply optimize it as we walk it, each time handling one.
     auto parsed = LocalConstraint::parse(curr);
     if (!parsed) {
       return;
@@ -472,8 +478,8 @@ struct ConstraintAnalysis
 
   // Given a predecessor and one of its successors, find new constraints that
   // can be added due to the flow to that specific successor.
-  std::optional<LocalConstraint> getBranchConstraints(BasicBlock* pred,
-                                                      BasicBlock* succ) {
+  ParsedAndedConstraints getBranchConstraints(BasicBlock* pred,
+                                              BasicBlock* succ) {
     auto* brancher = pred->contents.brancher;
     if (!brancher) {
       return {};
@@ -502,32 +508,31 @@ struct ConstraintAnalysis
     return {};
   }
 
-  std::optional<LocalConstraint> getConstraintsFromIf(If* iff,
-                                                      bool physicalSuccessor) {
-    auto parsed = LocalConstraint::parseCondition(iff->condition);
-    if (parsed && !physicalSuccessor) {
+  ParsedAndedConstraints getConstraintsFromIf(If* iff, bool physicalSuccessor) {
+    auto parsed = ParsedAndedConstraints::parseCondition(iff->condition);
+    if (!physicalSuccessor) {
       // We are in the ifFalse, so negate the condition.
-      parsed->constraint = parsed->constraint.negate();
+      parsed.negate();
     }
     return parsed;
   }
 
-  std::optional<LocalConstraint>
-  getConstraintsFromBreak(Break* br, bool physicalSuccessor) {
+  ParsedAndedConstraints getConstraintsFromBreak(Break* br,
+                                                 bool physicalSuccessor) {
     // We get here when there is more than one successor, so there must be a
     // condition.
     assert(br->condition);
 
-    auto parsed = LocalConstraint::parseCondition(br->condition);
-    if (parsed && physicalSuccessor) {
+    auto parsed = ParsedAndedConstraints::parseCondition(br->condition);
+    if (physicalSuccessor) {
       // The branch was not taken, so negate the condition.
-      parsed->constraint = parsed->constraint.negate();
+      parsed.negate();
     }
     return parsed;
   }
 
-  std::optional<LocalConstraint>
-  getConstraintsFromBrOn(BrOn* brOn, bool physicalSuccessor) {
+  ParsedAndedConstraints getConstraintsFromBrOn(BrOn* brOn,
+                                                bool physicalSuccessor) {
     // The constraint on that local depends on the op.
     // TODO: Handle BrOnCast* etc using subtyping operations.
     if (brOn->op != BrOnNull && brOn->op != BrOnNonNull) {
@@ -537,10 +542,10 @@ struct ConstraintAnalysis
     // parseCondition can parse more things than a local.get, which is all we
     // handle here, but there is no other valid IR that can appear there, so we
     // can reuse it.
-    auto parsed = LocalConstraint::parseCondition(brOn->ref);
+    auto parsed = ParsedAndedConstraints::parseCondition(brOn->ref);
     // Negate depending on the op and (similar to Break) the successor.
-    if (parsed && ((brOn->op == BrOnNull) ^ physicalSuccessor)) {
-      parsed->constraint = parsed->constraint.negate();
+    if ((brOn->op == BrOnNull) ^ physicalSuccessor) {
+      parsed.negate();
     }
     return parsed;
   }
@@ -664,17 +669,31 @@ struct ConstraintAnalysis
     return true;
   }
 
+  // Filters out constraints on irrelevant locals.
+  void filterRelevant(ParsedAndedConstraints& parsed) {
+    parsed.erase(std::remove_if(parsed.begin(),
+                                parsed.end(),
+                                [&](const LocalConstraint& pair) {
+                                  return !checkRelevancy(pair);
+                                }),
+                 parsed.end());
+  }
+
   // Apply branch constraints to the current set of constraints.
-  void applyBranchConstraints(const LocalConstraint& branch,
+  void applyBranchConstraints(const ParsedAndedConstraints& branch,
                               BasicBlockConstraintMap& constraints) {
-    // Extend the range of values in the "jump ahead" manner described in the
-    // top-level comment.
-    if (applyBranchRangeExtensionToConstraints(branch, constraints)) {
-      return;
-    }
+    for (auto& pair : branch) {
+      // Extend the range of values in the "jump ahead" manner described in the
+      // top-level comment.
+      if (!applyBranchRangeExtensionToConstraints(pair, constraints)) {
+        // Otherwise, apply the constraint normally.
+        constraints.approximateAnd(pair.local, pair.constraint);
+      }
 
-    // Otherwise, apply the constraint normally.
-    constraints.approximateAnd(branch.local, branch.constraint);
+      if (constraints.unreachable) {
+        return;
+      }
+    }
   }
 
   bool