Accept stringref arguments to .extract()

Author: devajithvs

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

@@ -79,8 +79,10 @@ class DynMatcher {
   DynMatcher *clone() const { return new DynMatcher(*this); }
 
   void setExtract(bool extractFunction) { ExtractFunction = extractFunction; };
+  void setFunctionName(StringRef functionName) { FunctionName = functionName; };
 
   bool getExtract() const { return ExtractFunction; };
+  StringRef getFunctionName() const { return FunctionName; };
 
 private:
   MLIRNodeKind SupportedKind;
@@ -91,6 +93,7 @@ class DynMatcher {
   MLIRNodeKind RestrictKind;
   llvm::IntrusiveRefCntPtr<DynMatcherInterface> Implementation;
   bool ExtractFunction;
+  StringRef FunctionName;
 };
 
 // Wrapper of a MatcherInterface<T> *

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

@@ -324,11 +324,13 @@ bool Parser::parseMatcherExpressionImpl(VariantValue *Value) {
   }
 
   bool extractFunction = false;
+  StringRef functionName;
   if (Tokenizer->peekNextToken().Kind == TokenInfo::TK_Period) {
     // Parse ".extract()"
     Tokenizer->consumeNextToken(); // consume the period.
     const TokenInfo ExtractToken = Tokenizer->consumeNextToken();
     const TokenInfo OpenToken = Tokenizer->consumeNextToken();
+    const TokenInfo NameToken = Tokenizer->consumeNextToken();
     const TokenInfo CloseToken = Tokenizer->consumeNextToken();
 
     // TODO: We could use different error codes for each/some to be more
@@ -342,18 +344,25 @@ bool Parser::parseMatcherExpressionImpl(VariantValue *Value) {
       Error->addError(OpenToken.Range, Error->ET_ParserMalformedBindExpr);
       return false;
     }
+    if (NameToken.Kind != TokenInfo::TK_Literal ||
+        !NameToken.Value.isString()) {
+      Error->addError(NameToken.Range, Error->ET_ParserMalformedBindExpr);
+      return false;
+    }
     if (CloseToken.Kind != TokenInfo::TK_CloseParen) {
       Error->addError(CloseToken.Range, Error->ET_ParserMalformedBindExpr);
       return false;
     }
+    functionName = NameToken.Value.getString();
     extractFunction = true;
+    LLVM_DEBUG(DBGS() << "Function Name: " << functionName << "\n");
   }
 
   // Merge the start and end infos.
   SourceRange MatcherRange = NameToken.Range;
   MatcherRange.End = EndToken.Range.End;
-  DynMatcher *Result = S->actOnMatcherExpression(NameToken.Text, MatcherRange,
-                                                 extractFunction, Args, Error);
+  DynMatcher *Result = S->actOnMatcherExpression(
+      NameToken.Text, MatcherRange, extractFunction, functionName, Args, Error);
 
   if (Result == nullptr) {
     // TODO: Add appropriate error.
@@ -415,10 +424,11 @@ class RegistrySema : public Parser::Sema {
   DynMatcher *actOnMatcherExpression(StringRef MatcherName,
                                      const SourceRange &NameRange,
                                      bool ExtractFunction,
+                                     StringRef FunctionName,
                                      ArrayRef<ParserValue> Args,
                                      Diagnostics *Error) override {
-    return Registry::constructMatcherWrapper(MatcherName, NameRange,
-                                             ExtractFunction, Args, Error);
+    return Registry::constructMatcherWrapper(
+        MatcherName, NameRange, ExtractFunction, FunctionName, Args, Error);
   }
 };
 

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

@@ -60,11 +60,10 @@ class Parser {
     // if an error occurred. In that case, Error will contain a
     // description of the error.
     // The caller takes ownership of the Matcher object returned.
-    virtual DynMatcher *actOnMatcherExpression(StringRef MatcherName,
-                                               const SourceRange &NameRange,
-                                               bool ExtractFunction,
-                                               ArrayRef<ParserValue> Args,
-                                               Diagnostics *Error) = 0;
+    virtual DynMatcher *
+    actOnMatcherExpression(StringRef MatcherName, const SourceRange &NameRange,
+                           bool ExtractFunction, StringRef FunctionName,
+                           ArrayRef<ParserValue> Args, Diagnostics *Error) = 0;
   };
 
   // Parse a matcher expression, creating matchers from the registry.

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

@@ -56,7 +56,7 @@ getMatches(Operation *rootOp, const matcher::DynMatcher *matcher) {
 
 // TODO: Only supports operation node type.
 Operation *extractFunction(std::vector<matcher::DynTypedNode> &nodes,
-                           OpBuilder builder) {
+                           OpBuilder builder, StringRef functionName) {
   std::vector<Operation *> slice;
   std::vector<Value> values;
 
@@ -81,7 +81,7 @@ Operation *extractFunction(std::vector<matcher::DynTypedNode> &nodes,
 
   auto loc = builder.getUnknownLoc();
   func::FuncOp funcOp = func::FuncOp::create(
-      loc, "extracted",
+      loc, functionName,
       builder.getFunctionType(ValueRange(values), resultType));
 
   loc = funcOp.getLoc();
@@ -119,10 +119,11 @@ bool MatchQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
   auto matches = getMatches(rootOp, matcher);
 
   if (matcher->getExtract()) {
+    auto functionName = matcher->getFunctionName();
     MLIRContext context;
     context.loadDialect<func::FuncDialect>();
     OpBuilder builder(&context);
-    Operation *function = extractFunction(matches, builder);
+    Operation *function = extractFunction(matches, builder, functionName);
     OS << "\n\n" << *function << "\n\n\n";
   } else {
     unsigned MatchCount = 0;

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

@@ -120,16 +120,15 @@ DynMatcher *Registry::constructMatcher(StringRef MatcherName,
 }
 
 // static
-DynMatcher *Registry::constructMatcherWrapper(StringRef MatcherName,
-                                              const SourceRange &NameRange,
-                                              bool ExtractFunction,
-                                              ArrayRef<ParserValue> Args,
-                                              Diagnostics *Error) {
+DynMatcher *Registry::constructMatcherWrapper(
+    StringRef MatcherName, const SourceRange &NameRange, bool ExtractFunction,
+    StringRef FunctionName, ArrayRef<ParserValue> Args, Diagnostics *Error) {
 
   DynMatcher *Out = constructMatcher(MatcherName, NameRange, Args, Error);
   if (!Out)
     return Out;
   Out->setExtract(ExtractFunction);
+  Out->setFunctionName(FunctionName);
   return Out;
 }
 

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

@@ -45,11 +45,10 @@ class Registry {
                                       const SourceRange &NameRange,
                                       ArrayRef<ParserValue> Args,
                                       Diagnostics *Error);
-  static DynMatcher *constructMatcherWrapper(StringRef MatcherName,
-                                             const SourceRange &NameRange,
-                                             bool ExtractFunction,
-                                             ArrayRef<ParserValue> Args,
-                                             Diagnostics *Error);
+  static DynMatcher *
+  constructMatcherWrapper(StringRef MatcherName, const SourceRange &NameRange,
+                          bool ExtractFunction, StringRef FunctionName,
+                          ArrayRef<ParserValue> Args, Diagnostics *Error);
 };
 
 } // namespace matcher