fedimint_core/encoding/
collections.rs

1use std::any::TypeId;
2use std::collections::{BTreeMap, BTreeSet, VecDeque};
3use std::fmt::Debug;
4
5use crate::module::registry::ModuleRegistry;
6use crate::{Decodable, DecodeError, Encodable, ModuleDecoderRegistry};
7
8impl<T> Encodable for &[T]
9where
10    T: Encodable + 'static,
11{
12    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
13        if TypeId::of::<T>() == TypeId::of::<u8>() {
14            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
15            let bytes = unsafe { std::mem::transmute::<&[T], &[u8]>(self) };
16
17            (bytes.len() as u64).consensus_encode(writer)?;
18            writer.write_all(bytes)?;
19            return Ok(());
20        }
21
22        (self.len() as u64).consensus_encode(writer)?;
23
24        for item in *self {
25            item.consensus_encode(writer)?;
26        }
27        Ok(())
28    }
29}
30
31impl<T> Encodable for Vec<T>
32where
33    T: Encodable + 'static,
34{
35    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
36        (self as &[T]).consensus_encode(writer)
37    }
38}
39
40/// Empty vectors are supported for all element types. Nonempty vectors of
41/// zero-sized types return a decode error without invoking the element decoder.
42/// A finite reader alone cannot bound iteration over elements that consume no
43/// bytes.
44impl<T> Decodable for Vec<T>
45where
46    T: Decodable + 'static,
47{
48    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
49        d: &mut D,
50        modules: &ModuleDecoderRegistry,
51    ) -> Result<Self, DecodeError> {
52        const CHUNK_SIZE: usize = 64 * 1024;
53
54        if TypeId::of::<T>() == TypeId::of::<u8>() {
55            let len =
56                u64::consensus_decode_partial_from_finite_reader(d, &ModuleRegistry::default())?;
57
58            let mut len: usize =
59                usize::try_from(len).map_err(|_| DecodeError::from_str("size exceeds memory"))?;
60
61            let mut bytes = vec![];
62
63            // Adapted from <https://github.com/rust-bitcoin/rust-bitcoin/blob/e2b9555070d9357fb552e56085fb6fb3f0274560/bitcoin/src/consensus/encode.rs#L667-L674>
64            while len > 0 {
65                let chunk_start = bytes.len();
66                let chunk_size = core::cmp::min(len, CHUNK_SIZE);
67                let chunk_end = chunk_start + chunk_size;
68                bytes.resize(chunk_end, 0u8);
69                d.read_exact(&mut bytes[chunk_start..chunk_end])?;
70                len -= chunk_size;
71            }
72
73            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
74            return Ok(unsafe { std::mem::transmute::<Vec<u8>, Self>(bytes) });
75        }
76        let len = u64::consensus_decode_partial_from_finite_reader(d, modules)?;
77
78        if len == 0 {
79            return Ok(Self::new());
80        }
81        if std::mem::size_of::<T>() == 0 {
82            return Err(DecodeError::from_str(
83                "cannot decode nonempty vector of zero-sized elements",
84            ));
85        }
86
87        // `collect` under the hood uses `FromIter::from_iter`, which can potentially be
88        // backed by code like:
89        // <https://github.com/rust-lang/rust/blob/fe03b46ee4688a99d7155b4f9dcd875b6903952d/library/alloc/src/vec/spec_from_iter_nested.rs#L31>
90        // This can take `size_hint` from input iterator and pre-allocate memory
91        // upfront with `Vec::with_capacity`. Because of that untrusted `len`
92        // should not be used directly.
93        let cap_len = std::cmp::min(8_000 / std::mem::size_of::<T>() as u64, len);
94
95        // Up to a cap, use the (potentially specialized for better perf in stdlib)
96        // `from_iter`.
97        let mut v: Self = (0..cap_len)
98            .map(|_| T::consensus_decode_partial_from_finite_reader(d, modules))
99            .collect::<Result<Self, DecodeError>>()?;
100
101        // Add any excess manually avoiding any surprises.
102        while (v.len() as u64) < len {
103            v.push(T::consensus_decode_partial_from_finite_reader(d, modules)?);
104        }
105
106        assert_eq!(v.len() as u64, len);
107
108        Ok(v)
109    }
110}
111
112impl<T> Encodable for VecDeque<T>
113where
114    T: Encodable + 'static,
115{
116    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
117        (self.len() as u64).consensus_encode(writer)?;
118        for i in self {
119            i.consensus_encode(writer)?;
120        }
121        Ok(())
122    }
123}
124
125/// Uses the same zero-sized element policy as [`Vec<T>`]: empty deques are
126/// supported, but nonempty deques of zero-sized types return a decode error.
127impl<T> Decodable for VecDeque<T>
128where
129    T: Decodable + 'static,
130{
131    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
132        d: &mut D,
133        modules: &ModuleDecoderRegistry,
134    ) -> Result<Self, DecodeError> {
135        Ok(Self::from(
136            Vec::<T>::consensus_decode_partial_from_finite_reader(d, modules)?,
137        ))
138    }
139}
140
141impl<T, const SIZE: usize> Encodable for [T; SIZE]
142where
143    T: Encodable + 'static,
144{
145    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
146        if TypeId::of::<T>() == TypeId::of::<u8>() {
147            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
148            let bytes = unsafe { std::mem::transmute::<&[T; SIZE], &[u8; SIZE]>(self) };
149            writer.write_all(bytes)?;
150            return Ok(());
151        }
152
153        for item in self {
154            item.consensus_encode(writer)?;
155        }
156        Ok(())
157    }
158}
159
160impl<T, const SIZE: usize> Decodable for [T; SIZE]
161where
162    T: Decodable + Debug + Default + Copy + 'static,
163{
164    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
165        d: &mut D,
166        modules: &ModuleDecoderRegistry,
167    ) -> Result<Self, DecodeError> {
168        // From <https://github.com/rust-lang/rust/issues/61956>
169        unsafe fn horribe_array_transmute_workaround<const N: usize, A, B>(
170            mut arr: [A; N],
171        ) -> [B; N] {
172            let ptr = std::ptr::from_mut(&mut arr).cast::<[B; N]>();
173            let res = unsafe { ptr.read() };
174            core::mem::forget(arr);
175            res
176        }
177
178        if TypeId::of::<T>() == TypeId::of::<u8>() {
179            let mut bytes = [0u8; SIZE];
180            d.read_exact(bytes.as_mut_slice())?;
181
182            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
183            return Ok(unsafe { horribe_array_transmute_workaround(bytes) });
184        }
185
186        // todo: impl without copy
187        let mut data = [T::default(); SIZE];
188        for item in &mut data {
189            *item = T::consensus_decode_partial_from_finite_reader(d, modules)?;
190        }
191        Ok(data)
192    }
193}
194
195impl<K, V> Encodable for BTreeMap<K, V>
196where
197    K: Encodable,
198    V: Encodable,
199{
200    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
201        (self.len() as u64).consensus_encode(writer)?;
202        for (k, v) in self {
203            k.consensus_encode(writer)?;
204            v.consensus_encode(writer)?;
205        }
206        Ok(())
207    }
208}
209
210impl<K, V> Decodable for BTreeMap<K, V>
211where
212    K: Decodable + Ord,
213    V: Decodable,
214{
215    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
216        d: &mut D,
217        modules: &ModuleDecoderRegistry,
218    ) -> Result<Self, DecodeError> {
219        let mut res = Self::new();
220        let len = u64::consensus_decode_partial_from_finite_reader(d, modules)?;
221        for _ in 0..len {
222            let k = K::consensus_decode_partial_from_finite_reader(d, modules)?;
223            if res
224                .last_key_value()
225                .is_some_and(|(prev_key, _v)| k <= *prev_key)
226            {
227                return Err(DecodeError::from_str("Non-canonical encoding"));
228            }
229            let v = V::consensus_decode_partial_from_finite_reader(d, modules)?;
230            if res.insert(k, v).is_some() {
231                return Err(DecodeError::from_str("Duplicate key"));
232            }
233        }
234        Ok(res)
235    }
236}
237
238impl<K> Encodable for BTreeSet<K>
239where
240    K: Encodable,
241{
242    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
243        (self.len() as u64).consensus_encode(writer)?;
244        for k in self {
245            k.consensus_encode(writer)?;
246        }
247        Ok(())
248    }
249}
250
251impl<K> Decodable for BTreeSet<K>
252where
253    K: Decodable + Ord,
254{
255    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
256        d: &mut D,
257        modules: &ModuleDecoderRegistry,
258    ) -> Result<Self, DecodeError> {
259        let mut res = Self::new();
260        let len = u64::consensus_decode_partial_from_finite_reader(d, modules)?;
261        for _ in 0..len {
262            let k = K::consensus_decode_partial_from_finite_reader(d, modules)?;
263            if res.last().is_some_and(|prev_key| k <= *prev_key) {
264                return Err(DecodeError::from_str("Non-canonical encoding"));
265            }
266            if !res.insert(k) {
267                return Err(DecodeError::from_str("Duplicate key"));
268            }
269        }
270        Ok(res)
271    }
272}
273
274#[cfg(test)]
275mod tests;
276
277#[cfg(test)]
278mod zst_tests;
fedimint_core/encoding/collections.rs

fedimint_core/encoding/
collections.rs
