Add anyOf and allOf matchers - Need fixing - Add polymorphic matcher

Author: devajithvs

mlir/lib/Tools/mlir-query/ExtraMatchers.h

@@ -10,8 +10,7 @@
 // goal is to move this to include/mlir/IR/Matchers.h after the initial testing
 // phase. The matchers in this file are:
 //
-// - operation(args...): Matches an operation that matches all of the matchers
-// in the vector `matchers`.
+// - operation(args...): Matches all of the matchers in the vector `matchers`.
 //
 // - argument(innerMatcher, index): Matches an operation argument that matches
 // `innerMatcher` at the given `index`.
@@ -39,11 +38,9 @@ namespace extramatcher {
 
 namespace detail {
 
-// TODO: Rename to AllOf
-// OperationMatcher takes a vector of DynMatchers and returns true if all
-// DynMatchers match the given operation.
-struct OperationMatcher {
-  OperationMatcher(std::vector<matcher::DynMatcher> matchers)
+// AllOf takes a vector of DynMatchers and returns true if all the DynMatchers match the given operation.
+struct AllOfMatcher {
+  AllOfMatcher(std::vector<matcher::DynMatcher> matchers)
       : matchers(matchers) {}
   bool match(Operation *op) {
     matcher::DynTypedNode node = matcher::DynTypedNode::create(op);
@@ -54,6 +51,19 @@ struct OperationMatcher {
   std::vector<matcher::DynMatcher> matchers;
 };
 
+// AnyOf takes a vector of DynMatchers and returns true if any of the DynMatchers match the given operation.
+struct AnyOfMatcher {
+  AnyOfMatcher(std::vector<matcher::DynMatcher> matchers)
+      : matchers(matchers) {}
+  bool match(Operation *op) {
+    matcher::DynTypedNode node = matcher::DynTypedNode::create(op);
+    return llvm::any_of(matchers, [&](const matcher::DynMatcher &matcher) {
+      return matcher.matches(node);
+    });
+  }
+  std::vector<matcher::DynMatcher> matchers;
+};
+
 // ArgumentMatcher matches the operand of an operation at a specific index.
 struct ArgumentMatcher {
   ArgumentMatcher(matcher::DynMatcher innerMatcher, unsigned index)
@@ -145,10 +155,14 @@ struct DefinedByMatcher {
 
 } // namespace detail
 
-// TODO: Rename to allOf()
-inline detail::OperationMatcher operation(matcher::DynMatcher args...) {
+inline detail::AllOfMatcher allOf(matcher::DynMatcher args...) {
+  std::vector<matcher::DynMatcher> matchers({args});
+  return detail::AllOfMatcher(matchers);
+}
+
+inline detail::AnyOfMatcher anyOf(matcher::DynMatcher args...) {
   std::vector<matcher::DynMatcher> matchers({args});
-  return detail::OperationMatcher(matchers);
+  return detail::AnyOfMatcher(matchers);
 }
 
 inline detail::ArgumentMatcher hasArgument(matcher::DynMatcher innerMatcher,

mlir/lib/Tools/mlir-query/MatchersInternal.h

@@ -122,7 +122,8 @@ class SingleMatcher : public MatcherInterface<T> {
   MatcherFn matcherFn;
 };
 
-// VariadicMatcher takes a vector of Matchers and returns true if all Matchers
+// TODO: Use a polymorphic matcher instead for this usecase
+// VariadicMatcher takes a vector of Matchers and returns true if any Matchers
 // match the given operation.
 template <typename T>
 class VariadicMatcher : public MatcherInterface<T> {
@@ -131,7 +132,7 @@ class VariadicMatcher : public MatcherInterface<T> {
 
   bool matches(T Node) override {
     DynTypedNode DynNode = DynTypedNode::create(Node);
-    return llvm::all_of(matchers, [&](const DynMatcher &matcher) {
+    return llvm::any_of(matchers, [&](const DynMatcher &matcher) {
       return matcher.matches(DynNode);
     });
   }

mlir/lib/Tools/mlir-query/Registry.cpp

@@ -68,7 +68,8 @@ RegistryMaps::RegistryMaps() {
   };
 
   // Register matchers using the template function
-  registerOpMatcher("operation", extramatcher::operation);
+  registerOpMatcher("allOf", extramatcher::allOf);
+  registerOpMatcher("anyOf", extramatcher::anyOf);
   registerOpMatcher("hasArgument", extramatcher::hasArgument);
   registerOpMatcher("definedBy", extramatcher::definedBy);
   registerOpMatcher("getDefinitions", extramatcher::getDefinitions);