Add new unary and binary matchers (#3168)

Adds new matchers that allow for matching any unary or binary operation and
optionally extracting it. The previous matchers only allowed matching specific
unary and binary operations. This should help simplify #3132.

3 files changed, 148 insertions, 47 deletions

diff --git a/src/ir/match.h b/src/ir/match.h
index 9b7d1ff0f..2ab3afe65 100644
--- a/src/ir/match.h
+++ b/src/ir/match.h
@@ -476,82 +476,108 @@ template<class S> inline decltype(auto) ConstMatcher(Const** binder, S&& s) {
   return Matcher<Const*, S>(binder, {}, s);
 }
 
-// Unary and AbstractUnary
-struct UnaryK {
+// Unary, UnaryOp and AbstractUnaryOp
+template<> struct NumComponents<Unary*> { static constexpr size_t value = 2; };
+template<> struct GetComponent<Unary*, 0> {
+  UnaryOp operator()(Unary* curr) { return curr->op; }
+};
+template<> struct GetComponent<Unary*, 1> {
+  Expression* operator()(Unary* curr) { return curr->value; }
+};
+struct UnaryOpK {
   using Op = UnaryOp;
   static UnaryOp getOp(Type, Op op) { return op; }
 };
-struct AbstractUnaryK {
+struct AbstractUnaryOpK {
   using Op = Abstract::Op;
   static UnaryOp getOp(Type type, Abstract::Op op) {
     return Abstract::getUnary(type, op);
   }
 };
-template<class T> struct UnaryKind {};
-template<class T> struct KindTypeRegistry<UnaryKind<T>> {
+template<class T> struct UnaryOpKind {};
+template<class T> struct KindTypeRegistry<UnaryOpKind<T>> {
   using matched_t = Unary*;
   using data_t = typename T::Op;
 };
-template<class T> struct MatchSelf<UnaryKind<T>> {
+template<class T> struct MatchSelf<UnaryOpKind<T>> {
   bool operator()(Unary* curr, typename T::Op op) {
     return curr->op == T::getOp(curr->value->type, op);
   }
 };
-template<class T> struct NumComponents<UnaryKind<T>> {
+template<class T> struct NumComponents<UnaryOpKind<T>> {
   static constexpr size_t value = 1;
 };
-template<class T> struct GetComponent<UnaryKind<T>, 0> {
+template<class T> struct GetComponent<UnaryOpKind<T>, 0> {
   Expression* operator()(Unary* curr) { return curr->value; }
 };
+template<class S1, class S2>
+inline decltype(auto) UnaryMatcher(Unary** binder, S1&& s1, S2&& s2) {
+  return Matcher<Unary*, S1, S2>(binder, {}, s1, s2);
+}
 template<class S>
-inline decltype(auto) UnaryMatcher(Unary** binder, UnaryOp op, S&& s) {
-  return Matcher<UnaryKind<UnaryK>, S>(binder, op, s);
+inline decltype(auto) UnaryOpMatcher(Unary** binder, UnaryOp op, S&& s) {
+  return Matcher<UnaryOpKind<UnaryOpK>, S>(binder, op, s);
 }
 template<class S>
 inline decltype(auto)
-AbstractUnaryMatcher(Unary** binder, Abstract::Op op, S&& s) {
-  return Matcher<UnaryKind<AbstractUnaryK>, S>(binder, op, s);
+AbstractUnaryOpMatcher(Unary** binder, Abstract::Op op, S&& s) {
+  return Matcher<UnaryOpKind<AbstractUnaryOpK>, S>(binder, op, s);
 }
 
-// Binary and AbstractBinary
-struct BinaryK {
+// Binary, BinaryOp and AbstractBinaryOp
+template<> struct NumComponents<Binary*> { static constexpr size_t value = 3; };
+template<> struct GetComponent<Binary*, 0> {
+  BinaryOp operator()(Binary* curr) { return curr->op; }
+};
+template<> struct GetComponent<Binary*, 1> {
+  Expression* operator()(Binary* curr) { return curr->left; }
+};
+template<> struct GetComponent<Binary*, 2> {
+  Expression* operator()(Binary* curr) { return curr->right; }
+};
+struct BinaryOpK {
   using Op = BinaryOp;
   static BinaryOp getOp(Type, Op op) { return op; }
 };
-struct AbstractBinaryK {
+struct AbstractBinaryOpK {
   using Op = Abstract::Op;
   static BinaryOp getOp(Type type, Abstract::Op op) {
     return Abstract::getBinary(type, op);
   }
 };
-template<class T> struct BinaryKind {};
-template<class T> struct KindTypeRegistry<BinaryKind<T>> {
+template<class T> struct BinaryOpKind {};
+template<class T> struct KindTypeRegistry<BinaryOpKind<T>> {
   using matched_t = Binary*;
   using data_t = typename T::Op;
 };
-template<class T> struct MatchSelf<BinaryKind<T>> {
+template<class T> struct MatchSelf<BinaryOpKind<T>> {
   bool operator()(Binary* curr, typename T::Op op) {
     return curr->op == T::getOp(curr->left->type, op);
   }
 };
-template<class T> struct NumComponents<BinaryKind<T>> {
+template<class T> struct NumComponents<BinaryOpKind<T>> {
   static constexpr size_t value = 2;
 };
-template<class T> struct GetComponent<BinaryKind<T>, 0> {
+template<class T> struct GetComponent<BinaryOpKind<T>, 0> {
   Expression* operator()(Binary* curr) { return curr->left; }
 };
-template<class T> struct GetComponent<BinaryKind<T>, 1> {
+template<class T> struct GetComponent<BinaryOpKind<T>, 1> {
   Expression* operator()(Binary* curr) { return curr->right; }
 };
+template<class S1, class S2, class S3>
+inline decltype(auto)
+BinaryMatcher(Binary** binder, S1&& s1, S2&& s2, S3&& s3) {
+  return Matcher<Binary*, S1, S2, S3>(binder, {}, s1, s2, s3);
+}
 template<class S1, class S2>
 inline decltype(auto)
-BinaryMatcher(Binary** binder, BinaryOp op, S1&& s1, S2&& s2) {
-  return Matcher<BinaryKind<BinaryK>, S1, S2>(binder, op, s1, s2);
+BinaryOpMatcher(Binary** binder, BinaryOp op, S1&& s1, S2&& s2) {
+  return Matcher<BinaryOpKind<BinaryOpK>, S1, S2>(binder, op, s1, s2);
 }
 template<class S1, class S2>
 inline decltype(auto)
-AbstractBinaryMatcher(Binary** binder, Abstract::Op op, S1&& s1, S2&& s2) {
-  return Matcher<BinaryKind<AbstractBinaryK>, S1, S2>(binder, op, s1, s2);
+AbstractBinaryOpMatcher(Binary** binder, Abstract::Op op, S1&& s1, S2&& s2) {
+  return Matcher<BinaryOpKind<AbstractBinaryOpK>, S1, S2>(binder, op, s1, s2);
 }
 
 // Select
@@ -751,37 +777,59 @@ inline decltype(auto) constant(Const** binder, int32_t x) {
 inline decltype(auto) any() { return Internal::Any<Expression*>(nullptr); }
 inline decltype(auto) any(Expression** binder) { return Internal::Any(binder); }
 
+template<class S> inline decltype(auto) unary(S&& s) {
+  return Internal::UnaryMatcher(nullptr, Internal::Any<UnaryOp>(nullptr), s);
+}
+template<class S> inline decltype(auto) unary(Unary** binder, S&& s) {
+  return Internal::UnaryMatcher(binder, Internal::Any<UnaryOp>(nullptr), s);
+}
+template<class S> inline decltype(auto) unary(UnaryOp* binder, S&& s) {
+  return Internal::UnaryMatcher(nullptr, Internal::Any<UnaryOp>(binder), s);
+}
 template<class S> inline decltype(auto) unary(UnaryOp op, S&& s) {
-  return Internal::UnaryMatcher(nullptr, op, s);
+  return Internal::UnaryOpMatcher(nullptr, op, s);
 }
 template<class S> inline decltype(auto) unary(Abstract::Op op, S&& s) {
-  return Internal::AbstractUnaryMatcher(nullptr, op, s);
+  return Internal::AbstractUnaryOpMatcher(nullptr, op, s);
 }
 template<class S>
 inline decltype(auto) unary(Unary** binder, UnaryOp op, S&& s) {
-  return Internal::UnaryMatcher(binder, op, s);
+  return Internal::UnaryOpMatcher(binder, op, s);
 }
 template<class S>
 inline decltype(auto) unary(Unary** binder, Abstract::Op op, S&& s) {
-  return Internal::AbstractUnaryMatcher(binder, op, s);
+  return Internal::AbstractUnaryOpMatcher(binder, op, s);
+}
+template<class S1, class S2> inline decltype(auto) binary(S1&& s1, S2&& s2) {
+  return Internal::BinaryMatcher(
+    nullptr, Internal::Any<BinaryOp>(nullptr), s1, s2);
+}
+template<class S1, class S2>
+inline decltype(auto) binary(Binary** binder, S1&& s1, S2&& s2) {
+  return Internal::BinaryMatcher(
+    binder, Internal::Any<BinaryOp>(nullptr), s1, s2);
+}
+template<class S1, class S2>
+inline decltype(auto) binary(BinaryOp* binder, S1&& s1, S2&& s2) {
+  return Internal::BinaryMatcher(
+    nullptr, Internal::Any<BinaryOp>(binder), s1, s2);
 }
-
 template<class S1, class S2>
 inline decltype(auto) binary(BinaryOp op, S1&& s1, S2&& s2) {
-  return Internal::BinaryMatcher(nullptr, op, s1, s2);
+  return Internal::BinaryOpMatcher(nullptr, op, s1, s2);
 }
 template<class S1, class S2>
 inline decltype(auto) binary(Abstract::Op op, S1&& s1, S2&& s2) {
-  return Internal::AbstractBinaryMatcher(nullptr, op, s1, s2);
+  return Internal::AbstractBinaryOpMatcher(nullptr, op, s1, s2);
 }
 template<class S1, class S2>
 inline decltype(auto) binary(Binary** binder, BinaryOp op, S1&& s1, S2&& s2) {
-  return Internal::BinaryMatcher(binder, op, s1, s2);
+  return Internal::BinaryOpMatcher(binder, op, s1, s2);
 }
 template<class S1, class S2>
 inline decltype(auto)
 binary(Binary** binder, Abstract::Op op, S1&& s1, S2&& s2) {
-  return Internal::AbstractBinaryMatcher(binder, op, s1, s2);
+  return Internal::AbstractBinaryOpMatcher(binder, op, s1, s2);
 }
 
 template<class S1, class S2, class S3>
diff --git a/test/example/match.cpp b/test/example/match.cpp
index e485f2355..21bbbab98 100644
--- a/test/example/match.cpp
+++ b/test/example/match.cpp
@@ -324,20 +324,44 @@ void test_internal_unary() {
   std::cout << "Testing Internal::UnaryMatcher\n";
 
   Unary* out = nullptr;
+  UnaryOp op;
 
-  auto eqz32Matcher =
-    Internal::UnaryMatcher(&out, EqZInt32, Internal::Any<Expression*>(nullptr));
+  auto unMatcher = Internal::UnaryMatcher(
+    &out, Internal::Any<UnaryOp>(&op), Internal::Any<Expression*>(nullptr));
+  assert(unMatcher.matches(eqz32));
+  assert(out == eqz32);
+  assert(op == EqZInt32);
+  assert(unMatcher.matches(eqz64));
+  assert(out == eqz64);
+  assert(op == EqZInt64);
+  assert(unMatcher.matches(clz));
+  assert(out == clz);
+  assert(op == ClzInt32);
+  assert(!unMatcher.matches(nop));
+
+  assert(matches(clz, unary(any())));
+  assert(matches(eqz64, unary(&out, any())));
+  assert(out == eqz64);
+  assert(matches(eqz32, unary(&op, any())));
+  assert(op == EqZInt32);
+
+  std::cout << "Testing Internal::UnaryOpMatcher\n";
+
+  out = nullptr;
+
+  auto eqz32Matcher = Internal::UnaryOpMatcher(
+    &out, EqZInt32, Internal::Any<Expression*>(nullptr));
   assert(eqz32Matcher.matches(eqz32));
   assert(out == eqz32);
   assert(!eqz32Matcher.matches(eqz64));
   assert(!eqz32Matcher.matches(clz));
   assert(!eqz32Matcher.matches(nop));
 
-  std::cout << "Testing Internal::AbstractUnaryMatcher\n";
+  std::cout << "Testing Internal::AbstractUnaryOpMatcher\n";
 
   out = nullptr;
 
-  auto eqzMatcher = Internal::AbstractUnaryMatcher(
+  auto eqzMatcher = Internal::AbstractUnaryOpMatcher(
     &out, Abstract::EqZ, Internal::Any<Expression*>(nullptr));
   assert(eqzMatcher.matches(eqz32));
   assert(out == eqz32);
@@ -362,27 +386,54 @@ void test_internal_binary() {
   std::cout << "Testing Internal::BinaryMatcher\n";
 
   Binary* out = nullptr;
+  BinaryOp op;
 
-  auto eq32Matcher =
+  auto binMatcher =
     Internal::BinaryMatcher(&out,
-                            EqInt32,
+                            Internal::Any<BinaryOp>(&op),
                             Internal::Any<Expression*>(nullptr),
                             Internal::Any<Expression*>(nullptr));
+  assert(binMatcher.matches(eq32));
+  assert(out == eq32);
+  assert(op == EqInt32);
+  assert(binMatcher.matches(eq64));
+  assert(out == eq64);
+  assert(op == EqInt64);
+  assert(binMatcher.matches(add));
+  assert(out == add);
+  assert(op == AddInt32);
+  assert(!binMatcher.matches(nop));
+
+  assert(matches(add, binary(any(), any())));
+  assert(matches(eq64, binary(&out, any(), any())));
+  assert(out == eq64);
+  assert(matches(eq32, binary(&op, any(), any())));
+  assert(op == EqInt32);
+
+  std::cout << "Testing Internal::BinaryOpMatcher\n";
+
+  out = nullptr;
+
+  auto eq32Matcher =
+    Internal::BinaryOpMatcher(&out,
+                              EqInt32,
+                              Internal::Any<Expression*>(nullptr),
+                              Internal::Any<Expression*>(nullptr));
   assert(eq32Matcher.matches(eq32));
   assert(out == eq32);
   assert(!eq32Matcher.matches(eq64));
   assert(!eq32Matcher.matches(add));
   assert(!eq32Matcher.matches(nop));
 
-  std::cout << "Testing Internal::AbstractBinaryMatcher\n";
+  std::cout << "Testing Internal::AbstractBinaryOpMatcher\n";
 
   out = nullptr;
 
   auto eqMatcher =
-    Internal::AbstractBinaryMatcher(&out,
-                                    Abstract::Eq,
-                                    Internal::Any<Expression*>(nullptr),
-                                    Internal::Any<Expression*>(nullptr));
+    Internal::AbstractBinaryOpMatcher(&out,
+                                      Abstract::Eq,
+                                      Internal::Any<Expression*>(nullptr),
+                                      Internal::Any<Expression*>(nullptr));
   assert(eqMatcher.matches(eq32));
   assert(out == eq32);
   assert(eqMatcher.matches(eq64));
diff --git a/test/example/match.txt b/test/example/match.txt
index 697d9e9f0..c4ad57d73 100644
--- a/test/example/match.txt
+++ b/test/example/match.txt
@@ -3,7 +3,9 @@ Testing Internal::Exact
 Testing Internal::{I32,I64,Int,F32,F64,Float}Lit
 Testing Internal::ConstantMatcher
 Testing Internal::UnaryMatcher
-Testing Internal::AbstractUnaryMatcher
+Testing Internal::UnaryOpMatcher
+Testing Internal::AbstractUnaryOpMatcher
 Testing Internal::BinaryMatcher
-Testing Internal::AbstractBinaryMatcher
+Testing Internal::BinaryOpMatcher
+Testing Internal::AbstractBinaryOpMatcher
 Testing Internal::SelectMatcher