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

primitives/types/heap_array/
array.rs
