Move to registryData

Author: devajithvs

mlir/include/mlir/Query/QuerySession.h

@@ -21,8 +21,7 @@ class QuerySession {
 public:
   QuerySession(Operation *rootOp,
                const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
-               unsigned bufferId,
-               const mlir::query::matcher::RegistryMaps &registryData)
+               unsigned bufferId, const matcher::RegistryMaps &registryData)
       : rootOp(rootOp), sourceMgr(sourceMgr), bufferId(bufferId),
         registryData(registryData), terminate(false) {}
 
@@ -33,7 +32,7 @@ class QuerySession {
   Operation *rootOp;
   const std::shared_ptr<llvm::SourceMgr> sourceMgr;
   unsigned bufferId;
-  const mlir::query::matcher::RegistryMaps &registryData;
+  const matcher::RegistryMaps &registryData;
   bool terminate;
   llvm::StringMap<matcher::VariantValue> namedValues;
 };

mlir/include/mlir/Tools/mlir-query/MlirQueryMain.h

@@ -14,13 +14,16 @@
 #ifndef MLIR_TOOLS_MLIRQUERY_MLIRQUERYMAIN_H
 #define MLIR_TOOLS_MLIRQUERY_MLIRQUERYMAIN_H
 
+#include "mlir/Query/Matcher/Registry.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 class MLIRContext;
 
-LogicalResult mlirQueryMain(int argc, char **argv, MLIRContext &context);
+LogicalResult
+mlirQueryMain(int argc, char **argv, MLIRContext &context,
+              const mlir::query::matcher::RegistryMaps &registryData);
 
 } // namespace mlir
 

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

@@ -13,8 +13,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-query/MlirQueryMain.h"
-#include "mlir/IR/Matchers.h"
-#include "mlir/Query/Matcher/Registry.h"
 #include "mlir/Query/QueryParser.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
@@ -28,8 +26,9 @@
 // Query Parser
 //===----------------------------------------------------------------------===//
 
-mlir::LogicalResult mlir::mlirQueryMain(int argc, char **argv,
-                                        MLIRContext &context) {
+mlir::LogicalResult
+mlir::mlirQueryMain(int argc, char **argv, MLIRContext &context,
+                    const mlir::query::matcher::RegistryMaps &registryData) {
 
   // Override the default '-h' and use the default PrintHelpMessage() which
   // won't print options in categories.
@@ -86,29 +85,6 @@ mlir::LogicalResult mlir::mlirQueryMain(int argc, char **argv,
   if (!opRef)
     return failure();
 
-  mlir::query::matcher::RegistryMaps registryData;
-
-  // This is needed because these matchers are defined as overloaded functions.
-  using IsConstantOp = mlir::detail::constant_op_matcher();
-  using HasOpAttrName = mlir::detail::AttrOpMatcher(StringRef);
-  using HasOpName = mlir::detail::NameOpMatcher(StringRef);
-
-  // Matchers registered in alphabetical order for consistency:
-  registryData.registerMatcher("hasOpAttrName",
-                               static_cast<HasOpAttrName *>(m_Attr));
-  registryData.registerMatcher("hasOpName", static_cast<HasOpName *>(m_Op));
-  registryData.registerMatcher("isConstantOp",
-                               static_cast<IsConstantOp *>(m_Constant));
-  registryData.registerMatcher("isNegInfFloat", m_NegInfFloat);
-  registryData.registerMatcher("isNegZeroFloat", m_NegZeroFloat);
-  registryData.registerMatcher("isNonZero", m_NonZero);
-  registryData.registerMatcher("isOne", m_One);
-  registryData.registerMatcher("isOneFloat", m_OneFloat);
-  registryData.registerMatcher("isPosInfFloat", m_PosInfFloat);
-  registryData.registerMatcher("isPosZeroFloat", m_PosZeroFloat);
-  registryData.registerMatcher("isZero", m_Zero);
-  registryData.registerMatcher("isZeroFloat", m_AnyZeroFloat);
-
   mlir::query::QuerySession QS(opRef.get(), sourceMgr, bufferId, registryData);
   if (!commands.empty()) {
     for (auto &command : commands) {

mlir/tools/mlir-query/mlir-query.cpp

@@ -13,11 +13,18 @@
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/Query/Matcher/Registry.h"
 #include "mlir/Tools/mlir-query/MlirQueryMain.h"
 
 using namespace mlir;
 
+// This is needed because these matchers are defined as overloaded functions.
+using HasOpAttrName = detail::AttrOpMatcher(StringRef);
+using HasOpName = detail::NameOpMatcher(StringRef);
+using IsConstantOp = detail::constant_op_matcher();
+
 namespace test {
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestDialect(DialectRegistry &);
@@ -28,10 +35,29 @@ int main(int argc, char **argv) {
 
   DialectRegistry registry;
   registerAllDialects(registry);
+
+  query::matcher::RegistryMaps registryData;
+
+  // Matchers registered in alphabetical order for consistency:
+  registryData.registerMatcher("hasOpAttrName",
+                               static_cast<HasOpAttrName *>(m_Attr));
+  registryData.registerMatcher("hasOpName", static_cast<HasOpName *>(m_Op));
+  registryData.registerMatcher("isConstantOp",
+                               static_cast<IsConstantOp *>(m_Constant));
+  registryData.registerMatcher("isNegInfFloat", m_NegInfFloat);
+  registryData.registerMatcher("isNegZeroFloat", m_NegZeroFloat);
+  registryData.registerMatcher("isNonZero", m_NonZero);
+  registryData.registerMatcher("isOne", m_One);
+  registryData.registerMatcher("isOneFloat", m_OneFloat);
+  registryData.registerMatcher("isPosInfFloat", m_PosInfFloat);
+  registryData.registerMatcher("isPosZeroFloat", m_PosZeroFloat);
+  registryData.registerMatcher("isZero", m_Zero);
+  registryData.registerMatcher("isZeroFloat", m_AnyZeroFloat);
+
 #ifdef MLIR_INCLUDE_TESTS
   test::registerTestDialect(registry);
 #endif
   MLIRContext context(registry);
 
-  return failed(mlirQueryMain(argc, argv, context));
+  return failed(mlirQueryMain(argc, argv, context, registryData));
 }