MLIR-QUERY: backwardSlice, forwardSlice & QueryOptions added

Author: dbudii

mlir/lib/Query/Matcher/Parser.cpp

@@ -293,13 +293,19 @@ bool Parser::parseIdentifierPrefixImpl(VariantValue *value) {
 
   if (tokenizer->nextTokenKind() != TokenKind::OpenParen) {
     // Parse as a named value.
-    auto namedValue =
-        namedValues ? namedValues->lookup(nameToken.text) : VariantValue();
+    if (auto namedValue = namedValues ? namedValues->lookup(nameToken.text)
+                                      : VariantValue()) {
 
-    if (!namedValue.isMatcher()) {
-      error->addError(tokenizer->peekNextToken().range,
-                      ErrorType::ParserNotAMatcher);
-      return false;
+      if (tokenizer->nextTokenKind() != TokenKind::Period) {
+        *value = namedValue;
+        return true;
+      }
+
+      if (!namedValue.isMatcher()) {
+        error->addError(tokenizer->peekNextToken().range,
+                        ErrorType::ParserNotAMatcher);
+        return false;
+      }
     }
 
     if (tokenizer->nextTokenKind() == TokenKind::NewLine) {

mlir/lib/Query/Query.cpp

@@ -127,6 +127,15 @@ LogicalResult QuitQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
   return mlir::success();
 }
 
+LogicalResult LetQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
+  if (value.hasValue()) {
+    qs.namedValues[name] = value;
+  } else {
+    qs.namedValues.erase(name);
+  }
+  return mlir::success();
+}
+
 LogicalResult MatchQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
   Operation *rootOp = qs.getRootOp();
 

mlir/lib/Query/QueryParser.cpp

@@ -166,6 +166,8 @@ QueryRef QueryParser::doParse() {
           .Case("", ParsedQueryKind::NoOp)
           .Case("#", ParsedQueryKind::Comment, /*isCompletion=*/false)
           .Case("help", ParsedQueryKind::Help)
+          .Case("l", ParsedQueryKind::Let, /*isCompletion=*/false)
+          .Case("let", ParsedQueryKind::Let, /*isCompletion=*/false)
           .Case("m", ParsedQueryKind::Match, /*isCompletion=*/false)
           .Case("set", ParsedQueryKind::Set)
           .Case("match", ParsedQueryKind::Match)
@@ -188,6 +190,27 @@ QueryRef QueryParser::doParse() {
   case ParsedQueryKind::Quit:
     return endQuery(new QuitQuery);
 
+  case ParsedQueryKind::Let: {
+    llvm::StringRef name = lexWord();
+
+    if (name.empty()) {
+      return new InvalidQuery("expected variable name");
+    }
+
+    if (completionPos) {
+      return completeMatcherExpression();
+    }
+
+    matcher::internal::Diagnostics diag;
+    matcher::VariantValue value;
+    if (!matcher::internal::Parser::parseExpression(
+            line, qs.getRegistryData(), &qs.namedValues, &value, &diag)) {
+      return makeInvalidQueryFromDiagnostics(diag);
+    }
+    QueryRef query = new LetQuery(name, value);
+    query->remainingContent = line;
+    return query;
+  }
   case ParsedQueryKind::Match: {
     if (completionPos) {
       return completeMatcherExpression();
@@ -204,6 +227,7 @@ QueryRef QueryParser::doParse() {
     }
     auto actualSource = origMatcherSource.substr(0, origMatcherSource.size() -
                                                         matcherSource.size());
+
     QueryRef query = new MatchQuery(actualSource, *matcher);
     query->remainingContent = matcherSource;
     return query;

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

@@ -4,7 +4,7 @@
 // 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
+// CHECK-NEXT:  return %[[MUL0]], %[[MUL1]], %[[MUL2]] : f32, f32, f32S
 
 func.func @mixedOperations(%a: f32, %b: f32, %c: f32) -> f32 {
   %sum0 = arith.addf %a, %b : f32
@@ -16,4 +16,4 @@ func.func @mixedOperations(%a: f32, %b: f32, %c: f32) -> f32 {
   %sum2 = arith.addf %mul1, %b : f32
   %mul2 = arith.mulf %sub2, %sum2 : f32
   return %mul2 : f32
-}
+    }

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

@@ -10,7 +10,6 @@
 // of the registered queries.
 //
 //===----------------------------------------------------------------------===//
-
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Matchers.h"