Skip to main content

primitives/utils/codec/
containers.rs

1//! [`InPlaceCodec`] for fixed-size container types.
2
3use std::{mem::MaybeUninit, sync::Arc};
4
5use hybrid_array::{Array, ArraySize};
6
7use super::{packed_size, read_packed_le_bytes, write_packed_le_bytes, InPlaceCodec};
8use crate::errors::PrimitiveError;
9
10// SAFETY: `Arc<T>` encodes exactly as its pointee `T` (same `ENCODED_SIZE`, same bytes), which is
11// architecture-independent, initializes every byte, and round-trips unbiasedly; decoding allocates
12// a fresh `Arc`.
13unsafe impl<T: InPlaceCodec> InPlaceCodec for Arc<T> {
14    const ENCODED_SIZE: usize = T::ENCODED_SIZE;
15
16    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
17        (**self).write_le_bytes(out);
18    }
19
20    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
21        T::read_le_bytes(bytes).map(Arc::new)
22    }
23}
24
25// SAFETY: a 2-tuple encodes as its two elements' `InPlaceCodec` encodings, back to back. Each half
26// occupies its element's `ENCODED_SIZE`; `write_le_bytes` initializes both (hence every byte) and
27// the round-trip is unbiased since each element's is.
28unsafe impl<A: InPlaceCodec, B: InPlaceCodec> InPlaceCodec for (A, B) {
29    const ENCODED_SIZE: usize = A::ENCODED_SIZE + B::ENCODED_SIZE;
30
31    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
32        let (a, b) = out.split_at_mut(A::ENCODED_SIZE);
33        self.0.write_le_bytes(a);
34        self.1.write_le_bytes(b);
35    }
36
37    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
38        let (a, b) = bytes.split_at(A::ENCODED_SIZE);
39        Ok((A::read_le_bytes(a)?, B::read_le_bytes(b)?))
40    }
41}
42
43// Fast (de)serialization for a native fixed-size array of `InPlaceCodec` elements. The `PACK`-aware
44// layout (`ceil(N / PACK)` full packs, trailing partial group padded with dummy elements) is
45// shared with `HeapArray<T, M>`/`Array<T, N>`'s impls via `packed_size`/`write_packed_le_bytes`/
46// `read_packed_le_bytes`.
47unsafe impl<T: InPlaceCodec, const N: usize> InPlaceCodec for [T; N] {
48    const ENCODED_SIZE: usize = packed_size::<T>(N);
49
50    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
51        write_packed_le_bytes(self, out);
52    }
53
54    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
55        let mut data: [MaybeUninit<T>; N] = std::array::from_fn(|_| MaybeUninit::uninit());
56        read_packed_le_bytes(bytes, &mut data)?;
57        // SAFETY: every slot of `data` was written above (`read_packed_le_bytes` only returns `Ok`
58        // once it has filled every slot).
59        Ok(data.map(|slot| unsafe { slot.assume_init() }))
60    }
61}
62
63// Fast (de)serialization for `hybrid_array::Array`s of `InPlaceCodec` elements. Same `PACK`-aware
64// layout as `[T; N]` above, shared via
65// `packed_size`/`write_packed_le_bytes`/`read_packed_le_bytes`.
66unsafe impl<T: InPlaceCodec, N: ArraySize> InPlaceCodec for Array<T, N> {
67    const ENCODED_SIZE: usize = packed_size::<T>(N::USIZE);
68
69    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
70        write_packed_le_bytes(self, out);
71    }
72
73    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
74        let mut data = Array::<MaybeUninit<T>, N>::uninit();
75        read_packed_le_bytes(bytes, &mut data)?;
76        // SAFETY: every slot of `data` was written above (`read_packed_le_bytes` only returns `Ok`
77        // once it has filled every slot).
78        Ok(unsafe { data.assume_init() })
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use hybrid_array::sizes::{U16, U20};
85
86    use super::*;
87    use crate::{
88        algebra::field::{binary::Gf2, mersenne::Mersenne107},
89        random::{test_rng, Random},
90    };
91
92    /// Mirrors `test_heap_array_pack_roundtrip_and_layout` in `types::heap_array::array`: `Gf2`
93    /// exercises sub-byte packing (8 -> 1); `Mersenne107` exercises the vectorized-but-unpacked
94    /// `PACK` path (byte-identical for full packs; an unpadded per-element tail for a remainder,
95    /// since padding would only add dead bytes here — see `tail_is_unpacked`).
96    #[test]
97    fn array_pack_roundtrip_and_layout() {
98        let mut rng = test_rng();
99
100        assert_eq!(
101            <Array<Gf2, U20> as InPlaceCodec>::ENCODED_SIZE,
102            20usize.div_ceil(8)
103        );
104        let gf2 = Array::<Gf2, U20>::from_fn(|_| Gf2::random(&mut rng));
105        let gf2_bytes = gf2.to_inplace_bytes();
106        assert_eq!(gf2_bytes.len(), 3);
107        assert_eq!(
108            Array::<Gf2, U20>::from_inplace_bytes(&gf2_bytes).unwrap(),
109            gf2
110        );
111
112        let mers16 = Array::<Mersenne107, U16>::from_fn(|_| Mersenne107::random(&mut rng));
113        let mers16_bytes = mers16.to_inplace_bytes();
114        let reference: Vec<u8> = mers16.iter().flat_map(|e| e.to_inplace_bytes()).collect();
115        assert_eq!(
116            mers16_bytes, reference,
117            "packed Mersenne107 must be byte-identical to the per-element encoding"
118        );
119        assert_eq!(
120            Array::<Mersenne107, U16>::from_inplace_bytes(&mers16_bytes).unwrap(),
121            mers16
122        );
123
124        // No compression from packing (`PACK_BYTES == PACK * ENCODED_SIZE`), so an unpadded tail
125        // keeps the total at exactly `20 * ENCODED_SIZE`, unlike `Gf2` above.
126        assert_eq!(
127            <Array<Mersenne107, U20> as InPlaceCodec>::ENCODED_SIZE,
128            20 * 14
129        );
130        let mers20 = Array::<Mersenne107, U20>::from_fn(|_| Mersenne107::random(&mut rng));
131        let mers20_bytes = mers20.to_inplace_bytes();
132        assert_eq!(mers20_bytes.len(), 20 * 14);
133        assert_eq!(
134            Array::<Mersenne107, U20>::from_inplace_bytes(&mers20_bytes).unwrap(),
135            mers20
136        );
137    }
138
139    /// Same as `array_pack_roundtrip_and_layout`, for the native `[T; N]` impl.
140    #[test]
141    fn fixed_array_pack_roundtrip_and_layout() {
142        let mut rng = test_rng();
143
144        assert_eq!(
145            <[Gf2; 20] as InPlaceCodec>::ENCODED_SIZE,
146            20usize.div_ceil(8)
147        );
148        let gf2: [Gf2; 20] = std::array::from_fn(|_| Gf2::random(&mut rng));
149        let gf2_bytes = gf2.to_inplace_bytes();
150        assert_eq!(gf2_bytes.len(), 3);
151        assert_eq!(<[Gf2; 20]>::from_inplace_bytes(&gf2_bytes).unwrap(), gf2);
152
153        let mers16: [Mersenne107; 16] = std::array::from_fn(|_| Mersenne107::random(&mut rng));
154        let mers16_bytes = mers16.to_inplace_bytes();
155        let reference: Vec<u8> = mers16.iter().flat_map(|e| e.to_inplace_bytes()).collect();
156        assert_eq!(
157            mers16_bytes, reference,
158            "packed Mersenne107 must be byte-identical to the per-element encoding"
159        );
160        assert_eq!(
161            <[Mersenne107; 16]>::from_inplace_bytes(&mers16_bytes).unwrap(),
162            mers16
163        );
164
165        assert_eq!(<[Mersenne107; 20] as InPlaceCodec>::ENCODED_SIZE, 20 * 14);
166        let mers20: [Mersenne107; 20] = std::array::from_fn(|_| Mersenne107::random(&mut rng));
167        let mers20_bytes = mers20.to_inplace_bytes();
168        assert_eq!(mers20_bytes.len(), 20 * 14);
169        assert_eq!(
170            <[Mersenne107; 20]>::from_inplace_bytes(&mers20_bytes).unwrap(),
171            mers20
172        );
173    }
174}