[mlir][mlir-query] Add function extraction feature to mlir-query

mlir/include/mlir/Query/Matcher/ErrorBuilder.h

struct SourceRange {
   SourceLocation end{};
 };
 
 // All errors from the system.
 enum class ErrorType {
   None,
 
   // Parser Errors
+  ParserChainedExprInvalidArg,
+  ParserChainedExprNoCloseParen,
+  ParserChainedExprNoOpenParen,
   ParserFailedToBuildMatcher,
   ParserInvalidToken,
+  ParserMalformedChainedExpr,
   ParserNoCloseParen,
   ParserNoCode,
   ParserNoComma,
   ParserNoOpenParen,
   ParserNotAMatcher,
   ParserOverloadedType,
   ParserStringError,
   ParserTrailingCode,

mlir/include/mlir/Query/Matcher/MatchersInternal.h

public:
   static std::unique_ptr<DynMatcher>
   constructDynMatcherFromMatcherFn(MatcherFn &matcherFn) {
     auto impl = std::make_unique<MatcherFnImpl<MatcherFn>>(matcherFn);
     return std::make_unique<DynMatcher>(impl.release());
   }
 
   bool match(Operation *op) const { return implementation->match(op); }
 
+  void setFunctionName(StringRef name) { functionName = name.str(); };
+
+  bool hasFunctionName() const { return !functionName.empty(); };
+
+  StringRef getFunctionName() const { return functionName; };
+
 private:
   llvm::IntrusiveRefCntPtr<MatcherInterface> implementation;
+  std::string functionName;
 };
 
 } // namespace mlir::query::matcher
 
 #endif // MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERSINTERNAL_H

mlir/lib/Query/Matcher/Diagnostics.cpp

static llvm::StringRef errorTypeToFormatString(ErrorType type) {
   case ErrorType::ParserNotAMatcher:
     return "Input value is not a matcher expression.";
   case ErrorType::ParserInvalidToken:
     return "Invalid token <$0> found when looking for a value.";
   case ErrorType::ParserTrailingCode:
     return "Unexpected end of code.";
   case ErrorType::ParserOverloadedType:
     return "Input value has unresolved overloaded type: $0";
+  case ErrorType::ParserMalformedChainedExpr:
+    return "Period not followed by valid chained call.";
+  case ErrorType::ParserChainedExprInvalidArg:
+    return "Missing/Invalid argument for the chained call.";
+  case ErrorType::ParserChainedExprNoCloseParen:
+    return "Missing ')' for the chained call.";
+  case ErrorType::ParserChainedExprNoOpenParen:
+    return "Missing '(' for the chained call.";
   case ErrorType::ParserFailedToBuildMatcher:
     return "Failed to build matcher: $0.";
 
   case ErrorType::None:
     return "<N/A>";
   }
   llvm_unreachable("Unknown ErrorType value.");
 }

mlir/lib/Query/Matcher/Parser.h

public:
   // the Sema instance passed into parseMatcherExpression() to handle all
   // matcher tokens.
   class Sema {
   public:
     virtual ~Sema();
 
     // Process a matcher expression. The caller takes ownership of the Matcher
     // object returned.
-    virtual VariantMatcher
-    actOnMatcherExpression(MatcherCtor ctor, SourceRange nameRange,
-                           llvm::ArrayRef<ParserValue> args,
-                           Diagnostics *error) = 0;
+    virtual VariantMatcher actOnMatcherExpression(
+        MatcherCtor ctor, SourceRange nameRange, llvm::StringRef functionName,
+        llvm::ArrayRef<ParserValue> args, Diagnostics *error) = 0;
 
     // Look up a matcher by name in the matcher name found by the parser.
     virtual std::optional<MatcherCtor>
     lookupMatcherCtor(llvm::StringRef matcherName) = 0;
 
     // Compute the list of completion types for Context.
     virtual std::vector<ArgKind> getAcceptedCompletionTypes(
         llvm::ArrayRef<std::pair<MatcherCtor, unsigned>> Context);
public:
   public:
     RegistrySema(const Registry &matcherRegistry)
         : matcherRegistry(matcherRegistry) {}
     ~RegistrySema() override;
 
     std::optional<MatcherCtor>
     lookupMatcherCtor(llvm::StringRef matcherName) override;
 
-    VariantMatcher actOnMatcherExpression(MatcherCtor ctor,
-                                          SourceRange nameRange,
-                                          llvm::ArrayRef<ParserValue> args,
-                                          Diagnostics *error) override;
+    VariantMatcher actOnMatcherExpression(MatcherCtor Ctor,
+                                          SourceRange NameRange,
+                                          StringRef functionName,
+                                          ArrayRef<ParserValue> Args,
+                                          Diagnostics *Error) override;
 
     std::vector<ArgKind> getAcceptedCompletionTypes(
         llvm::ArrayRef<std::pair<MatcherCtor, unsigned>> context) override;
 
     std::vector<MatcherCompletion>
     getMatcherCompletions(llvm::ArrayRef<ArgKind> acceptedTypes) override;
 
   private:

 private:
   class CodeTokenizer;
   struct ScopedContextEntry;
   struct TokenInfo;
 
   Parser(CodeTokenizer *tokenizer, const Registry &matcherRegistry,
          const NamedValueMap *namedValues, Diagnostics *error);
 
+  bool parseChainedExpression(std::string &argument);
+
   bool parseExpressionImpl(VariantValue *value);
 
   bool parseMatcherArgs(std::vector<ParserValue> &args, MatcherCtor ctor,
                         const TokenInfo &nameToken, TokenInfo &endToken);
 
   bool parseMatcherExpressionImpl(const TokenInfo &nameToken,
                                   const TokenInfo &openToken,
                                   std::optional<MatcherCtor> ctor,

mlir/lib/Query/Matcher/Parser.cpp

struct Parser::TokenInfo {
   TokenInfo() = default;
 
   // Method to set the kind and text of the token
   void set(TokenKind newKind, llvm::StringRef newText) {
     kind = newKind;
     text = newText;
   }
 
+  // known identifiers.
+  static const char *const ID_Extract;
+
   llvm::StringRef text;
   TokenKind kind = TokenKind::Eof;
   SourceRange range;
   VariantValue value;
 };
 
+const char *const Parser::TokenInfo::ID_Extract = "extract";
+
 class Parser::CodeTokenizer {
 public:
   // Constructor with matcherCode and error
   explicit CodeTokenizer(llvm::StringRef matcherCode, Diagnostics *error)
       : code(matcherCode), startOfLine(matcherCode), error(error) {
     nextToken = getNextToken();
   }
 
bool Parser::parseIdentifierPrefixImpl(VariantValue *value) {
   }
 
   std::optional<MatcherCtor> ctor = sema->lookupMatcherCtor(nameToken.text);
 
   // Parse as a matcher expression.
   return parseMatcherExpressionImpl(nameToken, openToken, ctor, value);
 }
 
+bool Parser::parseChainedExpression(std::string &argument) {
+  // Parse the parenthesized argument to .extract("foo")
+  const TokenInfo openToken = tokenizer->consumeNextToken();
+  const TokenInfo argumentToken = tokenizer->consumeNextTokenIgnoreNewlines();
+  const TokenInfo closeToken = tokenizer->consumeNextTokenIgnoreNewlines();
+
+  if (openToken.kind != TokenKind::OpenParen) {
+    error->addError(openToken.range, ErrorType::ParserChainedExprNoOpenParen);
+    return false;
+  }
+
+  if (argumentToken.kind != TokenKind::Literal ||
+      !argumentToken.value.isString()) {
+    error->addError(argumentToken.range,
+                    ErrorType::ParserChainedExprInvalidArg);
+    return false;
+  }
+
+  if (closeToken.kind != TokenKind::CloseParen) {
+    error->addError(closeToken.range, ErrorType::ParserChainedExprNoCloseParen);
+    return false;
+  }
+
+  // If all checks passed, extract the argument and return true.
+  argument = argumentToken.value.getString();
+  return true;
+}
+
 // Parse the arguments of a matcher
 bool Parser::parseMatcherArgs(std::vector<ParserValue> &args, MatcherCtor ctor,
                               const TokenInfo &nameToken, TokenInfo &endToken) {
   ScopedContextEntry sce(this, ctor);
 
   while (tokenizer->nextTokenKind() != TokenKind::Eof) {
     if (tokenizer->nextTokenKind() == TokenKind::CloseParen) {
       // end of args.
bool Parser::parseMatcherExpressionImpl(const TokenInfo &nameToken,
 
   // Check for the missing closing parenthesis
   if (endToken.kind != TokenKind::CloseParen) {
     error->addError(openToken.range, ErrorType::ParserNoCloseParen)
         << nameToken.text;
     return false;
   }
 
+  std::string functionName;
+  if (tokenizer->peekNextToken().kind == TokenKind::Period) {
+    tokenizer->consumeNextToken();
+    TokenInfo chainCallToken = tokenizer->consumeNextToken();
+    if (chainCallToken.kind == TokenKind::CodeCompletion) {
+      addCompletion(chainCallToken, MatcherCompletion("extract(\"", "extract"));
+      return false;
+    }
+
+    if (chainCallToken.kind != TokenKind::Ident ||
+        chainCallToken.text != TokenInfo::ID_Extract) {
+      error->addError(chainCallToken.range,
+                      ErrorType::ParserMalformedChainedExpr);
+      return false;
+    }
+
+    if (chainCallToken.text == TokenInfo::ID_Extract &&
+        !parseChainedExpression(functionName))
+      return false;
+  }
+
   if (!ctor)
     return false;
   // Merge the start and end infos.
   SourceRange matcherRange = nameToken.range;
   matcherRange.end = endToken.range.end;
-  VariantMatcher result =
-      sema->actOnMatcherExpression(*ctor, matcherRange, args, error);
+  VariantMatcher result = sema->actOnMatcherExpression(
+      *ctor, matcherRange, functionName, args, error);
   if (result.isNull())
     return false;
   *value = result;
   return true;
 }
 
 // If the prefix of this completion matches the completion token, add it to
 // completions minus the prefix.

 Parser::RegistrySema::~RegistrySema() = default;
 
 std::optional<MatcherCtor>
 Parser::RegistrySema::lookupMatcherCtor(llvm::StringRef matcherName) {
   return RegistryManager::lookupMatcherCtor(matcherName, matcherRegistry);
 }
 
 VariantMatcher Parser::RegistrySema::actOnMatcherExpression(
-    MatcherCtor ctor, SourceRange nameRange, llvm::ArrayRef<ParserValue> args,
-    Diagnostics *error) {
-  return RegistryManager::constructMatcher(ctor, nameRange, args, error);
+    MatcherCtor ctor, SourceRange nameRange, llvm::StringRef functionName,
+    llvm::ArrayRef<ParserValue> args, Diagnostics *error) {
+  return RegistryManager::constructMatcher(ctor, nameRange, functionName, args,
+                                           error);
 }
 
 std::vector<ArgKind> Parser::RegistrySema::getAcceptedCompletionTypes(
     llvm::ArrayRef<std::pair<MatcherCtor, unsigned>> context) {
   return RegistryManager::getAcceptedCompletionTypes(context);
 }
 
 std::vector<MatcherCompletion> Parser::RegistrySema::getMatcherCompletions(

mlir/lib/Query/Matcher/RegistryManager.h

static std::vector<ArgKind> getAcceptedCompletionTypes(
       llvm::ArrayRef<std::pair<MatcherCtor, unsigned>> context);
 
   static std::vector<MatcherCompletion>
   getMatcherCompletions(ArrayRef<ArgKind> acceptedTypes,
                         const Registry &matcherRegistry);
 
   static VariantMatcher constructMatcher(MatcherCtor ctor,
                                          internal::SourceRange nameRange,
+                                         llvm::StringRef functionName,
                                          ArrayRef<ParserValue> args,
                                          internal::Diagnostics *error);
 };
 
 } // namespace mlir::query::matcher
 
 #endif // MLIR_TOOLS_MLIRQUERY_MATCHER_REGISTRYMANAGER_H

mlir/lib/Query/Matcher/RegistryManager.cpp

for (const auto &m : matcherRegistry.constructors()) {
     completions.emplace_back(typedText, os.str());
   }
 
   return completions;
 }
 
 VariantMatcher RegistryManager::constructMatcher(
     MatcherCtor ctor, internal::SourceRange nameRange,
-    llvm::ArrayRef<ParserValue> args, internal::Diagnostics *error) {
-  return ctor->create(nameRange, args, error);
+    llvm::StringRef functionName, llvm::ArrayRef<ParserValue> args,
+    internal::Diagnostics *error) {
+
+  VariantMatcher out = ctor->create(nameRange, args, error);
+  if (functionName.empty() || out.isNull())
+    return out;
+
+  std::optional<DynMatcher> result = out.getDynMatcher();
+
+  if (result.has_value()) {
+    result->setFunctionName(functionName);
+    return VariantMatcher::SingleMatcher(*result);
+  }
+
+  // error->addError(NameRange, error->ET_RegistryNotBindable);
+  return out;
 }
 
 } // namespace mlir::query::matcher

mlir/lib/Query/Query.cpp

 //===---- Query.cpp - -----------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Query/Query.h"
 #include "QueryParser.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/Query/Matcher/MatchFinder.h"
 #include "mlir/Query/QuerySession.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir::query {
 
static void printMatch(llvm::raw_ostream &os, QuerySession &qs, 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");
 }
 
+static Operation *extractFunction(std::vector<Operation *> &ops,
+                                  MLIRContext *context,
+                                  llvm::StringRef functionName) {
+  context->loadDialect<func::FuncDialect>();
+  OpBuilder builder(context);
+
+  // Collect data for function creation
+  std::vector<Operation *> slice;
+  std::vector<Value> values;
+  std::vector<Type> outputTypes;
+
+  for (auto *op : ops) {
+
+    if (!isa<func::ReturnOp>(op))
+      slice.push_back(op);
+
+    outputTypes.insert(outputTypes.end(), op->getResults().getTypes().begin(),
+                       op->getResults().getTypes().end());
+
+    values.insert(values.end(), op->getOperands().begin(),
+                  op->getOperands().end());
+  }
+
+  auto loc = builder.getUnknownLoc();
+
+  // Create the function
+  FunctionType funcType =
+      builder.getFunctionType(ValueRange(values), outputTypes);
+  func::FuncOp funcOp = func::FuncOp::create(loc, functionName, funcType);
+
+  builder.setInsertionPointToEnd(funcOp.addEntryBlock());
+  builder.setInsertionPointToEnd(&funcOp.getBody().front());
+
+  // Map original values to function arguments
+  IRMapping mapper;
+  for (const auto &arg : llvm::enumerate(values))
+    mapper.map(arg.value(), funcOp.getArgument(arg.index()));
+
+  // Clone operations and build function body
+  std::vector<Operation *> clonedOps;
+  for (Operation *slicedOp : slice)
+    clonedOps.push_back(builder.clone(*slicedOp, mapper));
+
+  // Remove unused function arguments
+  size_t currentIndex = 0;
+  while (currentIndex < funcOp.getNumArguments()) {
+    if (funcOp.getArgument(currentIndex).getUses().empty()) {
+      funcOp.eraseArgument(currentIndex);
+    } else {
+      currentIndex++;
+    }
+  }
+
+  // Collect cloned values
+  std::vector<Value> clonedVals;
+  for (Operation *clonedOp : clonedOps)
+    clonedVals.insert(clonedVals.end(), clonedOp->result_begin(),
+                      clonedOp->result_end());
+
+  // Add return operation
+  builder.create<func::ReturnOp>(loc, clonedVals);
+
+  return funcOp;
+}
+
 Query::~Query() = default;
 
 mlir::LogicalResult InvalidQuery::run(llvm::raw_ostream &os,
                                       QuerySession &qs) const {
   os << errStr << "\n";
   return mlir::failure();
 }
 

 mlir::LogicalResult QuitQuery::run(llvm::raw_ostream &os,
                                    QuerySession &qs) const {
   qs.terminate = true;
   return mlir::success();
 }
 
 mlir::LogicalResult MatchQuery::run(llvm::raw_ostream &os,
                                     QuerySession &qs) const {
+  Operation *rootOp = qs.getRootOp();
   int matchCount = 0;
   std::vector<Operation *> matches =
-      matcher::MatchFinder().getMatches(qs.getRootOp(), matcher);
+      matcher::MatchFinder().getMatches(rootOp, matcher);
+
+  if (matcher.hasFunctionName()) {
+    auto functionName = matcher.getFunctionName();
+    Operation *function =
+        extractFunction(matches, rootOp->getContext(), functionName);
+    os << "\n" << *function << "\n\n";
+    return mlir::success();
+  }
+
   os << "\n";
   for (Operation *op : matches) {
     os << "Match #" << ++matchCount << ":\n\n";
     // Placeholder "root" binding for the initial draft.
     printMatch(os, qs, op, "root");
   }
   os << matchCount << (matchCount == 1 ? " match.\n\n" : " matches.\n\n");
 
   return mlir::success();
 }
 
 } // namespace mlir::query

mlir/test/mlir-query/function-extraction.mlir

+// RUN: mlir-query %s -c "m hasOpName(\"arith.mulf\").extract(\"testmul\")" | FileCheck %s
+
+// CHECK: func.func @testmul({{.*}}) -> (f32, f32, f32) {
+// CHECK:       %[[MUL0:.*]] = arith.mulf {{.*}} : f32
+// CHECK:       %[[MUL1:.*]] = arith.mulf {{.*}}, %[[MUL0]] : f32
+// CHECK:       %[[MUL2:.*]] = arith.mulf {{.*}} : f32
+// CHECK-NEXT:  return %[[MUL0]], %[[MUL1]], %[[MUL2]] : f32, f32, f32
+
+func.func @mixedOperations(%a: f32, %b: f32, %c: f32) -> f32 {
+  %sum0 = arith.addf %a, %b : f32
+  %sub0 = arith.subf %sum0, %c : f32
+  %mul0 = arith.mulf %a, %sub0 : f32
+  %sum1 = arith.addf %b, %c : f32
+  %mul1 = arith.mulf %sum1, %mul0 : f32
+  %sub2 = arith.subf %mul1, %a : f32
+  %sum2 = arith.addf %mul1, %b : f32
+  %mul2 = arith.mulf %sub2, %sum2 : f32
+  return %mul2 : f32
+}