Add range and mask operations

Signed-off-by: Anton Korobeynikov <anton@korobeynikov.info>

include/p4mlir/Dialect/P4HIR/P4HIR_ParserOps.td

@@ -161,4 +161,82 @@ def UniversalSetOp : P4HIR_Op<"universal_set",
   ];
 }
 
+def RangeOp : P4HIR_Op<"range", [Pure,
+  TypesMatchWith<
+       "the resulting type of a range operation must be the set of left-hand side operand types",
+       "result", "lhs", "cast<P4HIR::SetType>($_self).getElementType()">,
+  TypesMatchWith<
+       "the resulting type of a range operation must be the set of right-hand side operand types",
+       "result", "rhs", "cast<P4HIR::SetType>($_self).getElementType()">,
+  ParentOneOf<["ParserSelectCaseOp"]>, // TBD: ForInOp
+  DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+
+  let summary = "Range operation";
+  let description = [{
+    p4hir.range represents a set formed by all values numerically between the first and the second, inclusively.
+
+    It requires two input operands and has one result, with input operands
+    being of any integer type.
+
+    ```mlir
+    %3 = p4hir.range(%1, %2) : !p4hir.set<!p4hir.bit<42>>
+    ```
+  }];
+
+  let results = (outs SetType:$result);
+  let arguments = (ins AnyIntP4Type:$lhs, AnyIntP4Type:$rhs);
+
+  let assemblyFormat = [{
+    `(` $lhs `,` $rhs `)` `:` type($result) attr-dict
+  }];
+
+  let builders = [
+    OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs), [{
+      auto resultType = SetType::get($_builder.getContext(), lhs.getType());
+      build($_builder, $_state, resultType, lhs, rhs);
+    }]>
+  ];
+}
+
+def MaskOp : P4HIR_Op<"mask", [Pure,
+  TypesMatchWith<
+       "the resulting type of a mask operation must be the set of left-hand side operand types",
+       "result", "lhs", "cast<P4HIR::SetType>($_self).getElementType()">,
+  TypesMatchWith<
+       "the resulting type of a mask operation must be the set of right-hand side operand types",
+       "result", "rhs", "cast<P4HIR::SetType>($_self).getElementType()">,
+  ParentOneOf<["ParserSelectCaseOp"]>,       
+  DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+
+  let summary = "Mask operation";
+  let description = [{
+    p4hir.mask represents a set formed by all values defined by the mask.
+    More formally, the set denoted by a &&& b is defined as follows:
+    ```
+    a &&& b = { c where a & b = c & b }
+    ```
+
+    It requires two input operands and has one result, with input operands
+    being of any integer type.
+
+    ```mlir
+    %3 = p4hir.mask(%1, %2) : !p4hir.set<!p4hir.bit<42>>
+    ```
+  }];
+
+  let results = (outs SetType:$result);
+  let arguments = (ins AnyIntP4Type:$lhs, AnyIntP4Type:$rhs);
+
+  let assemblyFormat = [{
+    `(` $lhs `,` $rhs `)` `:` type($result) attr-dict
+  }];
+
+  let builders = [
+    OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs), [{
+      auto resultType = SetType::get($_builder.getContext(), lhs.getType());
+      build($_builder, $_state, resultType, lhs, rhs);
+    }]>
+  ];
+}
+
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_PARSEROPS_TD

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -997,6 +997,10 @@ void P4HIR::TupleExtractOp::build(OpBuilder &builder, OperationState &odsState,
     build(builder, odsState, tupleType.getType(fieldIndex), input, fieldIndex);
 }
 
+//===----------------------------------------------------------------------===//
+// ParserOp
+//===----------------------------------------------------------------------===//
+
 void P4HIR::ParserOp::print(mlir::OpAsmPrinter &printer) {
     auto funcName = getSymNameAttr().getValue();
 
@@ -1174,6 +1178,10 @@ void P4HIR::SetProductOp::build(mlir::OpBuilder &builder, mlir::OperationState &
     result.addOperands(values);
 }
 
+//===----------------------------------------------------------------------===//
+// UniversalSetOp
+//===----------------------------------------------------------------------===//
+
 void P4HIR::UniversalSetOp::build(mlir::OpBuilder &builder, mlir::OperationState &result) {
     result.addTypes(P4HIR::SetType::get(P4HIR::DontcareType::get(builder.getContext())));
 }
@@ -1182,6 +1190,22 @@ void P4HIR::UniversalSetOp::getAsmResultNames(function_ref<void(Value, StringRef
     setNameFn(getResult(), "everything");
 }
 
+//===----------------------------------------------------------------------===//
+// RangeOp
+//===----------------------------------------------------------------------===//
+
+void P4HIR::RangeOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
+    setNameFn(getResult(), "range");
+}
+
+//===----------------------------------------------------------------------===//
+// MaskOp
+//===----------------------------------------------------------------------===//
+
+void P4HIR::MaskOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
+    setNameFn(getResult(), "mask");
+}
+
 namespace {
 struct P4HIROpAsmDialectInterface : public OpAsmDialectInterface {
     using OpAsmDialectInterface::OpAsmDialectInterface;

tools/p4mlir-translate/translate.cpp

@@ -384,6 +384,8 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
     HANDLE_IN_POSTORDER(Declaration_Variable)
     HANDLE_IN_POSTORDER(ReturnStatement)
     HANDLE_IN_POSTORDER(ArrayIndex)
+    HANDLE_IN_POSTORDER(Range)
+    HANDLE_IN_POSTORDER(Mask)
 
 #undef HANDLE_IN_POSTORDER
 
@@ -1415,6 +1417,22 @@ void P4HIRConverter::postorder(const P4::IR::ArrayIndex *arr) {
     BUG("cannot handle this array yet: %1%", arr);
 }
 
+void P4HIRConverter::postorder(const P4::IR::Range *range) {
+    auto lhs = getValue(range->left);
+    auto rhs = getValue(range->right);
+
+    auto loc = getLoc(builder, range);
+    setValue(range, builder.create<P4HIR::RangeOp>(loc, lhs, rhs).getResult());
+}
+
+void P4HIRConverter::postorder(const P4::IR::Mask *range) {
+    auto lhs = getValue(range->left);
+    auto rhs = getValue(range->right);
+
+    auto loc = getLoc(builder, range);
+    setValue(range, builder.create<P4HIR::MaskOp>(loc, lhs, rhs).getResult());
+}
+
 bool P4HIRConverter::preorder(const P4::IR::P4Parser *parser) {
     auto applyType = mlir::cast<P4HIR::FuncType>(getOrCreateType(parser->getApplyMethodType()));
     auto ctorType =