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

commonware_codec/types/
btree_map.rs
