Add support for assignment statement

Signed-off-by: Anton Korobeynikov <[email protected]>
p4lang · Jan 31, 2025 · 2968abe · 2968abe
1 parent 2a42d0b
commit 2968abe
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/test/Translate/Ops/assign.p4 b/test/Translate/Ops/assign.p4
@@ -0,0 +1,30 @@
+// RUN: p4mlir-translate --typeinference-only %s | FileCheck %s
+
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+
+action assign() {
+    bit<10> res;
+
+    bit<10> lhs = 1;
+    bit<10> rhs = 2;
+
+    res = lhs;
+    res = lhs + rhs;
+}
+
+// CHECK-LABEL:   module
+// CHECK-NEXT:    %[[VAL_0:.*]] = p4hir.alloca !p4hir.bit<10> ["res"] : !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         %[[VAL_1:.*]] = p4hir.const #p4hir.int<1> : !p4hir.bit<10>
+// CHECK:         %[[VAL_2:.*]] = p4hir.cast(%[[VAL_1]] : !p4hir.bit<10>) : !p4hir.bit<10>
+// CHECK:         %[[VAL_3:.*]] = p4hir.alloca !p4hir.bit<10> ["lhs", init] : !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         p4hir.store %[[VAL_2]], %[[VAL_3]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         %[[VAL_4:.*]] = p4hir.const #p4hir.int<2> : !p4hir.bit<10>
+// CHECK:         %[[VAL_5:.*]] = p4hir.cast(%[[VAL_4]] : !p4hir.bit<10>) : !p4hir.bit<10>
+// CHECK:         %[[VAL_6:.*]] = p4hir.alloca !p4hir.bit<10> ["rhs", init] : !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         p4hir.store %[[VAL_5]], %[[VAL_6]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         %[[VAL_7:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         p4hir.store %[[VAL_7]], %[[VAL_0]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
+// CHECK:         %[[VAL_8:.*]] = p4hir.load %[[VAL_3]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_9:.*]] = p4hir.load %[[VAL_6]] : !p4hir.ref<!p4hir.bit<10>>, !p4hir.bit<10>
+// CHECK:         %[[VAL_10:.*]] = p4hir.binop(add, %[[VAL_8]], %[[VAL_9]]) : !p4hir.bit<10>
+// CHECK:         p4hir.store %[[VAL_10]], %[[VAL_0]] : !p4hir.bit<10>, !p4hir.ref<!p4hir.bit<10>>
diff --git a/tools/p4mlir-translate/translate.cpp b/tools/p4mlir-translate/translate.cpp
@@ -129,6 +129,7 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
 
     mlir::TypedAttr resolveConstant(const P4::IR::CompileTimeValue *ctv);
     mlir::TypedAttr resolveConstantExpr(const P4::IR::Expression *expr);
+    mlir::Value resolveReference(const P4::IR::Node *node);
 
  public:
     P4HIRConverter(mlir::OpBuilder &builder, const P4::TypeMap *typeMap)
@@ -295,7 +296,9 @@ class P4HIRConverter : public P4::Inspector, public P4::ResolutionContext {
     HANDLE_IN_POSTORDER(Declaration_Variable)
 
 #undef HANDLE_IN_POSTORDER
+
     bool preorder(const P4::IR::Declaration_Constant *decl) override;
+    bool preorder(const P4::IR::AssignmentStatement *assign) override;
 
     mlir::Value emitUnOp(const P4::IR::Operation_Unary *unop, P4HIR::UnaryOpKind kind);
     mlir::Value emitBinOp(const P4::IR::Operation_Binary *binop, P4HIR::BinOpKind kind);
@@ -347,6 +350,21 @@ bool P4TypeConverter::setType(const P4::IR::Type *type, mlir::Type mlirType) {
     return false;
 }
 
+mlir::Value P4HIRConverter::resolveReference(const P4::IR::Node *node) {
+    // If this is a PathExpression, resolve it
+    if (const auto *pe = node->to<P4::IR::PathExpression>()) {
+        node = resolvePath(pe->path, false)->checkedTo<P4::IR::Declaration>();
+    }
+
+    // The result is expected to be an l-value
+    auto ref = p4Values.lookup(node);
+    BUG_CHECK(ref, "expected %1% (aka %2%) to be converted", node, dbp(node));
+    BUG_CHECK(mlir::isa<P4HIR::ReferenceType>(ref.getType()),
+              "expected reference type for node %1%", node);
+
+    return ref;
+}
+
 mlir::TypedAttr P4HIRConverter::resolveConstant(const P4::IR::CompileTimeValue *ctv) {
     BUG("cannot resolve this constant yet %1%", ctv);
 }
@@ -495,6 +513,17 @@ CONVERT_BINOP(BXor, Xor);
 
 #undef CONVERT_BINOP
 
+bool P4HIRConverter::preorder(const P4::IR::AssignmentStatement *assign) {
+    ConversionTracer trace("Converting ", assign);
+
+    // TODO: Handle slice on LHS here
+    visit(assign->left);
+    visit(assign->right);
+    auto ref = resolveReference(assign->left);
+    builder.create<P4HIR::StoreOp>(getLoc(builder, assign), getValue(assign->right), ref);
+    return false;
+}
+
 }  // namespace
 
 namespace P4::P4MLIR {