bitcode/
int.rs

1use crate::coder::{Buffer, Decoder, Encoder, Result, View};
2use crate::error::err;
3use crate::fast::{CowSlice, NextUnchecked, PushUnchecked, SliceImpl, Unaligned, VecImpl};
4use crate::pack_ints::{pack_ints, unpack_ints, Int};
5use alloc::vec::Vec;
6use bytemuck::{CheckedBitPattern, NoUninit, Pod};
7use core::marker::PhantomData;
8use core::num::NonZeroUsize;
9
10#[derive(Default)]
11pub struct IntEncoder<T>(VecImpl<T>);
12
13/// Makes IntEncoder<u32> able to encode i32/f32/char.
14impl<T: Int, P: NoUninit> Encoder<P> for IntEncoder<T> {
15    #[inline(always)]
16    fn as_primitive(&mut self) -> Option<&mut VecImpl<P>> {
17        use core::mem::*;
18        assert_eq!(align_of::<T>(), align_of::<P>());
19        assert_eq!(size_of::<T>(), size_of::<P>());
20        // Safety: size/align are equal, T: Int implies Pod, and caller isn't reading P which may be NonZero.
21        unsafe { Some(transmute(&mut self.0)) }
22    }
23
24    #[inline(always)]
25    fn encode(&mut self, p: &P) {
26        let t = bytemuck::must_cast(*p);
27        unsafe { self.0.push_unchecked(t) };
28    }
29}
30
31impl<T: Int> Buffer for IntEncoder<T> {
32    fn collect_into(&mut self, out: &mut Vec<u8>) {
33        pack_ints(self.0.as_mut_slice(), out);
34        self.0.clear();
35    }
36
37    fn reserve(&mut self, additional: NonZeroUsize) {
38        self.0.reserve(additional.get());
39    }
40}
41
42#[derive(Default)]
43pub struct IntDecoder<'a, T: Int>(CowSlice<'a, T::Une>);
44
45impl<'a, T: Int> IntDecoder<'a, T> {
46    // For CheckedIntDecoder/LengthDecoder.
47    pub(crate) fn borrowed_clone<'me: 'a>(&'me self) -> IntDecoder<'me, T> {
48        let mut cow = CowSlice::default();
49        cow.set_borrowed_slice_impl(self.0.ref_slice().clone());
50        Self(cow)
51    }
52}
53
54impl<'a, T: Int> View<'a> for IntDecoder<'a, T> {
55    fn populate(&mut self, input: &mut &'a [u8], length: usize) -> Result<()> {
56        unpack_ints::<T>(input, length, &mut self.0)?;
57        Ok(())
58    }
59}
60
61// Makes IntDecoder<u32> able to decode i32/f32 (but not char since it can fail).
62impl<'a, T: Int, P: Pod> Decoder<'a, P> for IntDecoder<'a, T> {
63    #[inline(always)]
64    fn as_primitive(&mut self) -> Option<&mut SliceImpl<Unaligned<P>>> {
65        Some(self.0.mut_slice().cast())
66    }
67
68    #[inline(always)]
69    fn decode(&mut self) -> P {
70        let v = unsafe { self.0.mut_slice().next_unchecked() };
71        bytemuck::must_cast(v)
72    }
73}
74
75/// For NonZeroU32, char, etc.
76pub struct CheckedIntDecoder<'a, C, I: Int>(IntDecoder<'a, I>, PhantomData<C>);
77
78// Can't bound C: Default since NonZeroU32/char don't implement it.
79impl<C, I: Int> Default for CheckedIntDecoder<'_, C, I> {
80    fn default() -> Self {
81        Self(Default::default(), Default::default())
82    }
83}
84
85impl<'a, C: CheckedBitPattern, I: Int> View<'a> for CheckedIntDecoder<'a, C, I>
86where
87    <C as CheckedBitPattern>::Bits: Pod,
88{
89    fn populate(&mut self, input: &mut &'a [u8], length: usize) -> Result<()> {
90        self.0.populate(input, length)?;
91
92        let mut decoder = self.0.borrowed_clone();
93        // Optimizes much better than Iterator::any.
94        if (0..length)
95            .filter(|_| !C::is_valid_bit_pattern(&decoder.decode()))
96            .count()
97            != 0
98        {
99            return err("invalid bit pattern");
100        }
101        Ok(())
102    }
103}
104
105impl<'a, C: CheckedBitPattern + Send + Sync, I: Int> Decoder<'a, C> for CheckedIntDecoder<'a, C, I>
106where
107    <C as CheckedBitPattern>::Bits: Pod,
108{
109    #[inline(always)]
110    fn as_primitive(&mut self) -> Option<&mut SliceImpl<Unaligned<C>>> {
111        self.0
112            .as_primitive()
113            .map(|p: &mut SliceImpl<'_, Unaligned<I>>| {
114                let p = p.cast::<Unaligned<C::Bits>>();
115                // Safety: `Unaligned<C::Bits>` and `Unaligned<C>` have the same layout and populate
116                // ensured C's bit pattern is valid.
117                unsafe { core::mem::transmute(p) }
118            })
119    }
120
121    #[inline(always)]
122    fn decode(&mut self) -> C {
123        let v: I = self.0.decode();
124        let v: C::Bits = bytemuck::must_cast(v);
125        // Safety: C::Bits and C have the same layout and populate ensured C's bit pattern is valid.
126        unsafe { core::mem::transmute_copy(&v) }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use crate::{decode, encode};
133    use alloc::vec::Vec;
134    use core::num::NonZeroU32;
135
136    #[test]
137    fn non_zero_u32() {
138        assert!(decode::<NonZeroU32>(&encode(&0u32)).is_err());
139        assert!(decode::<NonZeroU32>(&encode(&1u32)).is_ok());
140    }
141
142    #[test]
143    fn char_() {
144        assert!(decode::<char>(&encode(&u32::MAX)).is_err());
145        assert!(decode::<char>(&encode(&0u32)).is_ok());
146    }
147
148    fn bench_data() -> Vec<u16> {
149        crate::random_data(1000)
150    }
151    crate::bench_encode_decode!(u16_vec: Vec<_>);
152}
153
154#[cfg(test)]
155mod test2 {
156    use alloc::vec::Vec;
157
158    fn bench_data() -> Vec<Vec<u16>> {
159        crate::random_data::<u8>(125)
160            .into_iter()
161            .map(|n| (0..n / 54).map(|_| n as u16 * 255).collect())
162            .collect()
163    }
164    crate::bench_encode_decode!(u16_vecs: Vec<Vec<_>>);
165}