Add ParserTransitionSelectOp::canonicalize and implement equivalent of SimplifySelectCases

Signed-off-by: Manolis Tsamis <tsamismanolis@gmail.com>

Author: mtsamis

include/p4mlir/Dialect/P4HIR/P4HIR_ParserOps.td

@@ -172,6 +172,7 @@ def ParserTransitionSelectOp : P4HIR_Op<"transition_select",
      return { StateIterator(selectCases.begin()), StateIterator(selectCases.end()) };
    }
  }];
+ let hasCanonicalizeMethod = 1;
 }
 
 def ParserSelectCaseOp : P4HIR_Op<"select_case",

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -1796,6 +1796,52 @@ mlir::ValueRange P4HIR::ParserSelectCaseOp::getSelectKeys() {
     return yield.getArgs();
 }
 
+//===----------------------------------------------------------------------===//
+// ParserTransitionSelectOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult P4HIR::ParserTransitionSelectOp::canonicalize(P4HIR::ParserTransitionSelectOp op,
+                                                            PatternRewriter &rewriter) {
+    mlir::Block *body = &op.getBody().back();
+    auto selectCases = op.selects();
+    auto it = llvm::find_if(selectCases, [](auto op) { return op.isDefault(); });
+
+    // No default case found.
+    if (it == selectCases.end()) return mlir::failure();
+
+    LogicalResult result = mlir::failure();
+
+    // Remove unreachable cases after last default case.
+    if (std::next(it) != selectCases.end()) {
+        mlir::Block *unreachableCases = rewriter.splitBlock(body, ++Block::iterator(*it));
+        rewriter.eraseBlock(unreachableCases);
+        result = mlir::success();
+    }
+
+    // Canonicalize tuple of universal sets to single universal set.
+    auto defaultCase = mlir::cast<P4HIR::ParserSelectCaseOp>(body->back());
+    auto defaultYield = mlir::cast<P4HIR::YieldOp>(defaultCase.getTerminator());
+    if (defaultYield.getArgs().size() > 1) {
+        rewriter.modifyOpInPlace(defaultYield, [&]() {
+            auto universalSetAttr = P4HIR::UniversalSetAttr::get(rewriter.getContext());
+            auto universalSet =
+                rewriter.create<P4HIR::ConstOp>(defaultYield.getLoc(), universalSetAttr);
+            defaultYield.getArgsMutable().assign(universalSet);
+        });
+
+        result = mlir::success();
+    }
+
+    // Replace select with single default case with direct transition.
+    if (body->getOperations().size() == 1) {
+        auto firstCase = *op.selects().begin();
+        rewriter.replaceOpWithNewOp<P4HIR::ParserTransitionOp>(op, firstCase.getState());
+        result = mlir::success();
+    }
+
+    return result;
+}
+
 //===----------------------------------------------------------------------===//
 // SetOp
 //===----------------------------------------------------------------------===//

test/Transforms/Folds/select.mlir

@@ -0,0 +1,80 @@
+// RUN: p4mlir-opt --canonicalize %s | FileCheck %s
+
+!b8i = !p4hir.bit<8>
+#everything = #p4hir.universal_set : !p4hir.set<!p4hir.dontcare>
+#false = #p4hir.bool<false> : !p4hir.bool
+#true = #p4hir.bool<true> : !p4hir.bool
+#int1_b8i = #p4hir.int<1> : !b8i
+#int2_b8i = #p4hir.int<2> : !b8i
+#int8_b8i = #p4hir.int<8> : !b8i
+#set_const_of_false = #p4hir.set<const : [#false]> : !p4hir.set<!p4hir.bool>
+#set_const_of_true = #p4hir.set<const : [#true]> : !p4hir.set<!p4hir.bool>
+#set_const_of_int1_b8i = #p4hir.set<const : [#int1_b8i]> : !p4hir.set<!b8i>
+#set_const_of_int2_b8i = #p4hir.set<const : [#int2_b8i]> : !p4hir.set<!b8i>
+#set_const_of_int8_b8i = #p4hir.set<const : [#int8_b8i]> : !p4hir.set<!b8i>
+
+// CHECK-LABEL: module
+module {
+  p4hir.parser @p1(%arg0: !b8i, %arg1: !p4hir.bool)() {
+    %set = p4hir.const #set_const_of_int8_b8i
+    %set_0 = p4hir.const #set_const_of_true
+    %set_1 = p4hir.const #set_const_of_int2_b8i
+    %set_2 = p4hir.const #set_const_of_false
+    %set_3 = p4hir.const #set_const_of_int1_b8i
+    %everything = p4hir.const #everything
+
+    // CHECK-LABEL: p4hir.state @start
+    p4hir.state @start {
+      // CHECK-COUNT-2: p4hir.select_case
+      // CHECK-NOT: p4hir.yield %everything, %everything : !p4hir.set<!p4hir.dontcare>, !p4hir.set<!p4hir.dontcare>
+      // CHECK: p4hir.yield %everything : !p4hir.set<!p4hir.dontcare>
+      // Remove dead cases; canonicalize default cases.
+      p4hir.transition_select %arg0, %arg1 : !b8i, !p4hir.bool {
+        p4hir.select_case {
+          p4hir.yield %set_3, %set_2 : !p4hir.set<!b8i>, !p4hir.set<!p4hir.bool>
+        } to @p1::@reject
+        p4hir.select_case {
+          p4hir.yield %everything, %everything : !p4hir.set<!p4hir.dontcare>, !p4hir.set<!p4hir.dontcare>
+        } to @p1::@next
+        p4hir.select_case {
+          p4hir.yield %set_1, %set_0 : !p4hir.set<!b8i>, !p4hir.set<!p4hir.bool>
+        } to @p1::@reject
+      }
+    }
+
+    // CHECK-LABEL: p4hir.state @next
+    p4hir.state @next {
+      // CHECK-NOT: p4hir.select_case
+      // CHECK: p4hir.transition to @p1::@final
+      // Replace select with one default case with direct transition.
+      p4hir.transition_select %arg0, %arg1 : !b8i, !p4hir.bool {
+        p4hir.select_case {
+          p4hir.yield %everything, %everything : !p4hir.set<!p4hir.dontcare>, !p4hir.set<!p4hir.dontcare>
+        } to @p1::@final
+        p4hir.select_case {
+          p4hir.yield %everything : !p4hir.set<!p4hir.dontcare>
+        } to @p1::@reject
+        p4hir.select_case {
+          p4hir.yield %set, %set_2 : !p4hir.set<!b8i>, !p4hir.set<!p4hir.bool>
+        } to @p1::@reject
+        p4hir.select_case {
+          p4hir.yield %everything : !p4hir.set<!p4hir.dontcare>
+        } to @p1::@reject
+      }
+    }
+
+    p4hir.state @final {
+      p4hir.transition to @p1::@accept
+    }
+
+    p4hir.state @accept {
+      p4hir.parser_accept
+    }
+
+    p4hir.state @reject {
+      p4hir.parser_reject
+    }
+
+    p4hir.transition to @p1::@start
+  }
+}