Skip to main content

harmos_domain/message/
maps.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Deterministic Standard Maps With Finite Entry Schemas And Bounded Decoding.
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16/* === Imports === */
17
18// Standard Library Imports
19use std::borrow::Cow;
20use std::collections::{BTreeMap, HashMap};
21use std::hash::{BuildHasher, BuildHasherDefault, Hash, RandomState};
22use std::marker::PhantomData;
23use std::mem::size_of;
24// Third Party Imports
25
26// Local Imports
27use 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
33/* === Constants === */
34
35/* === Macros === */
36
37macro_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
62/* === Type Aliases === */
63
64/* === Structs === */
65
66// Schema-Only Entries Have The Same Bytes As Consecutive Key And Value Payloads.
67pub struct Entry<K, V>(PhantomData<(K, V)>);
68struct Shape<K, V>(PhantomData<(K, V)>);
69
70/* === Enums === */
71
72/* === Traits === */
73
74// Protobuf-Compatible Keys Also Have A Stable Total Ordering.
75pub trait Key: Value + MemorySize + Described + Ord + Hash {}
76
77/* === Implementations === */
78
79impl 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
219/* === Public Functions === */
220
221/* === Private Functions === */
222
223fn 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
248// Charge Bucket Padding, The Control Group And A Conservative 7/8 Load Factor.
249fn 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/* === Tests === */