p4rs/
bitmath.rs

1// Copyright 2026 Oxide Computer Company
2
3use bitvec::prelude::*;
4use std::cmp::Ordering;
5
6pub fn add_be(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
7    let len = usize::max(a.len(), b.len());
8
9    // P4 spec says width limits are architecture defined, i here by define
10    // softnpu to have an architectural bit-type width limit of 128.
11    let x: u128 = a.load_be();
12    let y: u128 = b.load_be();
13    let z = x.wrapping_add(y);
14    let mut c = BitVec::new();
15    c.resize(len, false);
16    c.store_be(z);
17    c
18}
19
20pub fn add_le(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
21    let len = usize::max(a.len(), b.len());
22
23    // P4 spec says width limits are architecture defined, i here by define
24    // softnpu to have an architectural bit-type width limit of 128.
25    let x: u128 = a.load_le();
26    let y: u128 = b.load_le();
27    let z = x.wrapping_add(y);
28    let mut c = BitVec::new();
29    c.resize(len, false);
30    c.store_le(z);
31    c
32}
33
34pub fn sub_be(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
35    let len = usize::max(a.len(), b.len());
36    let x: u128 = a.load_be();
37    let y: u128 = b.load_be();
38    let z = x.wrapping_sub(y);
39    let mut c = BitVec::new();
40    c.resize(len, false);
41    c.store_be(z);
42    c
43}
44
45pub fn sub_le(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
46    let len = usize::max(a.len(), b.len());
47    let x: u128 = a.load_le();
48    let y: u128 = b.load_le();
49    let z = x.wrapping_sub(y);
50    let mut c = BitVec::new();
51    c.resize(len, false);
52    c.store_le(z);
53    c
54}
55
56// leaving here in case we have a need for a true arbitrary-width adder.
57#[allow(dead_code)]
58pub fn add_generic(
59    a: BitVec<u8, Msb0>,
60    b: BitVec<u8, Msb0>,
61) -> BitVec<u8, Msb0> {
62    if a.len() != b.len() {
63        panic!("bitvec add size mismatch");
64    }
65    let mut c = BitVec::new();
66    c.resize(a.len(), false);
67
68    for i in (1..a.len()).rev() {
69        let y = c[i];
70        let x = a[i] ^ b[i];
71        if !(a[i] | b[i]) {
72            continue;
73        }
74        c.set(i, x ^ y);
75        let mut z = (a[i] && b[i]) | y;
76        for j in (1..i).rev() {
77            if !z {
78                break;
79            }
80            z = c[j];
81            c.set(j, true);
82        }
83    }
84
85    c
86}
87
88pub fn mod_be(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
89    let len = usize::max(a.len(), b.len());
90
91    // P4 spec says width limits are architecture defined, i here by define
92    // softnpu to have an architectural bit-type width limit of 128.
93    let x: u128 = a.load_be();
94    let y: u128 = b.load_be();
95    let z = x % y;
96    let mut c = BitVec::new();
97    c.resize(len, false);
98    c.store_be(z);
99    c
100}
101
102/// Left shift `a` by `b` positions, little-endian byte order.
103/// Result width matches `a`. Shifts by the operand width or more produce zero.
104pub fn shl_le(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
105    if a.is_empty() {
106        return a;
107    }
108    let len = a.len();
109    let x: u128 = a.load_le();
110    let y: u128 = if b.is_empty() { 0 } else { b.load_le() };
111    let z = if y >= len as u128 { 0 } else { x << (y as u32) };
112    let mut c = BitVec::new();
113    c.resize(len, false);
114    c.store_le(z);
115    c
116}
117
118/// Right shift `a` by `b` positions, little-endian byte order.
119/// Result width matches `a`. Shifts by the operand width or more produce zero.
120pub fn shr_le(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
121    if a.is_empty() {
122        return a;
123    }
124    let len = a.len();
125    let x: u128 = a.load_le();
126    let y: u128 = if b.is_empty() { 0 } else { b.load_le() };
127    let z = if y >= len as u128 { 0 } else { x >> (y as u32) };
128    let mut c = BitVec::new();
129    c.resize(len, false);
130    c.store_le(z);
131    c
132}
133
134pub fn mod_le(a: BitVec<u8, Msb0>, b: BitVec<u8, Msb0>) -> BitVec<u8, Msb0> {
135    let len = usize::max(a.len(), b.len());
136
137    // P4 spec says width limits are architecture defined, i here by define
138    // softnpu to have an architectural bit-type width limit of 128.
139    let x: u128 = a.load_le();
140    let y: u128 = b.load_le();
141    let z = x % y;
142    let mut c = BitVec::new();
143    c.resize(len, false);
144    c.store_le(z);
145    c
146}
147
148/// An unsigned ordering view of a little-endian `bit<W>` value.
149///
150/// Implements the unsigned comparisons on `bit<W>` from P4-16 section [8.6].
151/// [`BitVec`]'s own ordering compares bits lexicographically, which is wrong
152/// for byte-reversed multi-byte fields.
153///
154/// Note: an empty bit-string compares as zero.
155///
156/// [8.6]: https://p4.org/wp-content/uploads/sites/53/2024/10/P4-16-spec-v1.2.5.html#sec-bit-ops
157pub struct Unsigned<'a>(pub &'a BitSlice<u8, Msb0>);
158
159impl Unsigned<'_> {
160    fn value(&self) -> u128 {
161        if self.0.is_empty() {
162            0
163        } else {
164            self.0.load_le()
165        }
166    }
167}
168
169impl Ord for Unsigned<'_> {
170    fn cmp(&self, other: &Self) -> Ordering {
171        self.value().cmp(&other.value())
172    }
173}
174
175impl PartialOrd for Unsigned<'_> {
176    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
177        Some(self.cmp(other))
178    }
179}
180
181impl PartialEq for Unsigned<'_> {
182    fn eq(&self, other: &Self) -> bool {
183        self.cmp(other) == Ordering::Equal
184    }
185}
186
187impl Eq for Unsigned<'_> {}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn unsigned_orders_by_value() {
195        let mut a = bitvec![u8, Msb0; 0; 16];
196        a.store_le(0x0100u16);
197        let mut b = bitvec![u8, Msb0; 0; 16];
198        b.store_le(0x00ffu16);
199        let mut one = bitvec![u8, Msb0; 0; 16];
200        one.store_le(1u16);
201        let zero = bitvec![u8, Msb0; 0; 16];
202        let empty = BitVec::<u8, Msb0>::new();
203
204        assert!(Unsigned(&a) > Unsigned(&b));
205        assert!(Unsigned(&empty) < Unsigned(&one));
206        assert!(Unsigned(&empty) == Unsigned(&zero));
207    }
208
209    #[test]
210    fn bitmath_add_be() {
211        let mut a = bitvec![mut u8, Msb0; 0; 16];
212        a.store_be(47);
213        let mut b = bitvec![mut u8, Msb0; 0; 16];
214        b.store_be(74);
215
216        println!("{:?}", a);
217        println!("{:?}", b);
218        let c = add_be(a, b);
219        println!("{:?}", c);
220
221        let cc: u128 = c.load_be();
222        assert_eq!(cc, 47u128 + 74u128);
223    }
224
225    #[test]
226    fn bitmath_add_le() {
227        let mut a = bitvec![mut u8, Msb0; 0; 16];
228        a.store_le(47);
229        let mut b = bitvec![mut u8, Msb0; 0; 16];
230        b.store_le(74);
231
232        println!("{:?}", a);
233        println!("{:?}", b);
234        let c = add_le(a, b);
235        println!("{:?}", c);
236
237        let cc: u128 = c.load_le();
238        assert_eq!(cc, 47u128 + 74u128);
239    }
240
241    #[test]
242    fn bitmath_sub_be() {
243        use super::*;
244        let mut a = bitvec![mut u8, Msb0; 0; 16];
245        a.store_be(74);
246        let mut b = bitvec![mut u8, Msb0; 0; 16];
247        b.store_be(47);
248
249        println!("{:?}", a);
250        println!("{:?}", b);
251        let c = sub_be(a, b);
252        println!("{:?}", c);
253
254        let cc: u128 = c.load_be();
255        assert_eq!(cc, 74u128 - 47u128);
256    }
257
258    #[test]
259    fn bitmath_sub_le() {
260        use super::*;
261        let mut a = bitvec![mut u8, Msb0; 0; 16];
262        a.store_le(74);
263        let mut b = bitvec![mut u8, Msb0; 0; 16];
264        b.store_le(47);
265
266        println!("{:?}", a);
267        println!("{:?}", b);
268        let c = sub_le(a, b);
269        println!("{:?}", c);
270
271        let cc: u128 = c.load_le();
272        assert_eq!(cc, 74u128 - 47u128);
273    }
274
275    #[test]
276    fn bitmath_add_mixed_size() {
277        use super::*;
278        let mut a = bitvec![mut u8, Msb0; 0; 8];
279        a.store_be(0xAB);
280        let mut b = bitvec![mut u8, Msb0; 0; 16];
281        b.store_be(0xCDE);
282
283        println!("{:?}", a);
284        println!("{:?}", b);
285        let c = add_be(a, b);
286        println!("{:?}", c);
287
288        let cc: u128 = c.load_be();
289        assert_eq!(cc, 0xABu128 + 0xCDEu128);
290    }
291
292    #[test]
293    fn bitmath_add_cascade() {
294        use super::*;
295        let mut a = bitvec![mut u8, Msb0; 0; 16];
296        a.store_be(47);
297        let mut b = bitvec![mut u8, Msb0; 0; 16];
298        b.store_be(74);
299        let mut c = bitvec![mut u8, Msb0; 0; 16];
300        c.store_be(123);
301        let mut d = bitvec![mut u8, Msb0; 0; 16];
302        d.store_be(9876);
303
304        let e = add_be(a, add_be(b, add_be(c, d)));
305
306        let ee: u128 = e.load_be();
307        assert_eq!(ee, 47u128 + 74u128 + 123u128 + 9876u128);
308    }
309
310    #[test]
311    fn bitmath_add_nest() {
312        use super::*;
313        let mut orig_l3_len = bitvec![mut u8, Msb0; 0; 16usize];
314        orig_l3_len.store_le(0xe9u128);
315        let x = add_le(
316            {
317                let mut x = bitvec![mut u8, Msb0; 0; 16usize];
318                x.store_le(14u128);
319                x
320            },
321            add_le(
322                orig_l3_len.clone(),
323                add_le(
324                    {
325                        let mut x = bitvec![mut u8, Msb0; 0; 16usize];
326                        x.store_le(8u128);
327                        x
328                    },
329                    {
330                        let mut x = bitvec![mut u8, Msb0; 0; 16usize];
331                        x.store_le(8u128);
332                        x
333                    },
334                ),
335            ),
336        );
337
338        let y: u128 = x.load_le();
339        assert_eq!(y, 0xe9 + 14 + 8 + 8);
340    }
341
342    #[test]
343    fn bitmath_mod() {
344        use super::*;
345        let mut a = bitvec![mut u8, Msb0; 0; 16];
346        a.store_be(47);
347        let mut b = bitvec![mut u8, Msb0; 0; 16];
348        b.store_be(7);
349
350        println!("{:?}", a);
351        println!("{:?}", b);
352        let c = mod_be(a, b);
353        println!("{:?}", c);
354
355        let cc: u128 = c.load_be();
356        assert_eq!(cc, 47u128 % 7u128);
357    }
358
359    #[test]
360    fn bitmath_shl_le() {
361        let mut a = bitvec![mut u8, Msb0; 0; 16];
362        a.store_le(1u128);
363        let mut b = bitvec![mut u8, Msb0; 0; 16];
364        b.store_le(4u128);
365
366        println!("{:?}", a);
367        println!("{:?}", b);
368        let c = shl_le(a, b);
369        println!("{:?}", c);
370
371        let cc: u128 = c.load_le();
372        assert_eq!(cc, 1u128 << 4);
373    }
374
375    #[test]
376    fn bitmath_shr_le() {
377        let mut a = bitvec![mut u8, Msb0; 0; 16];
378        a.store_le(0x8000u128);
379        let mut b = bitvec![mut u8, Msb0; 0; 16];
380        b.store_le(4u128);
381
382        println!("{:?}", a);
383        println!("{:?}", b);
384        let c = shr_le(a, b);
385        println!("{:?}", c);
386
387        let cc: u128 = c.load_le();
388        assert_eq!(cc, 0x8000u128 >> 4);
389    }
390
391    #[test]
392    fn bitmath_shifts_at_width_le() {
393        let mut a = bitvec![mut u8, Msb0; 0; 16];
394        a.store_le(1u128);
395        let mut b = bitvec![mut u8, Msb0; 0; 16];
396        b.store_le(16u128);
397
398        assert_eq!(shl_le(a.clone(), b.clone()).load_le::<u16>(), 0);
399        assert_eq!(shr_le(a, b).load_le::<u16>(), 0);
400    }
401
402    #[test]
403    fn bitmath_shifts_by_zero_width_count() {
404        let mut value = bitvec![u8, Msb0; 0; 16];
405        value.store_le(0x8123u16);
406
407        assert_eq!(shl_le(value.clone(), BitVec::new()), value);
408        assert_eq!(shr_le(value.clone(), BitVec::new()), value);
409    }
410
411    #[test]
412    fn bitmath_shifts_zero_width_value() {
413        let mut count = bitvec![u8, Msb0; 0; 8];
414        count.store_le(1u8);
415
416        for count in [BitVec::new(), count] {
417            assert!(shl_le(BitVec::new(), count.clone()).is_empty());
418            assert!(shr_le(BitVec::new(), count).is_empty());
419        }
420    }
421
422    #[test]
423    fn bitmath_shl_shr_roundtrip_le() {
424        let mut a = bitvec![mut u8, Msb0; 0; 32];
425        a.store_le(42u128);
426        let mut b = bitvec![mut u8, Msb0; 0; 32];
427        b.store_le(7u128);
428
429        println!("{:?}", a);
430        println!("{:?}", b);
431        let shifted = shl_le(a, b.clone());
432        println!("{:?}", shifted);
433        let back = shr_le(shifted, b);
434        println!("{:?}", back);
435
436        let result: u128 = back.load_le();
437        assert_eq!(result, 42u128);
438    }
439}