cml_crypto/typed_bytes/
builder.rs

1use crate::typed_bytes::ByteArray;
2use std::marker::PhantomData;
3use std::num::NonZeroUsize;
4
5/// A dynamically created buffer for T
6#[derive(Clone)]
7pub struct ByteBuilder<T> {
8    buffer: Vec<u8>,
9    phantom: PhantomData<T>,
10    expected: Option<NonZeroUsize>,
11}
12
13impl<T> From<ByteBuilder<T>> for Vec<u8> {
14    fn from(bb: ByteBuilder<T>) -> Vec<u8> {
15        bb.buffer
16    }
17}
18
19impl<T> ByteBuilder<T> {
20    /// Create an unconstrained Builder
21    pub fn new() -> Self {
22        ByteBuilder {
23            buffer: Vec::new(),
24            phantom: PhantomData,
25            expected: None,
26        }
27    }
28
29    /// Create a builder of fixed size
30    pub fn new_fixed(size: NonZeroUsize) -> Self {
31        ByteBuilder {
32            buffer: Vec::with_capacity(size.get()),
33            phantom: PhantomData,
34            expected: Some(size),
35        }
36    }
37
38    /// Append an u8 in the builder
39    pub fn u8(self, v: u8) -> Self {
40        let mut buf = self.buffer;
41        buf.push(v);
42        ByteBuilder {
43            buffer: buf,
44            phantom: self.phantom,
45            expected: self.expected,
46        }
47    }
48    /// Append bytes in the builder
49    pub fn bytes(self, v: &[u8]) -> Self {
50        let mut buf = self.buffer;
51        buf.extend_from_slice(v);
52        ByteBuilder {
53            buffer: buf,
54            phantom: self.phantom,
55            expected: self.expected,
56        }
57    }
58
59    /// fold over an iterator
60    pub fn fold<F, I>(self, l: I, f: F) -> Self
61    where
62        I: Iterator,
63        F: FnMut(Self, I::Item) -> Self,
64    {
65        l.fold(self, f)
66    }
67
68    /// write an iterator of maximum 256 items using the closure F
69    ///
70    /// note that the buffer contains a byte to represent the size
71    /// of the list
72    pub fn iter8<F, I>(self, l: I, f: F) -> Self
73    where
74        I: Iterator + ExactSizeIterator,
75        F: FnMut(Self, I::Item) -> Self,
76    {
77        assert!(l.len() < 256);
78        let bb = self.u8(l.len() as u8);
79        l.fold(bb, f)
80    }
81
82    /// write an iterator of maximum 2^16 items using the closure F
83    ///
84    /// note that the buffer contains 2 bytes to represent the size
85    /// of the list
86    pub fn iter16<F, I>(self, l: I, f: F) -> Self
87    where
88        I: Iterator + ExactSizeIterator,
89        F: FnMut(Self, I::Item) -> Self,
90    {
91        assert!(l.len() < 65536);
92        let bb = self.u16(l.len() as u16);
93        l.fold(bb, f)
94    }
95
96    #[allow(clippy::should_implement_trait)]
97    pub fn sub<F, U>(self, f: F) -> Self
98    where
99        F: Fn(ByteBuilder<U>) -> ByteBuilder<U>,
100    {
101        let res = f(ByteBuilder {
102            buffer: self.buffer,
103            phantom: PhantomData,
104            expected: None,
105        });
106        ByteBuilder {
107            buffer: res.buffer,
108            phantom: self.phantom,
109            expected: self.expected,
110        }
111    }
112
113    /// Append an u16 in the builder
114    pub fn u16(self, v: u16) -> Self {
115        self.bytes(&v.to_be_bytes())
116    }
117
118    /// Append an u32 in the builder
119    pub fn u32(self, v: u32) -> Self {
120        self.bytes(&v.to_be_bytes())
121    }
122
123    /// Append an u64 in the builder
124    pub fn u64(self, v: u64) -> Self {
125        self.bytes(&v.to_be_bytes())
126    }
127
128    /// Append an u128 in the builder
129    pub fn u128(self, v: u128) -> Self {
130        self.bytes(&v.to_be_bytes())
131    }
132
133    /// Finalize the buffer and return a fixed ByteArray of T
134    pub fn finalize(self) -> ByteArray<T> {
135        match self.expected {
136            None => ByteArray::from_vec(self.buffer),
137            Some(expected_sz) => {
138                if expected_sz.get() == self.buffer.len() {
139                    ByteArray::from_vec(self.buffer)
140                } else {
141                    panic!(
142                        "internal-error: bytebuilder: expected size {} but got {}",
143                        expected_sz.get(),
144                        self.buffer.len()
145                    )
146                }
147            }
148        }
149    }
150
151    /// Finalize the buffer and return a fixed ByteArray of T
152    pub fn finalize_as_vec(self) -> Vec<u8> {
153        self.buffer
154    }
155}
156
157impl<T> Default for ByteBuilder<T> {
158    fn default() -> Self {
159        Self::new()
160    }
161}