abin/implementation/internal/
stack_bin_builder.rs

1use crate::{SBin, StackBin};
2
3const MAX_LEN: usize = StackBin::max_len();
4
5/// Unfortunately `SmallVec` does not implement the size for 23.
6pub struct StackBinBuilder {
7    vec_excess_capacity: usize,
8    inner: Inner,
9}
10
11enum Inner {
12    Vec(Vec<u8>),
13    Stack { len: usize, array: [u8; MAX_LEN] },
14}
15
16impl StackBinBuilder {
17    #[inline]
18    pub fn new(vec_excess_capacity: usize) -> Self {
19        Self {
20            vec_excess_capacity,
21            inner: Inner::Stack {
22                len: 0,
23                array: [0; MAX_LEN],
24            },
25        }
26    }
27
28    #[inline]
29    pub fn extend_from_slice(&mut self, other: &[u8]) {
30        match &mut self.inner {
31            Inner::Vec(vec) => vec.extend_from_slice(other),
32            Inner::Stack { len, array } => {
33                let other_len = other.len();
34                let resulting_len = len.checked_add(other_len).unwrap();
35                if resulting_len > MAX_LEN {
36                    // we need to use a vec
37                    let mut vec = Vec::with_capacity(resulting_len + self.vec_excess_capacity);
38                    vec.extend_from_slice(&array[0..*len]);
39                    vec.extend_from_slice(other);
40                    self.inner = Inner::Vec(vec);
41                } else {
42                    // ok, still enough for the stack
43                    let start_index = *len;
44                    let end_index = start_index + other.len();
45                    (&mut array[start_index..end_index]).copy_from_slice(other);
46                    *len = resulting_len;
47                }
48            }
49        }
50    }
51
52    /// Tries to extend from slice. Returns `true` if there's still enough space to fit onto the
53    /// stack. Returns `false` if item has not been added, because it does not fit onto the stack.
54    pub fn try_extend_from_slice(&mut self, other: &[u8]) -> bool {
55        match &mut self.inner {
56            Inner::Vec(_vec) => false,
57            Inner::Stack { len, array } => {
58                let other_len = other.len();
59                let resulting_len = len.checked_add(other_len).unwrap();
60                if resulting_len > MAX_LEN {
61                    false
62                } else {
63                    // ok, still enough for the stack
64                    let start_index = *len;
65                    let end_index = start_index + other.len();
66                    (&mut array[start_index..end_index]).copy_from_slice(other);
67                    *len = resulting_len;
68                    true
69                }
70            }
71        }
72    }
73
74    /// only builds a binary if this fits onto the stack. Returns `None` otherwise.
75    pub fn build_stack_only(&self) -> Option<SBin> {
76        match &self.inner {
77            Inner::Vec(_vec) => None,
78            Inner::Stack { len, array } => Some(
79                StackBin::try_from(&array[0..*len])
80                    .expect("This MUST be small enough for the stack."),
81            ),
82        }
83    }
84
85    pub fn build(self) -> Result<SBin, Vec<u8>> {
86        match self.inner {
87            Inner::Vec(vec) => {
88                // too large
89                Err(vec)
90            }
91            Inner::Stack { len, array } => Ok(StackBin::try_from(&array[0..len])
92                .expect("This MUST be small enough for the stack.")),
93        }
94    }
95}