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
40impl<T> Decodable for Vec<T>
41where
42    T: Decodable + 'static,
43{
44    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
45        d: &mut D,
46        modules: &ModuleDecoderRegistry,
47    ) -> Result<Self, DecodeError> {
48        const CHUNK_SIZE: usize = 64 * 1024;
49
50        if TypeId::of::<T>() == TypeId::of::<u8>() {
51            let len =
52                u64::consensus_decode_partial_from_finite_reader(d, &ModuleRegistry::default())?;
53
54            let mut len: usize =
55                usize::try_from(len).map_err(|_| DecodeError::from_str("size exceeds memory"))?;
56
57            let mut bytes = vec![];
58
59            // Adapted from <https://github.com/rust-bitcoin/rust-bitcoin/blob/e2b9555070d9357fb552e56085fb6fb3f0274560/bitcoin/src/consensus/encode.rs#L667-L674>
60            while len > 0 {
61                let chunk_start = bytes.len();
62                let chunk_size = core::cmp::min(len, CHUNK_SIZE);
63                let chunk_end = chunk_start + chunk_size;
64                bytes.resize(chunk_end, 0u8);
65                d.read_exact(&mut bytes[chunk_start..chunk_end])?;
66                len -= chunk_size;
67            }
68
69            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
70            return Ok(unsafe { std::mem::transmute::<Vec<u8>, Self>(bytes) });
71        }
72        let len = u64::consensus_decode_partial_from_finite_reader(d, modules)?;
73
74        // `collect` under the hood uses `FromIter::from_iter`, which can potentially be
75        // backed by code like:
76        // <https://github.com/rust-lang/rust/blob/fe03b46ee4688a99d7155b4f9dcd875b6903952d/library/alloc/src/vec/spec_from_iter_nested.rs#L31>
77        // This can take `size_hint` from input iterator and pre-allocate memory
78        // upfront with `Vec::with_capacity`. Because of that untrusted `len`
79        // should not be used directly.
80        let cap_len = std::cmp::min(8_000 / std::mem::size_of::<T>() as u64, len);
81
82        // Up to a cap, use the (potentially specialized for better perf in stdlib)
83        // `from_iter`.
84        let mut v: Self = (0..cap_len)
85            .map(|_| T::consensus_decode_partial_from_finite_reader(d, modules))
86            .collect::<Result<Self, DecodeError>>()?;
87
88        // Add any excess manually avoiding any surprises.
89        while (v.len() as u64) < len {
90            v.push(T::consensus_decode_partial_from_finite_reader(d, modules)?);
91        }
92
93        assert_eq!(v.len() as u64, len);
94
95        Ok(v)
96    }
97}
98
99impl<T> Encodable for VecDeque<T>
100where
101    T: Encodable + 'static,
102{
103    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
104        (self.len() as u64).consensus_encode(writer)?;
105        for i in self {
106            i.consensus_encode(writer)?;
107        }
108        Ok(())
109    }
110}
111
112impl<T> Decodable for VecDeque<T>
113where
114    T: Decodable + 'static,
115{
116    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
117        d: &mut D,
118        modules: &ModuleDecoderRegistry,
119    ) -> Result<Self, DecodeError> {
120        Ok(Self::from(
121            Vec::<T>::consensus_decode_partial_from_finite_reader(d, modules)?,
122        ))
123    }
124}
125
126impl<T, const SIZE: usize> Encodable for [T; SIZE]
127where
128    T: Encodable + 'static,
129{
130    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
131        if TypeId::of::<T>() == TypeId::of::<u8>() {
132            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
133            let bytes = unsafe { std::mem::transmute::<&[T; SIZE], &[u8; SIZE]>(self) };
134            writer.write_all(bytes)?;
135            return Ok(());
136        }
137
138        for item in self {
139            item.consensus_encode(writer)?;
140        }
141        Ok(())
142    }
143}
144
145impl<T, const SIZE: usize> Decodable for [T; SIZE]
146where
147    T: Decodable + Debug + Default + Copy + 'static,
148{
149    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
150        d: &mut D,
151        modules: &ModuleDecoderRegistry,
152    ) -> Result<Self, DecodeError> {
153        // From <https://github.com/rust-lang/rust/issues/61956>
154        unsafe fn horribe_array_transmute_workaround<const N: usize, A, B>(
155            mut arr: [A; N],
156        ) -> [B; N] {
157            let ptr = std::ptr::from_mut(&mut arr).cast::<[B; N]>();
158            let res = unsafe { ptr.read() };
159            core::mem::forget(arr);
160            res
161        }
162
163        if TypeId::of::<T>() == TypeId::of::<u8>() {
164            let mut bytes = [0u8; SIZE];
165            d.read_exact(bytes.as_mut_slice())?;
166
167            // unsafe: we've just checked that T is `u8` so the transmute here is a no-op
168            return Ok(unsafe { horribe_array_transmute_workaround(bytes) });
169        }
170
171        // todo: impl without copy
172        let mut data = [T::default(); SIZE];
173        for item in &mut data {
174            *item = T::consensus_decode_partial_from_finite_reader(d, modules)?;
175        }
176        Ok(data)
177    }
178}
179
180impl<K, V> Encodable for BTreeMap<K, V>
181where
182    K: Encodable,
183    V: Encodable,
184{
185    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
186        (self.len() as u64).consensus_encode(writer)?;
187        for (k, v) in self {
188            k.consensus_encode(writer)?;
189            v.consensus_encode(writer)?;
190        }
191        Ok(())
192    }
193}
194
195impl<K, V> Decodable for BTreeMap<K, V>
196where
197    K: Decodable + Ord,
198    V: Decodable,
199{
200    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
201        d: &mut D,
202        modules: &ModuleDecoderRegistry,
203    ) -> Result<Self, DecodeError> {
204        let mut res = Self::new();
205        let len = u64::consensus_decode_partial_from_finite_reader(d, modules)?;
206        for _ in 0..len {
207            let k = K::consensus_decode_partial_from_finite_reader(d, modules)?;
208            if res
209                .last_key_value()
210                .is_some_and(|(prev_key, _v)| k <= *prev_key)
211            {
212                return Err(DecodeError::from_str("Non-canonical encoding"));
213            }
214            let v = V::consensus_decode_partial_from_finite_reader(d, modules)?;
215            if res.insert(k, v).is_some() {
216                return Err(DecodeError::from_str("Duplicate key"));
217            }
218        }
219        Ok(res)
220    }
221}
222
223impl<K> Encodable for BTreeSet<K>
224where
225    K: Encodable,
226{
227    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
228        (self.len() as u64).consensus_encode(writer)?;
229        for k in self {
230            k.consensus_encode(writer)?;
231        }
232        Ok(())
233    }
234}
235
236impl<K> Decodable for BTreeSet<K>
237where
238    K: Decodable + Ord,
239{
240    fn consensus_decode_partial_from_finite_reader<D: std::io::Read>(
241        d: &mut D,
242        modules: &ModuleDecoderRegistry,
243    ) -> Result<Self, DecodeError> {
244        let mut res = Self::new();
245        let len = u64::consensus_decode_partial_from_finite_reader(d, modules)?;
246        for _ in 0..len {
247            let k = K::consensus_decode_partial_from_finite_reader(d, modules)?;
248            if res.last().is_some_and(|prev_key| k <= *prev_key) {
249                return Err(DecodeError::from_str("Non-canonical encoding"));
250            }
251            if !res.insert(k) {
252                return Err(DecodeError::from_str("Duplicate key"));
253            }
254        }
255        Ok(res)
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262    use crate::encoding::tests::test_roundtrip_expected;
263
264    #[test_log::test]
265    fn test_lists() {
266        // The length of the list is encoded before the elements. It is encoded as a
267        // variable length integer, but for lists with a length less than 253, it's
268        // encoded as a single byte.
269        test_roundtrip_expected(&vec![1u8, 2, 3], &[3u8, 1, 2, 3]);
270        test_roundtrip_expected(&vec![1u16, 2, 3], &[3u8, 1, 2, 3]);
271        test_roundtrip_expected(&vec![1u32, 2, 3], &[3u8, 1, 2, 3]);
272        test_roundtrip_expected(&vec![1u64, 2, 3], &[3u8, 1, 2, 3]);
273
274        // Empty list should be encoded as a single byte 0.
275        test_roundtrip_expected::<Vec<u8>>(&vec![], &[0u8]);
276        test_roundtrip_expected::<Vec<u16>>(&vec![], &[0u8]);
277        test_roundtrip_expected::<Vec<u32>>(&vec![], &[0u8]);
278        test_roundtrip_expected::<Vec<u64>>(&vec![], &[0u8]);
279
280        // A length prefix greater than the number of elements should return an error.
281        let buf = [4u8, 1, 2, 3];
282        assert!(Vec::<u8>::consensus_decode_whole(&buf, &ModuleRegistry::default()).is_err());
283        assert!(Vec::<u16>::consensus_decode_whole(&buf, &ModuleRegistry::default()).is_err());
284        assert!(VecDeque::<u8>::consensus_decode_whole(&buf, &ModuleRegistry::default()).is_err());
285        assert!(VecDeque::<u16>::consensus_decode_whole(&buf, &ModuleRegistry::default()).is_err());
286
287        // A length prefix less than the number of elements should skip elements beyond
288        // the encoded length.
289        let buf = [2u8, 1, 2, 3];
290        assert_eq!(
291            Vec::<u8>::consensus_decode_partial(&mut &buf[..], &ModuleRegistry::default()).unwrap(),
292            vec![1u8, 2]
293        );
294        assert_eq!(
295            Vec::<u16>::consensus_decode_partial(&mut &buf[..], &ModuleRegistry::default())
296                .unwrap(),
297            vec![1u16, 2]
298        );
299        assert_eq!(
300            VecDeque::<u8>::consensus_decode_partial(&mut &buf[..], &ModuleRegistry::default())
301                .unwrap(),
302            vec![1u8, 2]
303        );
304        assert_eq!(
305            VecDeque::<u16>::consensus_decode_partial(&mut &buf[..], &ModuleRegistry::default())
306                .unwrap(),
307            vec![1u16, 2]
308        );
309    }
310
311    #[test_log::test]
312    fn test_btreemap() {
313        test_roundtrip_expected(
314            &BTreeMap::from([("a".to_string(), 1u32), ("b".to_string(), 2)]),
315            &[2, 1, 97, 1, 1, 98, 2],
316        );
317    }
318
319    #[test_log::test]
320    fn test_btreeset() {
321        test_roundtrip_expected(
322            &BTreeSet::from(["a".to_string(), "b".to_string()]),
323            &[2, 1, 97, 1, 98],
324        );
325    }
326}
fedimint_core/encoding/collections.rs

fedimint_core/encoding/
collections.rs
