bitcode/serde/
variant.rs

1use crate::coder::{Buffer, Decoder, Encoder, Result, View};
2use crate::fast::{CowSlice, NextUnchecked, PushUnchecked, VecImpl};
3use crate::pack::{pack_bytes, unpack_bytes};
4use alloc::vec::Vec;
5use core::marker::PhantomData;
6use core::num::NonZeroUsize;
7
8#[derive(Default)]
9pub struct VariantEncoder {
10    data: VecImpl<u8>,
11}
12
13impl Encoder<u8> for VariantEncoder {
14    #[inline(always)]
15    fn encode(&mut self, v: &u8) {
16        unsafe { self.data.push_unchecked(*v) };
17    }
18}
19
20impl Buffer for VariantEncoder {
21    fn collect_into(&mut self, out: &mut Vec<u8>) {
22        pack_bytes(self.data.as_mut_slice(), out);
23        self.data.clear();
24    }
25
26    fn reserve(&mut self, additional: NonZeroUsize) {
27        self.data.reserve(additional.get());
28    }
29}
30
31#[derive(Default)]
32pub struct VariantDecoder<'a> {
33    variants: CowSlice<'a, u8>,
34    histogram: Vec<usize>,
35    spooky: PhantomData<&'a ()>,
36}
37
38impl VariantDecoder<'_> {
39    pub fn length(&self, variant_index: u8) -> usize {
40        self.histogram[variant_index as usize]
41    }
42
43    /// Returns the max variant index if there were any variants.
44    pub fn max_variant_index(&self) -> Option<u8> {
45        self.histogram.len().checked_sub(1).map(|v| v as u8)
46    }
47}
48
49impl<'a> View<'a> for VariantDecoder<'a> {
50    fn populate(&mut self, input: &mut &'a [u8], length: usize) -> Result<()> {
51        unpack_bytes(input, length, &mut self.variants)?;
52        // Safety: unpack_bytes just initialized self.variants with length of `length`.
53        let variants = unsafe { self.variants.as_slice(length) };
54
55        let histogram = crate::histogram::histogram(variants);
56        let len = histogram
57            .iter()
58            .copied()
59            .rposition(|v| v != 0)
60            .map(|i| i + 1)
61            .unwrap_or(0);
62        self.histogram.clear();
63        self.histogram.extend_from_slice(&histogram[..len]);
64        Ok(())
65    }
66}
67
68impl<'a> Decoder<'a, u8> for VariantDecoder<'a> {
69    fn decode(&mut self) -> u8 {
70        unsafe { self.variants.mut_slice().next_unchecked() }
71    }
72}