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

commonware_codec/types/
btree_map.rs
