[multicast] Add multicast replication support for (upcoming) softnpu/a4x2

codegen/rust/src/expression.rs

@@ -111,22 +111,37 @@ impl<'a> ExpressionGenerator<'a> {
             }
             ExpressionKind::Index(lval, xpr) => {
                 let mut ts = self.generate_lvalue(lval);
-                ts.extend(self.generate_expression(xpr.as_ref()));
+                // When the index expression is a slice (e.g. `field[31:28]`),
+                // we need the parent field's bit width to adjust for the
+                // byte reversal in header.rs. The HLIR resolves lvalue types
+                // during type checking, so we look up the width here and
+                // pass it to generate_slice. For non-bit types (or if the
+                // lvalue is missing from the HLIR, which should not happen
+                // for well-typed programs), we fall back to width 0 which
+                // skips the adjustment.
+                if let ExpressionKind::Slice(begin, end) = &xpr.kind {
+                    let field_width = self
+                        .hlir
+                        .lvalue_decls
+                        .get(lval)
+                        .map(|ni| match &ni.ty {
+                            p4::ast::Type::Bit(w)
+                            | p4::ast::Type::Varbit(w)
+                            | p4::ast::Type::Int(w) => *w,
+                            _ => 0,
+                        })
+                        .unwrap_or(0);
+                    ts.extend(self.generate_slice(begin, end, field_width));
+                } else {
+                    ts.extend(self.generate_expression(xpr.as_ref()));
+                }
                 ts
             }
             ExpressionKind::Slice(begin, end) => {
-                let l = match &begin.kind {
-                    ExpressionKind::IntegerLit(v) => *v as usize,
-                    _ => panic!("slice ranges can only be integer literals"),
-                };
-                let l = l + 1;
-                let r = match &end.kind {
-                    ExpressionKind::IntegerLit(v) => *v as usize,
-                    _ => panic!("slice ranges can only be integer literals"),
-                };
-                quote! {
-                    [#r..#l]
-                }
+                // Bare slice outside ExpressionKind::Index context (should not
+                // occur in well-typed programs per hlir.rs validation). No
+                // field width available, so no byte-reversal adjustment.
+                self.generate_slice(begin, end, 0)
             }
             ExpressionKind::Call(call) => {
                 let lv: Vec<TokenStream> = call
@@ -158,6 +173,76 @@ impl<'a> ExpressionGenerator<'a> {
         }
     }
 
+    /// Generate a bitvec range for a P4 bit slice `[hi:lo]`.
+    ///
+    /// # Confused-endian byte reversal
+    ///
+    /// header.rs stores multi-byte fields (width > 8 bits) with bytes
+    /// reversed relative to wire order. For a 32-bit field containing
+    /// wire bytes `[0xE0, 0x01, 0x01, 0x01]` (224.1.1.1), the bitvec
+    /// holds `[0x01, 0x01, 0x01, 0xE0]`. Bit positions within each
+    /// byte are preserved (Msb0) and only the byte order is flipped.
+    ///
+    /// P4 uses MSB-first bit numbering: in a W-bit field, bit W-1 is
+    /// the MSB (first on wire) and bit 0 is the LSB. A naive mapping
+    /// from P4 bit `b` to bitvec index `W-1-b` is correct for wire
+    /// order but wrong after byte reversal.
+    ///
+    /// The correct mapping for a reversed W-bit field:
+    ///
+    /// ```text
+    /// wire_idx     = W - 1 - b           // P4 bit -> wire-order bitvec index
+    /// wire_byte    = wire_idx / 8        // which byte on the wire
+    /// bit_in_byte  = wire_idx % 8        // position within that byte (Msb0)
+    /// storage_byte = W/8 - 1 - wire_byte // byte reversal
+    /// bitvec_idx   = storage_byte * 8 + bit_in_byte
+    /// ```
+    ///
+    /// # Why full-byte MSB-aligned slices worked without this previously
+    ///
+    /// For a slice like `[127:120]` on a 128-bit field (extracting wire
+    /// byte 0), the mapping sends byte 0 to storage byte 15, producing
+    /// bitvec range `[120..128]`. The naive codegen also emits
+    /// `[120..128]`. The numeric coincidence holds whenever the slice
+    /// spans complete bytes at the MSB end: the reversed byte position
+    /// `W - 8 - start` equals the original `end` index. This breaks
+    /// for sub-byte slices (e.g. nibble extraction) or slices not
+    /// aligned to the MSB boundary.
+    ///
+    /// This confusion was discovered while implementing multicast support.
+    /// Multicast routing classifies packets by extracting the top nibble
+    /// of `ipv4.dst` via `ipv4.dst[31:28]` to check for the 0xE prefix
+    /// (IPv4 multicast range 224.0.0.0/4). The naive codegen emitted
+    /// `bitvec[28..32]`, which reads the low nibble of storage byte 3
+    /// (= 0x0 for 224.x.x.x) instead of the high nibble (= 0xE). The
+    /// correct range after byte-reversal adjustment is `bitvec[24..28]`.
+    fn generate_slice(
+        &self,
+        begin: &Expression,
+        end: &Expression,
+        field_width: FieldWidth,
+    ) -> TokenStream {
+        let hi: P4Bit = match &begin.kind {
+            ExpressionKind::IntegerLit(v) => *v as usize,
+            _ => panic!("slice ranges can only be integer literals"),
+        };
+        let lo: P4Bit = match &end.kind {
+            ExpressionKind::IntegerLit(v) => *v as usize,
+            _ => panic!("slice ranges can only be integer literals"),
+        };
+
+        if field_width > 8 {
+            let (r, l) = reversed_slice_range(hi, lo, field_width);
+            quote! { [#r..#l] }
+        } else {
+            // Fields <= 8 bits are not byte-reversed by header.rs,
+            // so the naive P4-to-bitvec mapping is correct.
+            let l = hi + 1;
+            let r = lo;
+            quote! { [#r..#l] }
+        }
+    }
+
     pub(crate) fn generate_bit_literal(
         &self,
         width: u16,
@@ -223,3 +308,150 @@ impl<'a> ExpressionGenerator<'a> {
         }
     }
 }
+
+/// P4 bit position (MSB-first index within a field).
+type P4Bit = usize;
+
+/// Width of a P4 header field in bits.
+type FieldWidth = usize;
+
+/// Half-open bitvec range `(start, end)` into the storage representation.
+type BitvecRange = (usize, usize);
+
+/// Compute the bitvec range `(start, end)` for a P4 slice `[hi:lo]` on a
+/// byte-reversed field of the given width.
+///
+/// header.rs reverses byte order for multi-byte fields. The mapping from
+/// P4 bit positions to storage positions depends on whether the slice
+/// stays within a single wire byte or spans multiple bytes.
+///
+/// For a single-byte slice, we map each endpoint through the byte reversal:
+/// the target byte moves to a new position but bit ordering within the byte
+/// is preserved (Msb0). This handles sub-byte nibble extractions like
+/// `ipv4.dst[31:28]`.
+///
+/// For a multi-byte slice, the byte reversal makes the endpoints
+/// non-contiguous (byte 0 maps to the far end, byte 1 maps next to it,
+/// etc.). However, if the slice is byte-aligned, the reversed bytes form
+/// a contiguous block at a different offset. We compute the storage byte
+/// range directly. Non-byte-aligned multi-byte slices cannot be represented
+/// as a single contiguous range after reversal and will panic.
+fn reversed_slice_range(
+    hi: P4Bit,
+    lo: P4Bit,
+    field_width: FieldWidth,
+) -> BitvecRange {
+    // Wire byte indices for the slice endpoints. P4 bit W-1 is in wire
+    // byte 0 (MSB-first), so higher bit numbers map to lower byte indices.
+    let wire_byte_hi = (field_width - 1 - hi) / 8;
+    let wire_byte_lo = (field_width - 1 - lo) / 8;
+
+    if wire_byte_hi == wire_byte_lo {
+        // Single-byte slice: map each endpoint individually.
+        let map_bit = |bit_pos: usize| -> usize {
+            let wire_idx = field_width - 1 - bit_pos;
+            let wire_byte = wire_idx / 8;
+            let bit_in_byte = wire_idx % 8;
+            let storage_byte = field_width / 8 - 1 - wire_byte;
+            storage_byte * 8 + bit_in_byte
+        };
+
+        let mapped_hi = map_bit(hi);
+        let mapped_lo = map_bit(lo);
+        (mapped_hi.min(mapped_lo), mapped_hi.max(mapped_lo) + 1)
+    } else {
+        // Multi-byte slice: the HLIR rejects non-byte-aligned cases
+        // during validation.
+        assert!(
+            (hi + 1).is_multiple_of(8) && lo.is_multiple_of(8),
+            "non-byte-aligned multi-byte slice [{hi}:{lo}] on \
+             {field_width}-bit field reached codegen",
+        );
+
+        // Reversed storage bytes form a contiguous block.
+        let storage_byte_start = field_width / 8 - 1 - wire_byte_lo;
+        let storage_byte_end = field_width / 8 - 1 - wire_byte_hi;
+        (storage_byte_start * 8, (storage_byte_end + 1) * 8)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    // Verify the reversed slice range mapping against the byte reversal
+    // in header.rs. For each case we check that the bitvec range lands
+    // on the correct bits in the reversed storage layout.
+
+    // Sub-byte slices within a single wire byte.
+
+    #[test]
+    fn slice_32bit_top_nibble() {
+        // P4 [31:28] on 32-bit: top nibble of wire byte 0.
+        // Storage: wire byte 0 -> storage byte 3.
+        // High nibble of storage byte 3 = bitvec [24..28].
+        assert_eq!(reversed_slice_range(31, 28, 32), (24, 28));
+    }
+
+    #[test]
+    fn slice_32bit_bottom_nibble() {
+        // P4 [3:0] on 32-bit: bottom nibble of wire byte 3.
+        // Storage: wire byte 3 -> storage byte 0.
+        // Low nibble (Msb0) of storage byte 0 = bitvec [4..8].
+        assert_eq!(reversed_slice_range(3, 0, 32), (4, 8));
+    }
+
+    #[test]
+    fn slice_16bit_top_nibble() {
+        // P4 [15:12] on 16-bit: top nibble of wire byte 0.
+        // Storage: wire byte 0 -> storage byte 1.
+        // High nibble of storage byte 1 = bitvec [8..12].
+        assert_eq!(reversed_slice_range(15, 12, 16), (8, 12));
+    }
+
+    // Full-byte slices (single byte).
+
+    #[test]
+    fn slice_128bit_top_byte() {
+        // P4 [127:120] on 128-bit: wire byte 0 -> storage byte 15.
+        // bitvec [120..128].
+        assert_eq!(reversed_slice_range(127, 120, 128), (120, 128));
+    }
+
+    #[test]
+    fn slice_16bit_low_byte() {
+        // P4 [7:0] on 16-bit: wire byte 1 -> storage byte 0.
+        // bitvec [0..8].
+        assert_eq!(reversed_slice_range(7, 0, 16), (0, 8));
+    }
+
+    #[test]
+    fn slice_32bit_middle_byte() {
+        // P4 [23:16] on 32-bit: wire byte 1 -> storage byte 2.
+        // bitvec [16..24].
+        assert_eq!(reversed_slice_range(23, 16, 32), (16, 24));
+    }
+
+    // Multi-byte byte-aligned slices.
+
+    #[test]
+    fn slice_128bit_top_two_bytes() {
+        // P4 [127:112] on 128-bit: wire bytes 0-1 -> storage bytes 14-15.
+        // bitvec [112..128].
+        assert_eq!(reversed_slice_range(127, 112, 128), (112, 128));
+    }
+
+    #[test]
+    fn slice_32bit_top_three_bytes() {
+        // P4 [31:8] on 32-bit: wire bytes 0-2 -> storage bytes 1-3.
+        // bitvec [8..32].
+        assert_eq!(reversed_slice_range(31, 8, 32), (8, 32));
+    }
+
+    #[test]
+    fn slice_32bit_bottom_two_bytes() {
+        // P4 [15:0] on 32-bit: wire bytes 2-3 -> storage bytes 0-1.
+        // bitvec [0..16].
+        assert_eq!(reversed_slice_range(15, 0, 32), (0, 16));
+    }
+}