primitives/types/heap_array/
array.rs

1use std::{
2    fmt::{Debug, Display},
3    marker::PhantomData,
4    mem::MaybeUninit,
5    ops::{Add, Mul, Sub},
6    sync::Arc,
7    vec,
8};
9
10use bytemuck::{box_bytes_of, from_box_bytes, BoxBytes, Pod};
11use derive_more::derive::{AsMut, AsRef, Deref, DerefMut, IntoIterator};
12use hybrid_array::{Array, ArraySize};
13use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
14use serde::{Deserialize, Serialize};
15use typenum::{Diff, Prod, Sum, Unsigned, U1, U2, U3, U5};
16
17use crate::{
18    errors::PrimitiveError,
19    types::Positive,
20    utils::codec::{packed_size, read_packed_le_bytes, write_packed_le_bytes, InPlaceCodec},
21};
22
23/// An array on the heap that encodes its length in the type system.
24#[derive(Deref, DerefMut, Clone, IntoIterator, AsRef, AsMut, Eq)]
25#[into_iterator(owned, ref, ref_mut)]
26pub struct HeapArray<T: Sized, M: Positive> {
27    #[deref]
28    #[deref_mut]
29    #[as_ref(forward)]
30    #[as_mut(forward)]
31    pub(super) data: Box<[T]>,
32    #[into_iterator(ignore)]
33    // `fn() -> M` is used instead of `M` so `HeapArray<T, M>` doesn't need `M` to implement `Send
34    // + Sync` to be `Send + Sync` itself. This would be the case if `M` was used directly.
35    pub(super) _len: PhantomData<fn() -> M>,
36}
37
38impl<T: Sized, M: Positive> HeapArray<T, M> {
39    pub(super) fn new(data: Box<[T]>) -> Self {
40        Self {
41            data,
42            _len: PhantomData,
43        }
44    }
45}
46
47impl<T: Sized, M: Positive> HeapArray<T, M> {
48    /// Zero-copy cast from a `#[repr(transparent)]` wrapper to its inner type.
49    pub fn peel_transparent<U>(self) -> HeapArray<U, M>
50    where
51        T: bytemuck::TransparentWrapper<U>,
52    {
53        use bytemuck::allocation::TransparentWrapperAlloc;
54        HeapArray::new(T::peel_vec(self.data.into_vec()).into_boxed_slice())
55    }
56
57    /// Zero-copy cast from a type to its `#[repr(transparent)]` wrapper.
58    pub fn wrap_transparent<W>(self) -> HeapArray<W, M>
59    where
60        W: bytemuck::TransparentWrapper<T>,
61    {
62        use bytemuck::allocation::TransparentWrapperAlloc;
63        HeapArray::new(W::wrap_vec(self.data.into_vec()).into_boxed_slice())
64    }
65}
66
67// bytemuck::BoxBytes transformation for copy-less casting
68impl<T: Pod, M: Positive> HeapArray<T, M> {
69    pub fn into_box_bytes(self) -> BoxBytes {
70        box_bytes_of(self.data)
71    }
72
73    pub fn from_box_bytes(buf: BoxBytes) -> Self {
74        Self {
75            data: from_box_bytes(buf),
76            _len: PhantomData,
77        }
78    }
79}
80
81// Fast, allocation-minimal (de)serialization for `HeapArray`s of `InPlaceCodec` types, used by
82// the `Serialize`/`Deserialize` impls below instead of per-element `serde` dispatch.
83// `to_inplace_bytes`/`from_inplace_bytes` come for free from the trait's default methods.
84unsafe impl<T: InPlaceCodec, M: Positive> InPlaceCodec for HeapArray<T, M> {
85    const ENCODED_SIZE: usize = packed_size::<T>(M::USIZE);
86
87    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
88        write_packed_le_bytes(&self.data, out);
89    }
90
91    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
92        let mut data = Box::<[T]>::new_uninit_slice(M::USIZE);
93        read_packed_le_bytes(bytes, &mut data)?;
94        // SAFETY: every slot of `data` was written above (`read_packed_le_bytes` only returns `Ok`
95        // once it has filled every slot).
96        Ok(Self::new(unsafe { data.assume_init() }))
97    }
98}
99
100impl<T: Sized, M: Positive> HeapArray<T, M> {
101    pub fn map<F, U>(self, f: F) -> HeapArray<U, M>
102    where
103        F: FnMut(T) -> U,
104    {
105        self.into_iter().map(f).collect()
106    }
107
108    /// Destructure into a fixed-size array `[T; K]`.
109    ///
110    /// `K` is inferred from the destructuring pattern at the call site
111    /// (e.g. `let [a, b, c] = ha.into_array();` fixes `K = 3`). A mismatch
112    /// between `K` and `M::USIZE` is a compile-time error.
113    pub fn into_array<const K: usize>(self) -> [T; K] {
114        const {
115            assert!(
116                M::USIZE == K,
117                "HeapArray length does not match destructured array length",
118            );
119        }
120        // SAFETY: the const assert above proves `self.data.len() == M::USIZE == K`,
121        // and `Box<[T]>` with length `K` has the same layout as `Box<[T; K]>`.
122        let raw: *mut [T] = Box::into_raw(self.data);
123        unsafe { *Box::from_raw(raw.cast::<[T; K]>()) }
124    }
125}
126
127impl<T: Sized + Default, M: Positive> HeapArray<T, M> {
128    pub fn from_single_value(val: T) -> Self {
129        Self {
130            data: vec![val].into_boxed_slice(),
131            _len: PhantomData,
132        }
133    }
134}
135
136impl<T: Sized + Default, M: Positive> Default for HeapArray<T, M> {
137    fn default() -> Self {
138        Self {
139            data: (0..M::USIZE).map(|_| T::default()).collect(),
140            _len: PhantomData,
141        }
142    }
143}
144
145impl<T: Sized + Debug, M: Positive> Debug for HeapArray<T, M> {
146    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        f.debug_struct(format!("HeapArray<{}>", M::USIZE).as_str())
148            .field("data", &self.data)
149            .finish()
150    }
151}
152
153impl<T: Sized + InPlaceCodec, M: Positive> Serialize for HeapArray<T, M> {
154    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
155        serializer.serialize_bytes(self.to_inplace_bytes().as_slice())
156    }
157}
158
159impl<'de, T: Sized + InPlaceCodec, M: Positive> Deserialize<'de> for HeapArray<T, M> {
160    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
161        // Zero-copy borrow from the input buffer when the format supports it.
162        let bytes: &[u8] = Deserialize::deserialize(deserializer)?;
163        HeapArray::<T, M>::from_inplace_bytes(bytes).map_err(serde::de::Error::custom)
164    }
165}
166
167impl<T: Sized, M: Positive> FromIterator<T> for HeapArray<T, M> {
168    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
169        let data = iter.into_iter().collect::<Box<_>>();
170        assert_eq!(data.len(), M::USIZE,);
171        Self {
172            data,
173            _len: PhantomData,
174        }
175    }
176}
177
178impl<T: Sized + Send, M: Positive> IntoParallelIterator for HeapArray<T, M> {
179    type Item = T;
180    type Iter = rayon::vec::IntoIter<T>;
181
182    fn into_par_iter(self) -> Self::Iter {
183        self.data.into_par_iter()
184    }
185}
186
187impl<T: Sized + Send, M: Positive> FromParallelIterator<T> for HeapArray<T, M> {
188    fn from_par_iter<I: IntoParallelIterator<Item = T>>(par_iter: I) -> Self {
189        let data: Box<[T]> = par_iter.into_par_iter().collect::<Vec<T>>().into();
190        assert_eq!(data.len(), M::USIZE);
191        Self {
192            data,
193            _len: PhantomData,
194        }
195    }
196}
197
198// -----------------------
199// |   Split and Merge   |
200// -----------------------
201
202impl<T: Sized + Copy, M: Positive> HeapArray<T, M> {
203    pub fn split<M1, M2>(&self) -> (HeapArray<T, M1>, HeapArray<T, M2>)
204    where
205        M1: Positive,
206        M2: Positive + Add<M1, Output = M>,
207    {
208        let (m1, m2) = self.split_at(M1::USIZE);
209        (
210            HeapArray::<T, M1> {
211                data: m1.into(),
212                _len: PhantomData,
213            },
214            HeapArray::<T, M2> {
215                data: m2.into(),
216                _len: PhantomData,
217            },
218        )
219    }
220
221    pub fn split_last_pos<M1>(self) -> (HeapArray<T, M1>, T)
222    where
223        M1: Positive + Add<typenum::B1, Output = M>,
224    {
225        let Self { data, .. } = self;
226        let mut data = data.into_vec();
227        let last = data.pop().expect("HeapArray is empty");
228        (
229            HeapArray::<T, M1> {
230                data: data.into_boxed_slice(),
231                _len: PhantomData,
232            },
233            last,
234        )
235    }
236
237    pub fn split_halves<MDiv2>(&self) -> (HeapArray<T, MDiv2>, HeapArray<T, MDiv2>)
238    where
239        MDiv2: Positive + Mul<U2, Output = M>,
240    {
241        let (m1, m2) = self.split_at(MDiv2::USIZE);
242        (
243            HeapArray::<T, MDiv2> {
244                data: m1.into(),
245                _len: PhantomData,
246            },
247            HeapArray::<T, MDiv2> {
248                data: m2.into(),
249                _len: PhantomData,
250            },
251        )
252    }
253
254    pub fn merge_halves(this: Self, other: Self) -> HeapArray<T, Prod<M, U2>>
255    where
256        M: Mul<U2, Output: Positive>,
257    {
258        let mut vec = this.data.into_vec();
259        vec.extend(other.data.into_vec());
260        HeapArray::<T, Prod<M, U2>> {
261            data: vec.into_boxed_slice(),
262            _len: PhantomData,
263        }
264    }
265
266    pub fn split3<M1, M2, M3>(&self) -> (HeapArray<T, M1>, HeapArray<T, M2>, HeapArray<T, M3>)
267    where
268        M1: Positive,
269        M2: Positive + Add<M1>,
270        M3: Positive + Add<Sum<M2, M1>, Output = M>,
271    {
272        let (m1, m_rest) = self.split_at(M1::USIZE);
273        let (m2, m3) = m_rest.split_at(M2::USIZE);
274        (
275            HeapArray::<T, M1> {
276                data: m1.into(),
277                _len: PhantomData,
278            },
279            HeapArray::<T, M2> {
280                data: m2.into(),
281                _len: PhantomData,
282            },
283            HeapArray::<T, M3> {
284                data: m3.into(),
285                _len: PhantomData,
286            },
287        )
288    }
289
290    pub fn split_thirds<MDiv3>(
291        &self,
292    ) -> (
293        HeapArray<T, MDiv3>,
294        HeapArray<T, MDiv3>,
295        HeapArray<T, MDiv3>,
296    )
297    where
298        MDiv3: Positive + Mul<U3, Output = M>,
299    {
300        let (m1, m_rest) = self.split_at(MDiv3::USIZE);
301        let (m2, m3) = m_rest.split_at(MDiv3::USIZE);
302        (
303            HeapArray::<T, MDiv3> {
304                data: m1.into(),
305                _len: PhantomData,
306            },
307            HeapArray::<T, MDiv3> {
308                data: m2.into(),
309                _len: PhantomData,
310            },
311            HeapArray::<T, MDiv3> {
312                data: m3.into(),
313                _len: PhantomData,
314            },
315        )
316    }
317
318    pub fn merge_thirds(first: Self, second: Self, third: Self) -> HeapArray<T, Prod<M, U3>>
319    where
320        M: Mul<U3, Output: Positive>,
321    {
322        let mut vec = first.data.into_vec();
323        vec.extend(second.data.into_vec());
324        vec.extend(third.data.into_vec());
325        HeapArray::<T, Prod<M, U3>> {
326            data: vec.into_boxed_slice(),
327            _len: PhantomData,
328        }
329    }
330
331    pub fn merge_fifths(
332        first: Self,
333        second: Self,
334        third: Self,
335        fourth: Self,
336        fifth: Self,
337    ) -> HeapArray<T, Prod<M, U5>>
338    where
339        M: Mul<U5, Output: Positive>,
340    {
341        let mut vec = first.data.into_vec();
342        vec.reserve_exact(M::USIZE * 4);
343        vec.extend(second.data.into_vec());
344        vec.extend(third.data.into_vec());
345        vec.extend(fourth.data.into_vec());
346        vec.extend(fifth.data.into_vec());
347        HeapArray::<T, Prod<M, U5>> {
348            data: vec.into_boxed_slice(),
349            _len: PhantomData,
350        }
351    }
352}
353
354pub struct HeapArrayTuple<T1: Sized, T2: Sized, M: Positive>(
355    pub HeapArray<T1, M>,
356    pub HeapArray<T2, M>,
357);
358
359impl<T1: Sized, T2: Sized, M: Positive> FromIterator<(T1, T2)> for HeapArrayTuple<T1, T2, M> {
360    fn from_iter<I: IntoIterator<Item = (T1, T2)>>(iter: I) -> Self {
361        let (data1, data2): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
362
363        assert_eq!(data1.len(), M::USIZE);
364        assert_eq!(data2.len(), M::USIZE);
365        HeapArrayTuple(
366            HeapArray::<T1, M> {
367                data: data1.into_boxed_slice(),
368                _len: PhantomData,
369            },
370            HeapArray::<T2, M> {
371                data: data2.into_boxed_slice(),
372                _len: PhantomData,
373            },
374        )
375    }
376}
377
378impl<T: Sized, M: Positive> TryFrom<Vec<T>> for HeapArray<T, M> {
379    type Error = PrimitiveError;
380
381    fn try_from(data: Vec<T>) -> Result<Self, PrimitiveError> {
382        if data.len() == M::USIZE {
383            Ok(Self {
384                data: data.into_boxed_slice(),
385                _len: PhantomData,
386            })
387        } else {
388            Err(PrimitiveError::InvalidSize(M::USIZE, data.len()))
389        }
390    }
391}
392
393impl<T: Sized, M: Positive> TryFrom<Box<[T]>> for HeapArray<T, M> {
394    type Error = PrimitiveError;
395
396    fn try_from(data: Box<[T]>) -> Result<Self, PrimitiveError> {
397        if data.len() == M::USIZE {
398            Ok(Self {
399                data,
400                _len: PhantomData,
401            })
402        } else {
403            Err(PrimitiveError::InvalidSize(M::USIZE, data.len()))
404        }
405    }
406}
407
408impl<T: Sized + Clone, M: Positive> TryFrom<Arc<[T]>> for HeapArray<T, M> {
409    type Error = PrimitiveError;
410
411    fn try_from(data: Arc<[T]>) -> Result<Self, PrimitiveError> {
412        if data.len() != M::USIZE {
413            return Err(PrimitiveError::InvalidSize(M::USIZE, data.len()));
414        }
415        Ok(Self {
416            data: data.to_vec().into_boxed_slice(),
417            _len: PhantomData,
418        })
419    }
420}
421
422impl<T: Sized> From<T> for HeapArray<T, U1> {
423    fn from(element: T) -> Self {
424        Self {
425            data: Box::new([element]),
426            _len: PhantomData,
427        }
428    }
429}
430
431impl<T: Sized, M: Positive> From<HeapArray<T, M>> for Vec<T> {
432    fn from(array: HeapArray<T, M>) -> Self {
433        array.data.into_vec()
434    }
435}
436
437impl<T: Sized + Clone, M: Positive + ArraySize> From<Array<T, M>> for HeapArray<T, M> {
438    fn from(array: Array<T, M>) -> Self {
439        Self {
440            data: array.to_vec().into_boxed_slice(),
441            _len: PhantomData,
442        }
443    }
444}
445
446impl<T: Sized, M: Positive> HeapArray<T, M> {
447    /// Take the first `N` elements and discard the rest. `N` may equal `M`
448    /// (no slack) or be strictly less.
449    pub fn truncate<N>(self) -> HeapArray<T, N>
450    where
451        N: Positive,
452        M: Sub<N, Output: Unsigned>,
453    {
454        let mut vec = self.data.into_vec();
455        vec.truncate(N::USIZE);
456        HeapArray {
457            data: vec.into_boxed_slice(),
458            _len: PhantomData,
459        }
460    }
461
462    pub fn split_last<N: Positive>(self) -> (HeapArray<T, Diff<M, N>>, HeapArray<T, N>)
463    where
464        M: Sub<N, Output: Positive>,
465    {
466        let mut vec = self.data.into_vec();
467        let last_n = vec.split_off(M::USIZE - N::USIZE);
468
469        (
470            HeapArray {
471                data: vec.into_boxed_slice(),
472                _len: PhantomData,
473            },
474            HeapArray {
475                data: last_n.into_boxed_slice(),
476                _len: PhantomData,
477            },
478        )
479    }
480
481    pub fn from_fn(f: impl FnMut(usize) -> T) -> Self {
482        Self {
483            data: (0..M::USIZE).map(f).collect::<Box<_>>(),
484            _len: PhantomData,
485        }
486    }
487
488    pub fn from_constant(c: T) -> Self
489    where
490        T: Copy,
491    {
492        Self {
493            data: (0..M::USIZE).map(|_| c).collect::<Box<_>>(),
494            _len: PhantomData,
495        }
496    }
497}
498
499impl<T: Sized, M: Positive> Display for HeapArray<T, M>
500where
501    T: Display,
502{
503    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
504        write!(f, "[")?;
505        for (i, item) in self.data.iter().enumerate() {
506            if i != 0 {
507                write!(f, ", ")?;
508            }
509            write!(f, "{item}")?;
510        }
511        write!(f, "]")
512    }
513}
514
515#[cfg(test)]
516pub mod tests {
517    use hybrid_array::sizes::{U2, U3, U6};
518    use typenum::{U1, U4};
519
520    use super::*;
521    use crate::utils::codec::bincode_io;
522
523    #[test]
524    fn test_heap_array() {
525        let array = HeapArray::<_, U3>::from_fn(|i| i);
526        assert_eq!(array.into_iter().collect::<Vec<_>>(), vec![0, 1, 2]);
527    }
528
529    #[test]
530    fn test_default() {
531        let array = HeapArray::<usize, U3>::default();
532        assert_eq!(array.len(), 3);
533    }
534
535    #[test]
536    fn test_heap_array_split_last() {
537        let array = HeapArray::<_, U6>::from_fn(|i| i);
538        let (first, last) = array.split_last::<U2>();
539        assert_eq!(first.into_iter().collect::<Vec<_>>(), vec![0, 1, 2, 3]);
540        assert_eq!(last.into_iter().collect::<Vec<_>>(), vec![4, 5]);
541    }
542
543    #[test]
544    fn test_heap_array_from_array() {
545        let array = Array::<_, U3>::from_fn(|i| i);
546        let heap_array = HeapArray::<_, U3>::from(array);
547        assert_eq!(heap_array.into_iter().collect::<Vec<_>>(), vec![0, 1, 2]);
548    }
549
550    #[test]
551    fn test_heap_array_from_vec() {
552        let vec = vec![0, 1, 2];
553        let heap_array = HeapArray::<_, U3>::try_from(vec).unwrap();
554        assert_eq!(heap_array.into_iter().collect::<Vec<_>>(), vec![0, 1, 2]);
555
556        let vec = vec![0, 1];
557        let heap_array = HeapArray::<_, U3>::try_from(vec);
558        assert!(heap_array.is_err());
559    }
560
561    #[test]
562    fn test_heap_array_from_iter() {
563        let heap_array = HeapArray::<_, U3>::from_fn(|i| i);
564        assert_eq!(heap_array.into_iter().collect::<Vec<_>>(), vec![0, 1, 2]);
565    }
566
567    #[test]
568    #[should_panic]
569    fn test_heap_array_from_iter_wrong_size() {
570        HeapArray::<_, U2>::from_iter(0..3);
571    }
572
573    #[test]
574    fn test_heap_array_deserialize() {
575        let array = HeapArray::<usize, U6>::from_fn(|i| i);
576        let serialized = bincode_io::serialize(&array).unwrap();
577        bincode_io::deserialize::<HeapArray<usize, U6>>(&serialized).unwrap();
578
579        // With the new tuple-based serialization (which doesn't include length),
580        // we can't detect length mismatches unless we use strict deserialization
581        // that rejects trailing bytes.
582        use bincode::Options;
583        let config = bincode::DefaultOptions::new()
584            .with_fixint_encoding()
585            .reject_trailing_bytes();
586
587        let wrong_deserialize = config.deserialize::<HeapArray<usize, U3>>(&serialized);
588        assert!(wrong_deserialize.is_err());
589    }
590
591    #[test]
592    fn test_heap_array_split() {
593        let array = HeapArray::<_, U6>::from_fn(|i| i);
594        let (first, second) = array.split::<U4, U2>();
595        assert_eq!(first.into_iter().collect::<Vec<_>>(), vec![0, 1, 2, 3]);
596        assert_eq!(second.into_iter().collect::<Vec<_>>(), vec![4, 5]);
597
598        let (first, second) = array.split_halves::<U3>();
599        assert_eq!(first.into_iter().collect::<Vec<_>>(), vec![0, 1, 2]);
600        assert_eq!(second.into_iter().collect::<Vec<_>>(), vec![3, 4, 5]);
601
602        // let (first, second) = array.split::<U4, U1>(); --> doesn't compile
603
604        // let (first, second) = array.split_halves::<U2>(); --> doesn't compile
605    }
606
607    #[test]
608    fn test_heap_array_split3() {
609        let array = HeapArray::<_, U6>::from_fn(|i| i);
610        let (first, second, third) = array.split3::<U3, U2, U1>();
611        assert_eq!(first.into_iter().collect::<Vec<_>>(), vec![0, 1, 2]);
612        assert_eq!(second.into_iter().collect::<Vec<_>>(), vec![3, 4]);
613        assert_eq!(third.into_iter().collect::<Vec<_>>(), vec![5]);
614
615        let (first, second, third) = array.split_thirds::<U2>();
616        assert_eq!(first.into_iter().collect::<Vec<_>>(), vec![0, 1]);
617        assert_eq!(second.into_iter().collect::<Vec<_>>(), vec![2, 3]);
618        assert_eq!(third.into_iter().collect::<Vec<_>>(), vec![4, 5]);
619
620        // let (a1, a2, a3) = array.split3::<U3, U2, U2>(); // doesn't compile
621
622        // let (a1, a2, a3) = array.split_thirds::<U2>(); // doesn't compile
623    }
624
625    #[test]
626    fn test_heap_array_wrap_transparent() {
627        #[repr(transparent)]
628        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
629        struct Wrap(u32);
630        // SAFETY: `Wrap` is `#[repr(transparent)]` over `u32`.
631        unsafe impl bytemuck::TransparentWrapper<u32> for Wrap {}
632
633        let array = HeapArray::<u32, U3>::from_fn(|i| (i as u32) * 10);
634        let wrapped: HeapArray<Wrap, U3> = array.wrap_transparent();
635        assert_eq!(
636            wrapped.into_iter().collect::<Vec<_>>(),
637            vec![Wrap(0), Wrap(10), Wrap(20)],
638        );
639    }
640
641    #[test]
642    fn test_heap_array_peel_transparent() {
643        #[repr(transparent)]
644        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
645        struct Wrap(u32);
646        // SAFETY: `Wrap` is `#[repr(transparent)]` over `u32`.
647        unsafe impl bytemuck::TransparentWrapper<u32> for Wrap {}
648
649        let array = HeapArray::<Wrap, U3>::from_fn(|i| Wrap((i as u32) * 10));
650        let peeled: HeapArray<u32, U3> = array.peel_transparent();
651        assert_eq!(peeled.into_iter().collect::<Vec<_>>(), vec![0, 10, 20]);
652    }
653
654    #[test]
655    fn test_heap_array_wrap_peel_roundtrip() {
656        #[repr(transparent)]
657        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
658        struct Wrap(u32);
659        // SAFETY: `Wrap` is `#[repr(transparent)]` over `u32`.
660        unsafe impl bytemuck::TransparentWrapper<u32> for Wrap {}
661
662        let original = HeapArray::<u32, U4>::from_fn(|i| i as u32 + 1);
663        let expected = original.clone().into_iter().collect::<Vec<_>>();
664        let roundtripped: HeapArray<u32, U4> =
665            original.wrap_transparent::<Wrap>().peel_transparent();
666        assert_eq!(roundtripped.into_iter().collect::<Vec<_>>(), expected);
667    }
668
669    #[test]
670    fn test_heap_array_inplace_roundtrip() {
671        use crate::algebra::field::mersenne::Mersenne107;
672
673        let array = HeapArray::<Mersenne107, U6>::from_fn(|i| Mersenne107::from(i as u64));
674        let inplace_ser = array.to_inplace_bytes();
675
676        // `bincode` round-trips via the same `InPlaceCodec` fast path (through `HeapArray`'s serde
677        // impl), though its bytes carry a length prefix so they differ from the raw buffer.
678        let bin_ser = bincode_io::serialize(&array).unwrap();
679        assert_eq!(
680            bincode_io::deserialize::<HeapArray<Mersenne107, U6>>(&bin_ser).unwrap(),
681            array
682        );
683
684        let deserialized = HeapArray::<Mersenne107, U6>::from_inplace_bytes(&inplace_ser).unwrap();
685        assert_eq!(deserialized, array);
686
687        // Wrong length is rejected.
688        assert!(HeapArray::<Mersenne107, U6>::from_inplace_bytes(&inplace_ser[..1]).is_err());
689
690        // Non-canonical encoding of the first element is rejected.
691        let mut bad_bytes = inplace_ser.clone();
692        bad_bytes[..14].copy_from_slice(&[0xFF; 14]);
693        assert!(HeapArray::<Mersenne107, U6>::from_inplace_bytes(&bad_bytes).is_err());
694
695        // Non-canonical encoding of the last *real* element is rejected too (exercises the batched
696        // canonicalization pass, not just a short-circuit on the first). `U6 < PACK` (8), so the
697        // whole array is one padded pack — target the last real element's 14-byte window, not the
698        // trailing padding.
699        let mut bad_bytes = inplace_ser.clone();
700        let last_real = 5 * 14;
701        bad_bytes[last_real..last_real + 14].copy_from_slice(&[0xFF; 14]);
702        assert!(HeapArray::<Mersenne107, U6>::from_inplace_bytes(&bad_bytes).is_err());
703    }
704
705    /// For every type with an `InPlaceCodec` impl, assert a `HeapArray` of it encodes to exactly
706    /// `ENCODED_SIZE` bytes and round-trips through the in-place encoding.
707    #[test]
708    fn test_field_extension_inplace_roundtrip() {
709        use crate::{
710            algebra::{
711                elliptic_curve::{BaseField, Curve25519Ristretto, Point, ScalarField},
712                field::{
713                    binary::{Gf2, Gf2_128},
714                    mersenne::Mersenne107,
715                    FieldElement,
716                    SubfieldElement,
717                },
718            },
719            random::{test_rng, Random},
720            utils::codec::InPlaceCodec,
721        };
722
723        fn check<T>()
724        where
725            T: InPlaceCodec + Random + PartialEq + std::fmt::Debug,
726        {
727            let name = std::any::type_name::<T>();
728            let mut rng = test_rng();
729            let array = HeapArray::<T, U6>::random(&mut rng);
730            let inplace = array.to_inplace_bytes();
731            assert_eq!(
732                inplace.len(),
733                <HeapArray<T, U6> as InPlaceCodec>::ENCODED_SIZE,
734                "encoded size mismatch for {name}"
735            );
736
737            let back = HeapArray::<T, U6>::from_inplace_bytes(&inplace).unwrap();
738            assert_eq!(back, array, "round-trip failed for {name}");
739        }
740
741        check::<Gf2>();
742        check::<Gf2_128>();
743        check::<Mersenne107>();
744        check::<ScalarField<Curve25519Ristretto>>();
745        check::<BaseField<Curve25519Ristretto>>();
746        check::<FieldElement<Gf2_128>>();
747        check::<SubfieldElement<Gf2_128>>();
748        check::<Point<Curve25519Ristretto>>();
749    }
750
751    /// Exercises the `PACK` bulk path with a length that has both full packs and a remainder:
752    /// `Gf2` (sub-byte packing, padded remainder) and `Mersenne107` (vectorization, byte-identical
753    /// full packs, unpadded per-element remainder).
754    #[test]
755    fn test_heap_array_pack_roundtrip_and_layout() {
756        use typenum::{U16, U20};
757
758        use crate::{
759            algebra::field::{binary::Gf2, mersenne::Mersenne107},
760            random::{test_rng, Random},
761            utils::codec::InPlaceCodec,
762        };
763
764        let mut rng = test_rng();
765
766        // Gf2: 20 elems -> `ceil(20 / 8) = 3` packed bytes. The trailing 4 elements are padded to
767        // a full 8-element/1-byte pack instead of a 4-byte per-element tail.
768        assert_eq!(
769            <HeapArray<Gf2, U20> as InPlaceCodec>::ENCODED_SIZE,
770            20usize.div_ceil(8)
771        );
772        let gf2 = HeapArray::<Gf2, U20>::random(&mut rng);
773        let gf2_bytes = gf2.to_inplace_bytes();
774        assert_eq!(gf2_bytes.len(), 3);
775        assert_eq!(
776            HeapArray::<Gf2, U20>::from_inplace_bytes(&gf2_bytes).unwrap(),
777            gf2
778        );
779
780        // Mersenne107 with a length that's an exact multiple of `PACK` (8): no padding involved,
781        // so the packed bytes stay byte-identical to a manual per-element concatenation.
782        let mers16 = HeapArray::<Mersenne107, U16>::random(&mut rng);
783        let mers16_bytes = mers16.to_inplace_bytes();
784        let reference: Vec<u8> = mers16.iter().flat_map(|e| e.to_inplace_bytes()).collect();
785        assert_eq!(
786            mers16_bytes, reference,
787            "packed Mersenne107 must be byte-identical to the per-element encoding"
788        );
789        assert_eq!(
790            HeapArray::<Mersenne107, U16>::from_inplace_bytes(&mers16_bytes).unwrap(),
791            mers16
792        );
793
794        // Mersenne107 with a remainder: `PACK_BYTES == PACK * ENCODED_SIZE` here (`PACK` is for
795        // vectorized speed, not compression), so the trailing partial group is encoded
796        // element-by-element rather than padded to a full pack — padding would only add dead
797        // bytes. Total stays exactly `20 * 14`.
798        assert_eq!(
799            <HeapArray<Mersenne107, U20> as InPlaceCodec>::ENCODED_SIZE,
800            20 * 14
801        );
802        let mers20 = HeapArray::<Mersenne107, U20>::random(&mut rng);
803        let mers20_bytes = mers20.to_inplace_bytes();
804        assert_eq!(mers20_bytes.len(), 20 * 14);
805        assert_eq!(
806            HeapArray::<Mersenne107, U20>::from_inplace_bytes(&mers20_bytes).unwrap(),
807            mers20
808        );
809    }
810}
primitives/types/heap_array/array.rs

primitives/types/heap_array/
array.rs
