Skip to main content

commonware_codec/types/
hash_map.rs

1//! Codec implementations for HashMap (requires std).
2//!
3//! For portability and consistency between architectures,
4//! the size of the map must fit within a [u32].
5
6use crate::{
7    Buf, RangeCfg,
8    codec::{BufsMut, EncodeSize, Read, Write},
9    error::Error,
10    types::read_ordered_map,
11};
12use bytes::BufMut;
13use std::{collections::HashMap, hash::Hash};
14
15const HASHMAP_TYPE: &str = "HashMap";
16
17impl<K: Ord + Write, V: Write, S> Write for HashMap<K, V, S> {
18    fn write(&self, buf: &mut impl BufMut) {
19        self.len().write(buf);
20
21        // Sort the keys to ensure deterministic encoding
22        let mut entries: Vec<_> = self.iter().collect();
23        entries.sort_unstable_by(|a, b| a.0.cmp(b.0));
24        for (k, v) in entries {
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        // Sort the keys to ensure deterministic encoding
34        let mut entries: Vec<_> = self.iter().collect();
35        entries.sort_unstable_by(|a, b| a.0.cmp(b.0));
36        for (k, v) in entries {
37            k.write_bufs(buf);
38            v.write_bufs(buf);
39        }
40    }
41}
42
43impl<K: EncodeSize, V: EncodeSize, S> EncodeSize for HashMap<K, V, S> {
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        // Note: Iteration order doesn't matter for size calculation.
50        for (k, v) in self {
51            size += k.encode_size();
52            size += v.encode_size();
53        }
54        size
55    }
56
57    fn encode_inline_size(&self) -> usize {
58        // Start with the size of the length prefix
59        let mut size = self.len().encode_size();
60
61        // Add the encoded size of each key and value
62        // Note: Iteration order doesn't matter for size calculation.
63        for (k, v) in self {
64            size += k.encode_inline_size();
65            size += v.encode_inline_size();
66        }
67        size
68    }
69}
70
71// Read implementation for HashMap
72impl<K: Read + Ord + Hash, V: Read> Read for HashMap<K, V> {
73    type Cfg = (RangeCfg<usize>, (K::Cfg, V::Cfg));
74
75    fn read_cfg(buf: &mut impl Buf, (range, (k_cfg, v_cfg)): &Self::Cfg) -> Result<Self, Error> {
76        // Read and validate the length prefix
77        let len = usize::read_cfg(buf, range)?;
78        let mut map = Self::with_capacity(len.min(buf.remaining()));
79
80        // Read items in ascending order
81        read_ordered_map(
82            buf,
83            len,
84            k_cfg,
85            v_cfg,
86            |k, v| map.insert(k, v),
87            HASHMAP_TYPE,
88        )?;
89
90        Ok(map)
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::{Decode, Encode, FixedSize, types::tests::TrackingWriteBuf};
98    use bytes::{Bytes, BytesMut};
99    use std::{collections::BTreeMap, fmt::Debug};
100
101    // Manual round trip test function for HashMap with non-default configs
102    fn round_trip_hash<K, V, KCfg, VCfg>(
103        map: &HashMap<K, V>,
104        range_cfg: RangeCfg<usize>,
105        k_cfg: KCfg,
106        v_cfg: VCfg,
107    ) where
108        K: Write + EncodeSize + Read<Cfg = KCfg> + Ord + Hash + Debug,
109        V: Write + EncodeSize + Read<Cfg = VCfg> + PartialEq + Debug,
110        HashMap<K, V>: Read<Cfg = (RangeCfg<usize>, (K::Cfg, V::Cfg))>
111            + Decode<Cfg = (RangeCfg<usize>, (K::Cfg, V::Cfg))>
112            + PartialEq
113            + Write
114            + EncodeSize,
115    {
116        let encoded = map.encode();
117        assert_eq!(encoded.len(), map.encode_size());
118        let config_tuple = (range_cfg, (k_cfg, v_cfg));
119        let decoded = HashMap::<K, V>::decode_cfg(encoded, &config_tuple)
120            .expect("decode_cfg failed for HashMap");
121        assert_eq!(map, &decoded);
122    }
123
124    // --- HashMap Tests ---
125
126    #[test]
127    fn test_empty_hashmap() {
128        let map = HashMap::<u32, u64>::new();
129        round_trip_hash(&map, (..).into(), (), ());
130        assert_eq!(map.encode_size(), 1);
131        let encoded = map.encode();
132        assert_eq!(encoded, 0usize.encode());
133    }
134
135    #[test]
136    fn test_simple_hashmap_u32_u64() {
137        let mut map = HashMap::new();
138        map.insert(1u32, 100u64);
139        map.insert(5u32, 500u64);
140        map.insert(2u32, 200u64);
141        round_trip_hash(&map, (..).into(), (), ());
142        assert_eq!(map.encode_size(), 1 + 3 * (u32::SIZE + u64::SIZE));
143    }
144
145    #[test]
146    fn test_large_hashmap() {
147        // Fixed-size items
148        let mut map = HashMap::new();
149        for i in 0..1000 {
150            map.insert(i as u16, i as u64 * 2);
151        }
152        round_trip_hash(&map, (0..=1000).into(), (), ());
153
154        // Variable-size items
155        let mut map = HashMap::new();
156        for i in 0..1000usize {
157            map.insert(i, 1000usize + i);
158        }
159        round_trip_hash(
160            &map,
161            (0..=1000).into(),
162            (..=1000).into(),
163            (1000..=2000).into(),
164        );
165    }
166
167    #[test]
168    fn test_hashmap_with_variable_values() {
169        let mut map = HashMap::new();
170        map.insert(Bytes::from_static(b"apple"), vec![1, 2]);
171        map.insert(Bytes::from_static(b"banana"), vec![3, 4, 5]);
172        map.insert(Bytes::from_static(b"cherry"), vec![]);
173
174        let map_range = RangeCfg::from(0..=10);
175        let key_range = RangeCfg::from(..=10);
176        let val_range = RangeCfg::from(0..=100);
177
178        round_trip_hash(&map, map_range, key_range, (val_range, ()));
179    }
180
181    #[test]
182    fn test_hashmap_decode_length_limit_exceeded() {
183        let mut map = HashMap::new();
184        map.insert(1u32, 100u64);
185        map.insert(5u32, 500u64);
186
187        let encoded = map.encode();
188        let config_tuple = ((0..=1).into(), ((), ()));
189
190        let result = HashMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
191        assert!(matches!(result, Err(Error::InvalidLength(2))));
192    }
193
194    #[test]
195    fn test_hashmap_decode_value_length_limit_exceeded() {
196        let mut map = HashMap::new();
197        map.insert(Bytes::from_static(b"key1"), vec![1u8, 2u8, 3u8, 4u8, 5u8]);
198
199        let key_range = RangeCfg::from(..=10);
200        let map_range = RangeCfg::from(0..=10);
201        let restrictive_val_range = RangeCfg::from(0..=3);
202
203        let encoded = map.encode();
204        let config_tuple = (map_range, (key_range, (restrictive_val_range, ())));
205        let result = HashMap::<Bytes, Vec<u8>>::decode_cfg(encoded, &config_tuple);
206
207        assert!(matches!(result, Err(Error::InvalidLength(5))));
208    }
209
210    #[test]
211    fn test_hashmap_decode_invalid_key_order() {
212        let mut encoded = BytesMut::new();
213        2usize.write(&mut encoded); // Map length = 2
214        5u32.write(&mut encoded); // Key 5
215        500u64.write(&mut encoded); // Value 500
216        2u32.write(&mut encoded); // Key 2 (out of order)
217        200u64.write(&mut encoded); // Value 200
218
219        let range = (..).into();
220        let config_tuple = (range, ((), ()));
221
222        let result = HashMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
223        assert!(matches!(
224            result,
225            Err(Error::Invalid("HashMap", "Keys must ascend"))
226        ));
227    }
228
229    #[test]
230    fn test_hashmap_decode_duplicate_key() {
231        let mut encoded = BytesMut::new();
232        2usize.write(&mut encoded); // Map length = 2
233        1u32.write(&mut encoded); // Key 1
234        100u64.write(&mut encoded); // Value 100
235        1u32.write(&mut encoded); // Duplicate Key 1
236        200u64.write(&mut encoded); // Value 200
237
238        let range = (..).into();
239        let config_tuple = (range, ((), ()));
240
241        let result = HashMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
242        assert!(matches!(
243            result,
244            Err(Error::Invalid("HashMap", "Duplicate key"))
245        ));
246    }
247
248    #[test]
249    fn test_hashmap_decode_end_of_buffer_key() {
250        let mut map = HashMap::new();
251        map.insert(1u32, 100u64);
252        map.insert(5u32, 500u64);
253
254        let mut encoded = map.encode();
255        encoded.truncate(map.encode_size() - 10); // Truncate during last key/value pair
256
257        let range = (..).into();
258        let config_tuple = (range, ((), ()));
259        let result = HashMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
260        assert!(matches!(result, Err(Error::EndOfBuffer)));
261    }
262
263    #[test]
264    fn test_hashmap_decode_end_of_buffer_value() {
265        let mut map = HashMap::new();
266        map.insert(1u32, 100u64);
267        map.insert(5u32, 500u64);
268
269        let mut encoded = map.encode();
270        encoded.truncate(map.encode_size() - 4); // Truncate during last value
271
272        let range = RangeCfg::from(..);
273        let config_tuple = (range, ((), ()));
274        let result = HashMap::<u32, u64>::decode_cfg(encoded, &config_tuple);
275        assert!(matches!(result, Err(Error::EndOfBuffer)));
276    }
277
278    #[test]
279    fn test_hashmap_decode_extra_data() {
280        let mut map = HashMap::new();
281        map.insert(1u32, 100u64);
282
283        let mut encoded = map.encode_mut();
284        encoded.put_u8(0xFF); // Add extra byte
285
286        // Use decode_cfg which enforces buffer is fully consumed
287        let config_tuple = ((..).into(), ((), ()));
288        let result = HashMap::<u32, u64>::decode_cfg(encoded.clone(), &config_tuple);
289        assert!(matches!(result, Err(Error::ExtraData(1))));
290
291        // Verify that read_cfg would succeed (doesn't check for extra data)
292        let read_result = HashMap::<u32, u64>::read_cfg(&mut encoded, &config_tuple);
293        assert!(read_result.is_ok());
294        let decoded_map = read_result.unwrap();
295        assert_eq!(decoded_map.len(), 1);
296        assert_eq!(decoded_map.get(&1u32), Some(&100u64));
297    }
298
299    #[test]
300    fn test_hashmap_deterministic_encoding() {
301        // In-order
302        let mut map2 = HashMap::new();
303        (0..=1000u32).for_each(|i| {
304            map2.insert(i, i * 2);
305        });
306
307        // Reverse order
308        let mut map1 = HashMap::new();
309        (0..=1000u32).rev().for_each(|i| {
310            map1.insert(i, i * 2);
311        });
312
313        assert_eq!(map1.encode(), map2.encode());
314    }
315
316    #[test]
317    fn test_hashmap_encoding_matches_ordered_map() {
318        // Ordered maps provide the canonical encoding across collection sizes.
319        for len in [0, 1, 20, 128, 1024] {
320            let ordered: BTreeMap<_, _> = (0..len)
321                .map(|key: u32| (key, Bytes::copy_from_slice(&key.to_le_bytes())))
322                .collect();
323            let expected = ordered.encode();
324
325            // Insertion order must not affect the encoded bytes.
326            for reverse in [false, true] {
327                let mut entries: Vec<_> = ordered.iter().collect();
328                if reverse {
329                    entries.reverse();
330                }
331                let map: HashMap<_, _> = entries
332                    .into_iter()
333                    .map(|(&key, value)| (key, value.clone()))
334                    .collect();
335                assert_eq!(map.encode(), expected);
336
337                // Bytes values exercise external chunks alongside inline keys and lengths.
338                let mut buf = TrackingWriteBuf::new();
339                map.write_bufs(&mut buf);
340                assert_eq!(buf.freeze(), expected);
341            }
342        }
343    }
344
345    #[test]
346    fn test_hashmap_conformity() {
347        // Case 1: Empty HashMap<u8, u16>
348        let map1 = HashMap::<u8, u16>::new();
349        let mut expected1 = BytesMut::new();
350        0usize.write(&mut expected1); // Length 0
351        assert_eq!(map1.encode(), expected1.freeze());
352
353        // Case 2: Simple HashMap<u8, u16>
354        // Keys are sorted for encoding: 1, 2
355        let mut map2 = HashMap::<u8, u16>::new();
356        map2.insert(2u8, 0xBBBBu16); // Inserted out of order
357        map2.insert(1u8, 0xAAAAu16);
358
359        let mut expected2 = BytesMut::new();
360        2usize.write(&mut expected2); // Length 2
361        1u8.write(&mut expected2); // Key 1
362        0xAAAAu16.write(&mut expected2); // Value for key 1
363        2u8.write(&mut expected2); // Key 2
364        0xBBBBu16.write(&mut expected2); // Value for key 2
365        assert_eq!(map2.encode(), expected2.freeze());
366
367        // Case 3: HashMap<u16, bool>
368        // Keys are sorted for encoding: 0x0101, 0x0202, 0x0303
369        let mut map3 = HashMap::<u16, bool>::new();
370        map3.insert(0x0303u16, true);
371        map3.insert(0x0101u16, false);
372        map3.insert(0x0202u16, true);
373
374        let mut expected3 = BytesMut::new();
375        3usize.write(&mut expected3); // Length 3
376        0x0101u16.write(&mut expected3); // Key 0x0101
377        false.write(&mut expected3); // Value false (0x00)
378        0x0202u16.write(&mut expected3); // Key 0x0202
379        true.write(&mut expected3); // Value true (0x01)
380        0x0303u16.write(&mut expected3); // Key 0x0303
381        true.write(&mut expected3); // Value true (0x01)
382        assert_eq!(map3.encode(), expected3.freeze());
383
384        // Case 4: HashMap with Bytes as key and Vec<u8> as value
385        // Keys are sorted for encoding: "a", "b"
386        let mut map4 = HashMap::<Bytes, Vec<u8>>::new();
387        map4.insert(Bytes::from_static(b"b"), vec![20u8, 21u8]);
388        map4.insert(Bytes::from_static(b"a"), vec![10u8]);
389
390        let mut expected4 = BytesMut::new();
391        2usize.write(&mut expected4); // Map length = 2
392
393        // Key "a" (length 1, 'a')
394        Bytes::from_static(b"a").write(&mut expected4);
395        // Value vec![10u8] (length 1, 10u8)
396        vec![10u8].write(&mut expected4);
397
398        // Key "b" (length 1, 'b')
399        Bytes::from_static(b"b").write(&mut expected4);
400        // Value vec![20u8, 21u8] (length 2, 20u8, 21u8)
401        vec![20u8, 21u8].write(&mut expected4);
402
403        assert_eq!(map4.encode(), expected4.freeze());
404    }
405
406    #[test]
407    fn test_hashmap_malicious_allocation() {
408        // A length prefix advertising the maximum wire-encodable length (u32::MAX) with a
409        // single-entry payload must fail with [Error::EndOfBuffer] without attempting the
410        // advertised allocation.
411        let mut malicious_buf = BytesMut::new();
412        (u32::MAX as usize).write(&mut malicious_buf);
413        1u32.write(&mut malicious_buf);
414        100u64.write(&mut malicious_buf);
415
416        let result =
417            HashMap::<u32, u64>::decode_cfg(malicious_buf.freeze(), &((..).into(), ((), ())));
418        assert!(matches!(result, Err(Error::EndOfBuffer)));
419    }
420
421    #[cfg(feature = "arbitrary")]
422    mod conformance {
423        use super::*;
424        use crate::conformance::CodecConformance;
425
426        commonware_conformance::conformance_tests! {
427            CodecConformance<HashMap<u32, u32>>,
428        }
429    }
430}