1use 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 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 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 let mut size = self.len().encode_size();
47
48 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 let mut size = self.len().encode_size();
60
61 for (k, v) in self {
64 size += k.encode_inline_size();
65 size += v.encode_inline_size();
66 }
67 size
68 }
69}
70
71impl<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 let len = usize::read_cfg(buf, range)?;
78 let mut map = Self::with_capacity(len.min(buf.remaining()));
79
80 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 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 #[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 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 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); 5u32.write(&mut encoded); 500u64.write(&mut encoded); 2u32.write(&mut encoded); 200u64.write(&mut encoded); 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); 1u32.write(&mut encoded); 100u64.write(&mut encoded); 1u32.write(&mut encoded); 200u64.write(&mut encoded); 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); 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); 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); 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 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 let mut map2 = HashMap::new();
303 (0..=1000u32).for_each(|i| {
304 map2.insert(i, i * 2);
305 });
306
307 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 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 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 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 let map1 = HashMap::<u8, u16>::new();
349 let mut expected1 = BytesMut::new();
350 0usize.write(&mut expected1); assert_eq!(map1.encode(), expected1.freeze());
352
353 let mut map2 = HashMap::<u8, u16>::new();
356 map2.insert(2u8, 0xBBBBu16); map2.insert(1u8, 0xAAAAu16);
358
359 let mut expected2 = BytesMut::new();
360 2usize.write(&mut expected2); 1u8.write(&mut expected2); 0xAAAAu16.write(&mut expected2); 2u8.write(&mut expected2); 0xBBBBu16.write(&mut expected2); assert_eq!(map2.encode(), expected2.freeze());
366
367 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); 0x0101u16.write(&mut expected3); false.write(&mut expected3); 0x0202u16.write(&mut expected3); true.write(&mut expected3); 0x0303u16.write(&mut expected3); true.write(&mut expected3); assert_eq!(map3.encode(), expected3.freeze());
383
384 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); Bytes::from_static(b"a").write(&mut expected4);
395 vec![10u8].write(&mut expected4);
397
398 Bytes::from_static(b"b").write(&mut expected4);
400 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 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}