commonware_codec/types/
btree_map.rs

1//! Codec implementations for no_std compatible map types.
2//!
3//! For portability and consistency between architectures,
4//! the size of the map must fit within a [u32].
5
6extern crate alloc;
7
8use crate::{
9    codec::{EncodeSize, Read, Write},
10    error::Error,
11    RangeCfg,
12};
13use alloc::collections::BTreeMap;
14use bytes::{Buf, BufMut};
15use core::cmp::Ordering;
16
17const BTREEMAP_TYPE: &str = "BTreeMap";
18
19/// Read keyed items from [Buf] in ascending order.
20fn read_ordered_map<K, V, F>(
21    buf: &mut impl Buf,
22    len: usize,
23    k_cfg: &K::Cfg,
24    v_cfg: &V::Cfg,
25    mut insert: F,
26    map_type: &'static str,
27) -> Result<(), Error>
28where
29    K: Read + Ord,
30    V: Read,
31    F: FnMut(K, V) -> Option<V>,
32{
33    let mut last: Option<(K, V)> = None;
34    for _ in 0..len {
35        // Read key
36        let key = K::read_cfg(buf, k_cfg)?;
37
38        // Check if keys are in ascending order relative to the previous key
39        if let Some((ref last_key, _)) = last {
40            match key.cmp(last_key) {
41                Ordering::Equal => return Err(Error::Invalid(map_type, "Duplicate key")),
42                Ordering::Less => return Err(Error::Invalid(map_type, "Keys must ascend")),
43                _ => {}
44            }
45        }
46
47        // Read value
48        let value = V::read_cfg(buf, v_cfg)?;
49
50        // Add previous item, if exists
51        if let Some((last_key, last_value)) = last.take() {
52            insert(last_key, last_value);
53        }
54        last = Some((key, value));
55    }
56
57    // Add last item, if exists
58    if let Some((last_key, last_value)) = last {
59        insert(last_key, last_value);
60    }
61
62    Ok(())
63}
64
65// ---------- BTreeMap ----------
66
67impl<K: Ord + Eq + Write, V: Write> Write for BTreeMap<K, V> {
68    fn write(&self, buf: &mut impl BufMut) {
69        self.len().write(buf);
70
71        // Keys are already sorted in BTreeMap, so we can iterate directly
72        for (k, v) in self {
73            k.write(buf);
74            v.write(buf);
75        }
76    }
77}
78
79impl<K: Ord + Eq + EncodeSize, V: EncodeSize> EncodeSize for BTreeMap<K, V> {
80    fn encode_size(&self) -> usize {
81        // Start with the size of the length prefix
82        let mut size = self.len().encode_size();
83
84        // Add the encoded size of each key and value
85        for (k, v) in self {
86            size += k.encode_size();
87            size += v.encode_size();
88        }
89        size
90    }
91}
92
93impl<K: Read + Clone + Ord + Eq, V: Read + Clone> Read for BTreeMap<K, V> {
94    type Cfg = (RangeCfg<usize>, (K::Cfg, V::Cfg));
95
96    fn read_cfg(buf: &mut impl Buf, (range, (k_cfg, v_cfg)): &Self::Cfg) -> Result<Self, Error> {
97        // Read and validate the length prefix
98        let len = usize::read_cfg(buf, range)?;
99        let mut map = BTreeMap::new();
100
101        // Read items in ascending order
102        read_ordered_map(
103            buf,
104            len,
105            k_cfg,
106            v_cfg,
107            |k, v| map.insert(k, v),
108            BTREEMAP_TYPE,
109        )?;
110
111        Ok(map)
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::{Decode, Encode, FixedSize};
119    use bytes::{Bytes, BytesMut};
120    use core::fmt::Debug;
121
122    // Manual round trip test function for BTreeMap with non-default configs
123    fn round_trip_btree<K, V, KCfg, VCfg>(
124        map: &BTreeMap<K, V>,
125        range_cfg: RangeCfg<usize>,
126        k_cfg: KCfg,
127        v_cfg: VCfg,
128    ) where
129        K: Write + EncodeSize + Read<Cfg = KCfg> + Clone + Ord + Eq + PartialEq + Debug,
130        V: Write + EncodeSize + Read<Cfg = VCfg> + Clone + PartialEq + Debug,
131        BTreeMap<K, V>: Read<Cfg = (RangeCfg<usize>, (K::Cfg, V::Cfg))>
132            + Decode<Cfg = (RangeCfg<usize>, (K::Cfg, V::Cfg))>
133            + PartialEq
134            + Write
135            + EncodeSize,
136    {
137        let encoded = map.encode();
138        assert_eq!(encoded.len(), map.encode_size());
139        let config_tuple = (range_cfg, (k_cfg, v_cfg));
140        let decoded = BTreeMap::<K, V>::decode_cfg(encoded, &config_tuple)
141            .expect("decode_cfg failed for BTreeMap");
142        assert_eq!(map, &decoded);
143    }
144
145    #[test]
146    fn test_empty_btreemap() {
147        let map = BTreeMap::<u32, u64>::new();
148        round_trip_btree(&map, (..).into(), (), ());
149        assert_eq!(map.encode_size(), 1); // varint 0
150        let encoded = map.encode();
151        assert_eq!(encoded, Bytes::from_static(&[0]));
152    }
153
154    #[test]
155    fn test_simple_btreemap_u32_u64() {
156        let mut map = BTreeMap::new();
157        map.insert(1u32, 100u64);
158        map.insert(5u32, 500u64);
159        map.insert(2u32, 200u64);
160        round_trip_btree(&map, (..).into(), (), ());
161        assert_eq!(map.encode_size(), 1 + 3 * (u32::SIZE + u64::SIZE));
162        // Check encoding order (BTreeMap guarantees sorted keys: 1, 2, 5)
163        let mut expected = BytesMut::new();
164        3usize.write(&mut expected); // Map length = 3
165        1u32.write(&mut expected);
166        100u64.write(&mut expected);
167        2u32.write(&mut expected);
168        200u64.write(&mut expected);
169        5u32.write(&mut expected);
170        500u64.write(&mut expected);
171        assert_eq!(map.encode(), expected.freeze());
172    }
173
174    #[test]
175    fn test_large_btreemap() {
176        // Fixed-size items
177        let mut map = BTreeMap::new();
178        for i in 0..1000 {
179            map.insert(i as u16, i as u64 * 2);
180        }
181        round_trip_btree(&map, (0..=1000).into(), (), ());
182
183        // Variable-size items
184        let mut map = BTreeMap::new();
185        for i in 0..1000usize {
186            map.insert(i, 1000usize + i);
187        }
188        round_trip_btree(
189            &map,
190            (0..=1000).into(),
191            (..=1000).into(),
192            (1000..=2000).into(),
193        );
194    }
195
196    #[test]
197    fn test_btreemap_decode_length_limit_exceeded() {
198        let mut map = BTreeMap::new();
199        map.insert(1u32, 100u64);
200        map.insert(5u32, 500u64);
201
202        let encoded = map.encode();
203        let config_tuple = ((0..=1).into(), ((), ()));
204
205        let result = BTreeMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
206        assert!(matches!(result, Err(Error::InvalidLength(2))));
207    }
208
209    #[test]
210    fn test_btreemap_decode_invalid_key_order() {
211        let mut encoded = BytesMut::new();
212        2usize.write(&mut encoded); // Map length = 2
213        5u32.write(&mut encoded); // Key 5
214        500u64.write(&mut encoded); // Value 500
215        2u32.write(&mut encoded); // Key 2 (out of order)
216        200u64.write(&mut encoded); // Value 200
217
218        let range = (..).into();
219        let config_tuple = (range, ((), ()));
220
221        let result = BTreeMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
222        assert!(matches!(
223            result,
224            Err(Error::Invalid("BTreeMap", "Keys must ascend"))
225        ));
226    }
227
228    #[test]
229    fn test_btreemap_decode_duplicate_key() {
230        let mut encoded = BytesMut::new();
231        2usize.write(&mut encoded); // Map length = 2
232        1u32.write(&mut encoded); // Key 1
233        100u64.write(&mut encoded); // Value 100
234        1u32.write(&mut encoded); // Duplicate Key 1
235        200u64.write(&mut encoded); // Value 200
236
237        let range = (..).into();
238        let config_tuple = (range, ((), ()));
239
240        let result = BTreeMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
241        assert!(matches!(
242            result,
243            Err(Error::Invalid("BTreeMap", "Duplicate key"))
244        ));
245    }
246
247    #[test]
248    fn test_btreemap_deterministic_encoding() {
249        // In-order
250        let mut map2 = BTreeMap::new();
251        (0..=1000u32).for_each(|i| {
252            map2.insert(i, i * 2);
253        });
254
255        // Reverse order
256        let mut map1 = BTreeMap::new();
257        (0..=1000u32).rev().for_each(|i| {
258            map1.insert(i, i * 2);
259        });
260
261        assert_eq!(map1.encode(), map2.encode());
262    }
263}
commonware_codec/types/btree_map.rs

commonware_codec/types/
btree_map.rs
