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

commonware_codec/types/
btree_map.rs
