[Comb] Add support for header assignment.

p4_symbolic/bmv2/bmv2.proto

@@ -413,6 +413,8 @@ enum StatementOp {
   remove_header = 5;
   // Terminating the execution of all blocks currently executing.
   exit = 6;
+  // Assign a header instance to another of the same type.
+  assign_header = 7;
 }
 
 // This specifies the exact expressions that are supported by our ir.

p4_symbolic/ir/BUILD.bazel

@@ -67,6 +67,7 @@ cc_library(
         ":ir_cc_proto",
         ":table_entries",
         "//gutil:status",
+        "//p4_pdpi:ir_cc_proto",
         "//p4_symbolic/bmv2:bmv2_cc_proto",
         "@com_github_google_glog//:glog",
         "@com_github_p4lang_p4runtime//:p4info_cc_proto",

p4_symbolic/ir/ir.cc

@@ -36,6 +36,7 @@
 #include "google/protobuf/struct.pb.h"
 #include "gutil/status.h"
 #include "p4/config/v1/p4info.pb.h"
+#include "p4_pdpi/ir.pb.h"
 #include "p4_symbolic/bmv2/bmv2.pb.h"
 #include "p4_symbolic/ir/cfg.h"
 #include "p4_symbolic/ir/ir.pb.h"
@@ -65,14 +66,15 @@ absl::StatusOr<bmv2::StatementOp> StatementOpToEnum(const std::string &op) {
   static const std::unordered_map<std::string, bmv2::StatementOp> op_table = {
       {"assign", bmv2::StatementOp::assign},
       {"mark_to_drop", bmv2::StatementOp::mark_to_drop},
-      {"clone_ingress_pkt_to_egress",  // clone3(...)
+      {"clone_ingress_pkt_to_egress", // clone3(...)
        bmv2::StatementOp::clone_ingress_pkt_to_egress},
-      {"modify_field_with_hash_based_offset",  // hash(...)
+      {"modify_field_with_hash_based_offset", // hash(...)
        bmv2::StatementOp::modify_field_with_hash_based_offset},
-      {"add_header",  // hdr.SetValid()
+      {"add_header", // hdr.SetValid()
        bmv2::StatementOp::add_header},
-      {"remove_header",  // hdr.SetInvalid()
+      {"remove_header", // hdr.SetInvalid()
        bmv2::StatementOp::remove_header},
+      {"assign_header", bmv2::StatementOp::assign_header},
       {"exit", bmv2::StatementOp::exit}};
 
   if (op_table.count(op) != 1) {
@@ -661,6 +663,32 @@ absl::StatusOr<Statement> ExtractStatement(
       return statement;
     }
 
+    case bmv2::StatementOp::assign_header: {
+      HeaderAssignmentStatement &header_assignment =
+          *statement.mutable_header_assignment();
+      const google::protobuf::Value &params =
+          action_primitive.fields().at("parameters");
+      if (!params.has_list_value() || params.list_value().values_size() != 2) {
+        return gutil::InvalidArgumentErrorBuilder()
+               << "Header assignment must contain 2 parameters, found "
+               << action_primitive.DebugString();
+      }
+
+      ASSIGN_OR_RETURN(RValue left, ExtractRValue(params.list_value().values(0),
+                                                  parameter_map));
+      ASSIGN_OR_RETURN(
+          RValue right,
+          ExtractRValue(params.list_value().values(1), parameter_map));
+      if (!left.has_header_value() || !right.has_header_value()) {
+        return gutil::InvalidArgumentErrorBuilder()
+               << "Header assignment must be passed header instances, found "
+               << action_primitive.DebugString();
+      }
+      header_assignment.set_allocated_left(left.release_header_value());
+      header_assignment.set_allocated_right(right.release_header_value());
+      return statement;
+    }
+
     case bmv2::StatementOp::exit: {
       statement.mutable_exit();
       return statement;

p4_symbolic/ir/ir.proto

@@ -327,6 +327,7 @@ message Statement {
     CloneStatement clone = 4;
     HashStatement hash = 5;
     ExitStatement exit = 6;
+    HeaderAssignmentStatement header_assignment = 7;
   }
 }
 
@@ -365,6 +366,13 @@ message HashStatement {
 // Statement that terminates the execution of all blocks currently executing.
 message ExitStatement {}
 
+// An assignment statement of the form <left> = <right>, where both <left> and
+// <right> describe a header instance.
+message HeaderAssignmentStatement {
+  HeaderValue left = 1;
+  HeaderValue right = 2;
+}
+
 // The structure of an abstract RValue.
 message RValue {
   oneof rvalue {

p4_symbolic/symbolic/action.cc

@@ -45,24 +45,25 @@ namespace action {
 
 absl::Status EvaluateStatement(const ir::Statement &statement,
                                SymbolicPerPacketState &headers,
-                               ActionContext *context, z3::context &z3_context,
+                               SolverState &state, ActionContext *context,
                                const z3::expr &guard) {
   switch (statement.statement_case()) {
     case ir::Statement::kAssignment: {
-       return EvaluateAssignmentStatement(statement.assignment(), headers,
-                                         context, z3_context, guard);
+      return EvaluateAssignmentStatement(statement.assignment(), headers,
+                                         context, *state.context.z3_context,
+                                         guard);
     }
     case ir::Statement::kClone: {
       // TODO: Add support for cloning.
       return headers.Set(std::string(kGotClonedPseudoField),
-                         z3_context.bool_val(true), guard);
+                         state.context.z3_context->bool_val(true), guard);
     }
     case ir::Statement::kDrop: {
       // https://github.com/p4lang/p4c/blob/7ee76d16da63883c5092ab0c28321f04c2646759/p4include/v1model.p4#L435
       const std::string &header_name = statement.drop().header().header_name();
-      RETURN_IF_ERROR(
-          headers.Set(absl::StrFormat("%s.egress_spec", header_name),
-                      v1model::EgressSpecDroppedValue(z3_context), guard));
+      RETURN_IF_ERROR(headers.Set(
+          absl::StrFormat("%s.egress_spec", header_name),
+          v1model::EgressSpecDroppedValue(*state.context.z3_context), guard));
       return absl::OkStatus();
     }
     case ir::Statement::kHash: {
@@ -86,6 +87,10 @@ absl::Status EvaluateStatement(const ir::Statement &statement,
       LOG(ERROR) << "exit statement ignored since it is unsupported";
       return absl::OkStatus();
     }
+    case ir::Statement::kHeaderAssignment: {
+      return EvaluateHeaderAssignmentStatement(statement.header_assignment(),
+                                               headers, state, guard);
+    }
     case ir::Statement::STATEMENT_NOT_SET:
       break;
   }
@@ -94,6 +99,57 @@ absl::Status EvaluateStatement(const ir::Statement &statement,
       statement.DebugString()));
 }
 
+absl::Status EvaluateHeaderAssignmentStatement(
+    const ir::HeaderAssignmentStatement &statement,
+    SymbolicPerPacketState &headers, SolverState &state,
+    const z3::expr &guard) {
+  const std::string &left_header_name = statement.left().header_name();
+  const std::string &right_header_name = statement.right().header_name();
+  auto left_it = state.program.headers().find(left_header_name);
+  auto right_it = state.program.headers().find(right_header_name);
+  if (left_it == state.program.headers().end()) {
+    return gutil::NotFoundErrorBuilder()
+           << "Header '" << left_header_name << "' not found";
+  }
+  if (right_it == state.program.headers().end()) {
+    return gutil::NotFoundErrorBuilder()
+           << "Header '" << right_header_name << "' not found";
+  }
+  const ir::HeaderType &left_header = left_it->second;
+  const ir::HeaderType &right_header = right_it->second;
+
+  // Check the two headers have the same header type.
+  if (left_header.id() != right_header.id()) {
+    return gutil::InvalidArgumentErrorBuilder()
+           << "Headers '" << left_header_name << "' and '" << right_header_name
+           << "' have different header types. Header assignments expect the "
+              "same header type.";
+  }
+
+  // Assign the value of each header field from right to left.
+  for (const std::string &field_name : right_header.ordered_fields_list()) {
+    // Check if the `field_name` also exists in the left header.
+    if (!left_header.fields().contains(field_name)) {
+      return gutil::InvalidArgumentErrorBuilder()
+             << "Field '" << field_name << "' not found in header '"
+             << left_header_name << "' during header assignment.";
+    }
+
+    ASSIGN_OR_RETURN(z3::expr right_value,
+                     headers.Get(right_header_name, field_name));
+    RETURN_IF_ERROR(
+        headers.Set(left_header_name, field_name, right_value, guard));
+  }
+
+  // Assign the header validity from right to left.
+  ASSIGN_OR_RETURN(z3::expr right_valid,
+                   headers.Get(right_header_name, kValidPseudoField));
+  RETURN_IF_ERROR(
+      headers.Set(left_header_name, kValidPseudoField, right_valid, guard));
+
+  return absl::OkStatus();
+}
+
 // Constructs a symbolic expression for the assignment value, and either
 // constrains it in an enclosing assignment expression, or stores it in
 // the action scope.
@@ -335,11 +391,11 @@ absl::Status EvaluateConcreteAction(
                      " called with incompatible number of parameters"));
   }
 
-  // Find each parameter value in arguments by argument name. We should not rely
-  // on argument order matching param definition order, because the P4 runtime
-  // spec does not enforce this assumption in implementations, and furthermore
-  // the spec explicitly states that read entries do not have to preserve the
-  // order of repeated fields in written entries.
+  // Find each parameter value in arguments by argument name. We should not
+  // rely on argument order matching param definition order, because the P4
+  // runtime spec does not enforce this assumption in implementations, and
+  // furthermore the spec explicitly states that read entries do not have to
+  // preserve the order of repeated fields in written entries.
   for (const auto &arg : args) {
     absl::string_view arg_name = arg.name();
     ASSIGN_OR_RETURN(const pdpi::IrActionDefinition::IrActionParamDefinition
@@ -356,8 +412,8 @@ absl::Status EvaluateConcreteAction(
 
   // Iterate over the body in order, and evaluate each statement.
   for (const auto &statement : action.action_implementation().action_body()) {
-    RETURN_IF_ERROR(EvaluateStatement(statement, headers, &context,
-                                      *state.context.z3_context, guard));
+    RETURN_IF_ERROR(
+        EvaluateStatement(statement, headers, state, &context, guard));
   }
 
   return absl::OkStatus();
@@ -387,8 +443,8 @@ absl::Status EvaluateSymbolicAction(const ir::Action &action,
 
 // Iterate over the body in order, and evaluate each statement.
   for (const auto &statement : action.action_implementation().action_body()) {
-    RETURN_IF_ERROR(EvaluateStatement(statement, headers, &context,
-                                      *state.context.z3_context, guard));
+    RETURN_IF_ERROR(
+        EvaluateStatement(statement, headers, state, &context, guard));
   }
 
  return absl::OkStatus();

p4_symbolic/symbolic/action.h

@@ -50,9 +50,16 @@ struct ActionContext {
 // appropriate function.
 absl::Status EvaluateStatement(const ir::Statement &statement,
                                SymbolicPerPacketState &headers,
-                               ActionContext *context, z3::context &z3_context,
+                               SolverState &state, ActionContext *context,
                                const z3::expr &guard);
 
+// Assigns the values of all header fields of the RHS header instance to the
+// corresponding header fields of the LHS header instance. Returns an error if
+// the headers have different header types.
+absl::Status EvaluateHeaderAssignmentStatement(
+    const ir::HeaderAssignmentStatement &statement,
+    SymbolicPerPacketState &headers, SolverState &state, const z3::expr &guard);
+
 // Constructs a symbolic expression for the assignment value, and either
 // constrains it in an enclosing assignment expression, or stores it in
 // the action scope.

sai_p4/fixed/tunnel_termination.p4

@@ -78,20 +78,12 @@ control tunnel_termination(inout headers_t headers,
       headers.ipv6.setInvalid();
       if (headers.inner_ipv4.isValid()) {
         headers.ethernet.ether_type = ETHERTYPE_IPV4;
-        // TODO: Support header assignment in
-        // p4-symbolic and remove this guard.
-#ifndef PLATFORM_P4SYMBOLIC
         headers.ipv4 = headers.inner_ipv4;
-#endif
         headers.inner_ipv4.setInvalid();
       }
       if (headers.inner_ipv6.isValid()) {
         headers.ethernet.ether_type = ETHERTYPE_IPV6;
-        // TODO: Support header assignment in
-        // p4-symbolic and remove this guard.
-#ifndef PLATFORM_P4SYMBOLIC
         headers.ipv6 = headers.inner_ipv6;
-#endif
         headers.inner_ipv6.setInvalid();
       }
     }