1use std::borrow::Cow;
20use std::collections::{BTreeMap, HashMap};
21use std::hash::{BuildHasher, BuildHasherDefault, Hash, RandomState};
22use std::marker::PhantomData;
23use std::mem::size_of;
24use super::postcard::{Reader, Value};
28use super::{MemorySize, MessageError, btree_map_memory, decoding_add};
29use crate::schema::declaration::{Definition, Name};
30use crate::schema::graph::{Field, Kind};
31use crate::schema::{Cardinality, Described};
32
33macro_rules! keys {
38 ($($ty:ty),* $(,)?) => { $(impl Key for $ty {})* };
39}
40
41keys!(bool, u32, u64, i32, i64, String);
42
43macro_rules! map_schema {
44 ($container:ident $(, $hasher:ident)?) => {
45 impl<K, V $(, $hasher)?> Described for $container<K, V $(, $hasher)?>
46 where
47 K: Key,
48 V: Described,
49 $($hasher: 'static,)?
50 {
51 const NAME: &'static str = Shape::<K, V>::NAME;
52 const FIELDS: &'static [Field<Definition>] = Shape::<K, V>::FIELDS;
53 const MARKER: &'static [u8] = Shape::<K, V>::MARKER;
54 const DEFINITION: Definition = Shape::<K, V>::DEFINITION;
55 }
56 };
57}
58
59map_schema!(BTreeMap);
60map_schema!(HashMap, S);
61
62pub struct Entry<K, V>(PhantomData<(K, V)>);
68struct Shape<K, V>(PhantomData<(K, V)>);
69
70pub trait Key: Value + MemorySize + Described + Ord + Hash {}
76
77impl MemorySize for RandomState {
80 const INLINE_ONLY: bool = true;
81 fn memory_size(&self) -> usize {
82 size_of::<Self>()
83 }
84}
85
86impl<T> MemorySize for BuildHasherDefault<T> {
87 const INLINE_ONLY: bool = true;
88 fn memory_size(&self) -> usize {
89 size_of::<Self>()
90 }
91}
92
93impl<K: Key, V: Described> Described for Entry<K, V> {
94 const NAME: &'static str = Name::new("harmos::MapEntry", &[K::NAME, V::NAME]).as_str();
95 const FIELDS: &'static [Field<Definition>] = &[
96 Field {
97 name: Cow::Borrowed("key"),
98 tag: 1,
99 kind: K::FIELD_KIND,
100 cardinality: Cardinality::Required,
101 },
102 Field {
103 name: Cow::Borrowed("value"),
104 tag: 2,
105 kind: V::FIELD_KIND,
106 cardinality: Cardinality::Required,
107 },
108 ];
109 const MARKER: &'static [u8] = crate::schema::declaration::encode_buffer::<2048>(Self::NAME, Self::FIELDS)
110 .split_at(crate::schema::declaration::encoded_len(Self::NAME, Self::FIELDS))
111 .0;
112}
113
114impl<K: Key, V: Described> Described for Shape<K, V> {
115 const NAME: &'static str = Name::new("harmos::Map", &[K::NAME, V::NAME]).as_str();
116 const FIELDS: &'static [Field<Definition>] = &[Field {
117 name: Cow::Borrowed("entries"),
118 tag: 1,
119 kind: Kind::Message { definition: Entry::<K, V>::DEFINITION },
120 cardinality: Cardinality::Repeated,
121 }];
122 const MARKER: &'static [u8] = crate::schema::declaration::encode_buffer::<2048>(Self::NAME, Self::FIELDS)
123 .split_at(crate::schema::declaration::encoded_len(Self::NAME, Self::FIELDS))
124 .0;
125}
126
127impl<K: Key, V: Value> Value for BTreeMap<K, V> {
128 fn write(&self, bytes: &mut Vec<u8>) {
129 self.len().write(bytes);
130 for (key, value) in self {
131 key.write(bytes);
132 value.write(bytes);
133 }
134 }
135
136 fn size(&self) -> usize {
137 self.iter().fold(self.len().size(), |size, (key, value)| size.saturating_add(key.size()).saturating_add(value.size()))
138 }
139
140 fn memory(input: &mut Reader<'_>) -> Result<usize, MessageError> {
141 memory::<K, V>(input, |count| btree_map_memory::<K, V>(count))
142 }
143
144 fn read(input: &mut Reader<'_>) -> Result<Self, MessageError> {
145 input.nested(|input| {
146 let count = count(input)?;
147 let mut values = Self::new();
148 for _ in 0..count {
149 let (key, value) = input.nested(|input| Ok((K::read(input)?, V::read(input)?)))?;
150 if values.insert(key, value).is_some() {
151 return Err(MessageError::borrowed("duplicate map key"));
152 }
153 }
154 Ok(values)
155 })
156 }
157}
158
159impl<K: Key, V: Value, S: BuildHasher + Default + MemorySize + Send + Sync> Value for HashMap<K, V, S> {
160 fn write(&self, bytes: &mut Vec<u8>) {
161 self.len().write(bytes);
162 let mut entries: Vec<_> = self.iter().collect();
163 entries.sort_unstable_by_key(|(key, _)| *key);
164 for (key, value) in entries {
165 key.write(bytes);
166 value.write(bytes);
167 }
168 }
169
170 fn size(&self) -> usize {
171 self.iter().fold(self.len().size(), |size, (key, value)| size.saturating_add(key.size()).saturating_add(value.size()))
172 }
173
174 fn memory(input: &mut Reader<'_>) -> Result<usize, MessageError> {
175 memory::<K, V>(input, |count| {
176 size_of::<Self>()
177 .saturating_add(table_memory::<K, V>(count))
178 .saturating_add(S::default().memory_size().saturating_sub(size_of::<S>()))
179 })
180 }
181
182 fn read(input: &mut Reader<'_>) -> Result<Self, MessageError> {
183 input.nested(|input| {
184 let count = count(input)?;
185 let mut values = Self::with_hasher(S::default());
186 values.try_reserve(count).map_err(|_| MessageError::borrowed("map allocation exceeds its addressable capacity"))?;
187 for _ in 0..count {
188 let (key, value) = input.nested(|input| Ok((K::read(input)?, V::read(input)?)))?;
189 if values.insert(key, value).is_some() {
190 return Err(MessageError::borrowed("duplicate map key"));
191 }
192 }
193 Ok(values)
194 })
195 }
196}
197
198impl<K: MemorySize, V: MemorySize, S: MemorySize> MemorySize for HashMap<K, V, S> {
199 fn memory_size(&self) -> usize {
200 let hasher = self.hasher().memory_size();
201 if hasher == usize::MAX {
202 return usize::MAX;
203 }
204 let mut bytes =
205 size_of::<Self>().saturating_add(table_memory::<K, V>(self.capacity())).saturating_add(hasher.saturating_sub(size_of::<S>()));
206 for (key, value) in self {
207 let key_bytes = key.memory_size();
208 let value_bytes = value.memory_size();
209 if key_bytes == usize::MAX || value_bytes == usize::MAX {
210 return usize::MAX;
211 }
212 bytes =
213 bytes.saturating_add(key_bytes.saturating_sub(size_of::<K>())).saturating_add(value_bytes.saturating_sub(size_of::<V>()));
214 }
215 bytes
216 }
217}
218
219fn count(input: &mut Reader<'_>) -> Result<usize, MessageError> {
224 let count: usize = input.value()?;
225 input.claim(count)?;
226 Ok(count)
227}
228
229fn memory<K: Value, V: Value>(input: &mut Reader<'_>, allocation: impl FnOnce(usize) -> usize) -> Result<usize, MessageError> {
230 input.nested(|input| {
231 let count = count(input)?;
232 let mut bytes = allocation(count);
233 if bytes == usize::MAX {
234 return Err(MessageError::borrowed("map allocation exceeds its addressable capacity"));
235 }
236 for _ in 0..count {
237 let owned = input.nested(|input| {
238 let key = K::memory(input)?.saturating_sub(size_of::<K>());
239 let value = V::memory(input)?.saturating_sub(size_of::<V>());
240 decoding_add(key, value)
241 })?;
242 bytes = decoding_add(bytes, owned)?;
243 }
244 Ok(bytes)
245 })
246}
247
248fn table_memory<K, V>(capacity: usize) -> usize {
250 if capacity == 0 {
251 return 0;
252 }
253 let buckets = capacity
254 .checked_mul(8)
255 .and_then(|value| value.checked_add(6))
256 .map(|value| value / 7)
257 .and_then(|value| value.max(4).checked_next_power_of_two());
258 buckets
259 .and_then(|buckets| buckets.checked_mul(size_of::<(K, V)>().saturating_add(1)))
260 .and_then(|bytes| bytes.checked_add(16 + std::mem::align_of::<(K, V)>()))
261 .unwrap_or(usize::MAX)
262}
263
264