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

[Jacques Pienaar]

================
@@ -7,32 +7,60 @@
 //===----------------------------------------------------------------------===//
 //
 // 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;
+private:
+  static void printMatch(llvm::raw_ostream &os, QuerySession &qs,
+                         mlir::Operation *op, const std::string &binding) {
+    auto fileLoc = op->getLoc()->findInstanceOf<FileLineColLoc>();
+    auto smloc = qs.getSourceManager().FindLocForLineAndColumn(
+        qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
+    qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,
+                                       "\"" + binding + "\" binds here");
+  };
 
-    // Simple match finding with walk.
+public:
+  static std::vector<Operation *>
+  getMatches(Operation *root, QueryOptions &options, DynMatcher matcher,
+             llvm::raw_ostream &os, QuerySession &qs) {
+    unsigned matchCount = 0;
+    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);
+        os << "Match #" << ++matchCount << ":\n\n";
+        printMatch(os, qs, subOp, "root");
+      } else {
+        SmallVector<Operation *> printingOps;
+        if (matcher.match(subOp, tempStorage, options)) {
+          os << "Match #" << ++matchCount << ":\n\n";
+          SmallVector<Operation *> printingOps(tempStorage.takeVector());
+          for (auto op : printingOps) {
+            printMatch(os, qs, op, "root");
----------------
jpienaar wrote:

Would this result in marking these all as root?

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