tree_buf/internal/types/
hashmap.rs

1use crate::prelude::*;
2use std::collections::HashMap;
3use std::hash::{BuildHasher, Hash};
4use std::vec::IntoIter;
5
6#[cfg(feature = "encode")]
7impl<K: Encodable, V: Encodable, S: Default + BuildHasher> Encodable for HashMap<K, V, S> {
8    type EncoderArray = HashMapArrayEncoder<K::EncoderArray, V::EncoderArray, S>;
9    fn encode_root<O: EncodeOptions>(&self, stream: &mut EncoderStream<'_, O>) -> RootTypeId {
10        profile_method!(encode_root);
11
12        encode_usize(self.len(), stream);
13        match self.len() {
14            0 => {}
15            1 => {
16                for key in self.keys() {
17                    stream.encode_with_id(|stream| key.encode_root(stream));
18                }
19                for value in self.values() {
20                    stream.encode_with_id(|stream| value.encode_root(stream));
21                }
22            }
23            _ => {
24                let mut keys_encoder = K::EncoderArray::default();
25                for key in self.keys() {
26                    keys_encoder.buffer_one(key);
27                }
28                stream.encode_with_id(|stream| keys_encoder.flush(stream));
29                // TODO: Verify that iterating keys and values is always the same order!
30                let mut values_encoder = V::EncoderArray::default();
31                for value in self.values() {
32                    values_encoder.buffer_one(value);
33                }
34                stream.encode_with_id(|stream| values_encoder.flush(stream));
35            }
36        }
37
38        RootTypeId::Map
39    }
40}
41
42#[cfg(feature = "decode")]
43impl<K: Decodable + Hash + Eq + Send, V: Decodable + Send, S: Default + BuildHasher> Decodable for HashMap<K, V, S>
44where
45    // Overly verbose because of `?` requiring `From` See also ec4fa3ba-def5-44eb-9065-e80b59530af6
46    DecodeError: From<<<K as Decodable>::DecoderArray as DecoderArray>::Error>,
47    // Overly verbose because of `?` requiring `From` See also ec4fa3ba-def5-44eb-9065-e80b59530af6
48    DecodeError: From<<<V as Decodable>::DecoderArray as DecoderArray>::Error>,
49{
50    type DecoderArray = Option<HashMapArrayDecoder<K::DecoderArray, V::DecoderArray, S>>;
51    fn decode(sticks: DynRootBranch<'_>, options: &impl DecodeOptions) -> DecodeResult<Self> {
52        profile_method!(decode);
53
54        let mut v = Default::default(); // TODO: (Performance) Capacity
55        match sticks {
56            DynRootBranch::Map0 => Ok(v),
57            DynRootBranch::Map1 { key, value } => {
58                let (key, value) = parallel(move || K::decode(*key, options), move || V::decode(*value, options), options);
59                v.insert(key?, value?);
60                Ok(v)
61            }
62            DynRootBranch::Map { len, keys, values } => {
63                let (keys, values) = parallel(|| K::DecoderArray::new(keys, options), || V::DecoderArray::new(values, options), options);
64                let mut keys = keys?;
65                let mut values = values?;
66                for _ in 0..len {
67                    if v.insert(keys.decode_next()?, values.decode_next()?).is_some() {
68                        return Err(DecodeError::InvalidFormat);
69                    }
70                }
71                Ok(v)
72            }
73            _ => Err(DecodeError::SchemaMismatch),
74        }
75    }
76}
77
78#[cfg(feature = "encode")]
79#[derive(Debug, Default)]
80pub struct HashMapArrayEncoder<K, V, S> {
81    len: <u64 as Encodable>::EncoderArray,
82    items: Option<(K, V)>,
83    _marker: Unowned<S>,
84}
85
86#[cfg(feature = "decode")]
87pub struct HashMapArrayDecoder<K, V, S> {
88    len: IntoIter<u64>,
89    keys: K,
90    values: V,
91    _marker: Unowned<S>,
92}
93
94#[cfg(feature = "encode")]
95impl<K: Encodable, V: Encodable, S: Default + BuildHasher> EncoderArray<HashMap<K, V, S>> for HashMapArrayEncoder<K::EncoderArray, V::EncoderArray, S> {
96    fn buffer_one<'a, 'b: 'a>(&'a mut self, value: &'b HashMap<K, V, S>) {
97        profile_method!(EncoderArray::buffer);
98        self.len.buffer_one(&(value.len() as u64));
99        let (keys, values) = self.items.get_or_insert_with(Default::default);
100        for (key, value) in value.iter() {
101            keys.buffer_one(key);
102            values.buffer_one(value);
103        }
104    }
105    fn flush<O: EncodeOptions>(self, stream: &mut EncoderStream<'_, O>) -> ArrayTypeId {
106        profile_method!(flush);
107        let Self { len, items, _marker } = self;
108        if let Some((keys, values)) = items {
109            stream.encode_with_id(|stream| len.flush(stream));
110            stream.encode_with_id(|stream| keys.flush(stream));
111            stream.encode_with_id(|stream| values.flush(stream));
112        } else {
113            stream.encode_with_id(|_| ArrayTypeId::Void);
114        }
115        ArrayTypeId::Map
116    }
117}
118
119#[cfg(feature = "decode")]
120impl<K: DecoderArray, V: DecoderArray, S: Default + BuildHasher> DecoderArray for Option<HashMapArrayDecoder<K, V, S>>
121where
122    K::Decode: Hash + Eq,
123    // Overly verbose because of `?` requiring `From` See also ec4fa3ba-def5-44eb-9065-e80b59530af6
124    DecodeError: From<K::Error>,
125    // Overly verbose because of `?` requiring `From` See also ec4fa3ba-def5-44eb-9065-e80b59530af6
126    DecodeError: From<V::Error>,
127{
128    type Decode = HashMap<K::Decode, V::Decode, S>;
129    type Error = DecodeError;
130
131    fn new(sticks: DynArrayBranch<'_>, options: &impl DecodeOptions) -> DecodeResult<Self> {
132        profile_method!(DecoderArray::new);
133
134        match sticks {
135            DynArrayBranch::Map0 => Ok(None),
136            DynArrayBranch::Map { len, keys, values } => {
137                let (keys, (values, len)) = parallel(
138                    || K::new(*keys, options),
139                    || {
140                        parallel(
141                            || V::new(*values, options),
142                            || <<u64 as Decodable>::DecoderArray as DecoderArray>::new(*len, options),
143                            options,
144                        )
145                    },
146                    options,
147                );
148                let keys = keys?;
149                let values = values?;
150                let len = len?;
151                Ok(Some(HashMapArrayDecoder {
152                    len,
153                    keys,
154                    values,
155                    _marker: Unowned::new(),
156                }))
157            }
158            _ => Err(DecodeError::SchemaMismatch),
159        }
160    }
161    fn decode_next(&mut self) -> Result<Self::Decode, Self::Error> {
162        if let Some(inner) = self {
163            let len = inner.len.decode_next_infallible();
164            let mut result = <Self::Decode as Default>::default(); // TODO: (Performance) capacity
165            for _ in 0..len {
166                let key = inner.keys.decode_next()?;
167                let value = inner.values.decode_next()?;
168                // TODO: decode_next was made infallable for performance reasons,
169                // but duplicate keys would seem a reason to fail. Ideally this could
170                // have a Result<T, !> and perform well in the future.
171                if result.insert(key, value).is_some() {
172                    return Err(DecodeError::InvalidFormat);
173                };
174            }
175            Ok(result)
176        } else {
177            Ok(Default::default())
178        }
179    }
180}