[Mlir-commits] [mlir] [mlir] MLIR-QUERY slice-matchers implementation (PR #115670)

[Denzel-Brian Budii]

================
@@ -7,32 +7,93 @@
 //===----------------------------------------------------------------------===//
 //
 // This file contains the MatchFinder class, which is used to find operations
-// that match a given matcher.
+// that match a given matcher and print them.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
 #define MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
 
 #include "MatchersInternal.h"
+#include "mlir/Query/QuerySession.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace mlir::query::matcher {
 
-// MatchFinder is used to find all operations that match a given matcher.
 class MatchFinder {
-public:
-  // Returns all operations that match the given matcher.
-  static std::vector<Operation *> getMatches(Operation *root,
-                                             DynMatcher matcher) {
-    std::vector<Operation *> matches;
 
-    // Simple match finding with walk.
+public:
+  //
+  // getMatches walks the IR and prints operations as soon as it matches them
+  // if a matcher is to be further extracted into the function, then it does not
+  // print operations
+  //
+  static std::vector<Operation *>
+  getMatches(Operation *root, QueryOptions &options, DynMatcher matcher,
+             llvm::raw_ostream &os, QuerySession &qs) {
+    int matchCount = 0;
+    bool printMatchingOps = true;
+    // If matcher is to be extracted to a function, we don't want to print
+    // matching ops to sdout
+    if (matcher.hasFunctionName()) {
+      printMatchingOps = false;
+    }
+    std::vector<Operation *> matchedOps;
+    SetVector<Operation *> tempStorage;
+    os << "\n";
     root->walk([&](Operation *subOp) {
-      if (matcher.match(subOp))
-        matches.push_back(subOp);
+      if (matcher.match(subOp)) {
+        matchedOps.push_back(subOp);
+        if (printMatchingOps) {
+          os << "Match #" << ++matchCount << ":\n\n";
+          printMatch(os, qs, subOp, "root");
+        }
+      } else {
+        SmallVector<Operation *> printingOps;
+        if (matcher.match(subOp, tempStorage, options)) {
+          if (printMatchingOps) {
+            os << "Match #" << ++matchCount << ":\n\n";
+          }
+          SmallVector<Operation *> printingOps(tempStorage.takeVector());
+          for (auto op : printingOps) {
+            if (printMatchingOps) {
+              printMatch(os, qs, op, "root");
----------------
chios202 wrote:

This should use the overloaded print method with no argument for the binding. I accidentally reverted this back to using the old method, but will update it accordingly with the next changes. 

https://github.com/llvm/llvm-project/pull/115670