object_rainbow_hamt/
lib.rs

1use std::pin::Pin;
2
3use object_rainbow::{
4    Enum, Fetch, Hash, Inline, InlineOutput, ListHashes, MaybeHasNiche, Output, Parse, ParseInline,
5    PointInput, PointVisitor, Size, Tagged, Tags, ToOutput, Topological, Traversible, assert_impl,
6};
7use object_rainbow_array_map::ArrayMap;
8use object_rainbow_point::{IntoPoint, Point};
9
10type OptionFuture<'a, T> =
11    Pin<Box<dyn 'a + Send + Future<Output = object_rainbow::Result<Option<T>>>>>;
12
13trait Amt<K> {
14    type V: Send + Sync;
15    fn insert(&mut self, key: K, value: Self::V) -> OptionFuture<'_, Self::V>;
16    fn from_pair(a: (K, Self::V), b: (K, Self::V)) -> Self;
17    fn get(&self, key: K) -> OptionFuture<'_, Self::V>;
18}
19
20#[derive(ToOutput, Tagged, ListHashes, Topological, Parse, Clone)]
21/// what are you even doing at this point
22struct DeepestLeaf<V = ()>(ArrayMap<V>);
23
24impl<V: Send + Sync + Clone> Amt<u8> for DeepestLeaf<V> {
25    type V = V;
26
27    fn insert(&mut self, key: u8, value: Self::V) -> OptionFuture<'_, Self::V> {
28        Box::pin(async move { Ok(self.0.insert(key, value)) })
29    }
30
31    fn from_pair(a: (u8, Self::V), b: (u8, Self::V)) -> Self {
32        Self([a, b].into())
33    }
34
35    fn get(&self, key: u8) -> OptionFuture<'_, Self::V> {
36        Box::pin(async move { Ok(self.0.get(key).cloned()) })
37    }
38}
39
40#[derive(Enum, ToOutput, InlineOutput, Tagged, ListHashes, Topological, Parse, ParseInline)]
41enum SubTree<T, K, V = <T as Amt<K>>::V> {
42    Leaf(K, V),
43    SubTree(Point<T>),
44}
45
46impl<T, K: Clone, V: Clone> Clone for SubTree<T, K, V> {
47    fn clone(&self) -> Self {
48        match self {
49            Self::Leaf(key, value) => Self::Leaf(key.clone(), value.clone()),
50            Self::SubTree(point) => Self::SubTree(point.clone()),
51        }
52    }
53}
54
55impl<T: Amt<K, V: Clone> + Clone + Traversible, K: Send + Sync + PartialEq + Clone> Amt<K>
56    for SubTree<T, K>
57{
58    type V = T::V;
59
60    fn insert(&mut self, key: K, value: Self::V) -> OptionFuture<'_, Self::V> {
61        Box::pin(async move {
62            match self {
63                SubTree::Leaf(xkey, xvalue) => Ok(if *xkey != key {
64                    *self = Self::from_pair((xkey.clone(), xvalue.clone()), (key, value));
65                    None
66                } else {
67                    Some(std::mem::replace(xvalue, value))
68                }),
69                SubTree::SubTree(sub) => sub.fetch_mut().await?.insert(key, value).await,
70            }
71        })
72    }
73
74    fn from_pair(a: (K, Self::V), b: (K, Self::V)) -> Self {
75        Self::SubTree(T::from_pair(a, b).point())
76    }
77
78    fn get(&self, key: K) -> OptionFuture<'_, Self::V> {
79        Box::pin(async move {
80            match self {
81                SubTree::Leaf(existing, value) => Ok((*existing == key).then(|| value.clone())),
82                SubTree::SubTree(sub) => sub.fetch().await?.get(key).await,
83            }
84        })
85    }
86}
87
88#[derive(ToOutput, Tagged, ListHashes, Topological, Parse)]
89struct SetNode<T, K, V = <T as Amt<K>>::V>(ArrayMap<SubTree<T, K, V>>);
90
91impl<T, K: Clone, V: Clone> Clone for SetNode<T, K, V> {
92    fn clone(&self) -> Self {
93        Self(self.0.clone())
94    }
95}
96
97impl<T, K, V> Default for SetNode<T, K, V> {
98    fn default() -> Self {
99        Self(Default::default())
100    }
101}
102
103impl<T: Amt<K, V: Clone> + Clone + Traversible, K: Send + Sync + PartialEq + Clone> Amt<(u8, K)>
104    for SetNode<T, K>
105{
106    type V = T::V;
107
108    fn insert(&mut self, (key, rest): (u8, K), value: Self::V) -> OptionFuture<'_, Self::V> {
109        Box::pin(async move {
110            if let Some(sub) = self.0.get_mut(key) {
111                sub.insert(rest, value).await
112            } else {
113                assert!(self.0.insert(key, SubTree::Leaf(rest, value)).is_none());
114                Ok(None)
115            }
116        })
117    }
118
119    fn from_pair(
120        ((a, rest_a), value_a): ((u8, K), Self::V),
121        ((b, rest_b), value_b): ((u8, K), Self::V),
122    ) -> Self {
123        if a == b {
124            Self([(a, SubTree::from_pair((rest_a, value_a), (rest_b, value_b)))].into())
125        } else {
126            Self(
127                [
128                    (a, SubTree::Leaf(rest_a, value_a)),
129                    (b, SubTree::Leaf(rest_b, value_b)),
130                ]
131                .into(),
132            )
133        }
134    }
135
136    fn get(&self, (key, rest): (u8, K)) -> OptionFuture<'_, Self::V> {
137        Box::pin(async move {
138            if let Some(sub) = self.0.get(key) {
139                sub.get(rest).await
140            } else {
141                Ok(None)
142            }
143        })
144    }
145}
146
147type K1 = u8;
148type K2 = (u8, K1);
149type K3 = (u8, K2);
150type K4 = (u8, K3);
151type K5 = (u8, K4);
152type K6 = (u8, K5);
153type K7 = (u8, K6);
154type K8 = (u8, K7);
155type K9 = (u8, K8);
156type K10 = (u8, K9);
157type K11 = (u8, K10);
158type K12 = (u8, K11);
159type K13 = (u8, K12);
160type K14 = (u8, K13);
161type K15 = (u8, K14);
162type K16 = (u8, K15);
163type K17 = (u8, K16);
164type K18 = (u8, K17);
165type K19 = (u8, K18);
166type K20 = (u8, K19);
167type K21 = (u8, K20);
168type K22 = (u8, K21);
169type K23 = (u8, K22);
170type K24 = (u8, K23);
171type K25 = (u8, K24);
172type K26 = (u8, K25);
173type K27 = (u8, K26);
174type K28 = (u8, K27);
175type K29 = (u8, K28);
176type K30 = (u8, K29);
177type K31 = (u8, K30);
178type K32 = (u8, K31);
179
180mod private {
181    use super::*;
182    type N1<V = ()> = DeepestLeaf<V>;
183
184    macro_rules! next_node {
185        ($prev:ident, $next:ident, $pk:ident, $k:ident) => {
186            #[derive(Clone)]
187            pub struct $next<V = ()>(SetNode<$prev<V>, $pk, V>);
188
189            impl<V> Default for $next<V> {
190                fn default() -> Self {
191                    Self(Default::default())
192                }
193            }
194
195            impl<V: ParseInline<I> + Inline<I::Extra>, I: PointInput<Extra: Send + Sync>> Parse<I>
196                for $next<V>
197            {
198                fn parse(input: I) -> object_rainbow::Result<Self> {
199                    Ok(Self(input.parse()?))
200                }
201            }
202
203            impl<V: Tagged> Tagged for $next<V> {
204                const TAGS: Tags = V::TAGS;
205            }
206
207            impl<V: InlineOutput> ToOutput for $next<V> {
208                fn to_output(&self, output: &mut dyn Output) {
209                    self.0.to_output(output);
210                }
211            }
212
213            impl<V: ListHashes> ListHashes for $next<V> {
214                fn list_hashes(&self, f: &mut impl FnMut(Hash)) {
215                    self.0.list_hashes(f)
216                }
217            }
218
219            impl<V: Traversible + InlineOutput> Topological for $next<V> {
220                fn traverse(&self, visitor: &mut impl PointVisitor) {
221                    self.0.traverse(visitor)
222                }
223            }
224
225            impl<V: Traversible + InlineOutput + Clone> Amt<$k> for $next<V> {
226                type V = V;
227
228                fn insert(&mut self, key: $k, value: Self::V) -> OptionFuture<'_, Self::V> {
229                    self.0.insert(key, value)
230                }
231
232                fn from_pair(a: ($k, Self::V), b: ($k, Self::V)) -> Self {
233                    Self(Amt::from_pair(a, b))
234                }
235
236                fn get(&self, key: $k) -> OptionFuture<'_, Self::V> {
237                    self.0.get(key)
238                }
239            }
240        };
241    }
242
243    next_node!(N1, N2, K1, K2);
244    next_node!(N2, N3, K2, K3);
245    next_node!(N3, N4, K3, K4);
246    next_node!(N4, N5, K4, K5);
247    next_node!(N5, N6, K5, K6);
248    next_node!(N6, N7, K6, K7);
249    next_node!(N7, N8, K7, K8);
250    next_node!(N8, N9, K8, K9);
251    next_node!(N9, N10, K9, K10);
252    next_node!(N10, N11, K10, K11);
253    next_node!(N11, N12, K11, K12);
254    next_node!(N12, N13, K12, K13);
255    next_node!(N13, N14, K13, K14);
256    next_node!(N14, N15, K14, K15);
257    next_node!(N15, N16, K15, K16);
258    next_node!(N16, N17, K16, K17);
259    next_node!(N17, N18, K17, K18);
260    next_node!(N18, N19, K18, K19);
261    next_node!(N19, N20, K19, K20);
262    next_node!(N20, N21, K20, K21);
263    next_node!(N21, N22, K21, K22);
264    next_node!(N22, N23, K22, K23);
265    next_node!(N23, N24, K23, K24);
266    next_node!(N24, N25, K24, K25);
267    next_node!(N25, N26, K25, K26);
268    next_node!(N26, N27, K26, K27);
269    next_node!(N27, N28, K27, K28);
270    next_node!(N28, N29, K28, K29);
271    next_node!(N29, N30, K29, K30);
272    next_node!(N30, N31, K30, K31);
273    next_node!(N31, N32, K31, K32);
274}
275
276#[derive(
277    ToOutput, InlineOutput, Tagged, ListHashes, Topological, Parse, ParseInline, Size, MaybeHasNiche,
278)]
279pub struct HamtMap<V>(Point<private::N32<V>>);
280
281assert_impl!(
282    impl<V, E> Inline<E> for HamtMap<V>
283    where
284        E: 'static + Send + Sync + Clone,
285        V: Inline<E>,
286    {
287    }
288);
289
290impl<V> Clone for HamtMap<V> {
291    fn clone(&self) -> Self {
292        Self(self.0.clone())
293    }
294}
295
296impl<V> PartialEq for HamtMap<V> {
297    fn eq(&self, other: &Self) -> bool {
298        self.0 == other.0
299    }
300}
301
302impl<V> Eq for HamtMap<V> {}
303
304impl<V: Traversible + InlineOutput + Clone> Default for HamtMap<V> {
305    fn default() -> Self {
306        Self(Default::default())
307    }
308}
309
310impl<V: Traversible + InlineOutput + Clone> HamtMap<V> {
311    pub fn new() -> Self {
312        Self::default()
313    }
314
315    pub async fn insert(&mut self, hash: Hash, value: V) -> object_rainbow::Result<Option<V>> {
316        self.0
317            .fetch_mut()
318            .await?
319            .insert(hash_key(hash), value)
320            .await
321    }
322
323    pub async fn get(&self, hash: Hash) -> object_rainbow::Result<Option<V>> {
324        self.0.fetch().await?.get(hash_key(hash)).await
325    }
326}
327
328fn hash_key(hash: Hash) -> K32 {
329    let [
330        x0,
331        x1,
332        x2,
333        x3,
334        x4,
335        x5,
336        x6,
337        x7,
338        x8,
339        x9,
340        x10,
341        x11,
342        x12,
343        x13,
344        x14,
345        x15,
346        x16,
347        x17,
348        x18,
349        x19,
350        x20,
351        x21,
352        x22,
353        x23,
354        x24,
355        x25,
356        x26,
357        x27,
358        x28,
359        x29,
360        x30,
361        x31,
362    ] = hash.into_bytes();
363    (x0,(x1,(x2,(x3,(x4,(x5,(x6,(x7,(x8,(x9,(x10,(x11,(x12,(x13,(x14,(x15,(x16,(x17,(x18,(x19,(x20,(x21,(x22,(x23,(x24,(x25,(x26,(x27,(x28,(x29,(x30,(x31))))))))))))))))))))))))))))))))
364}
365
366#[derive(
367    ToOutput,
368    InlineOutput,
369    Tagged,
370    ListHashes,
371    Topological,
372    Parse,
373    ParseInline,
374    Size,
375    MaybeHasNiche,
376    Clone,
377    Default,
378    PartialEq,
379    Eq,
380)]
381pub struct HamtSet(HamtMap<()>);
382
383assert_impl!(
384    impl<E> Inline<E> for HamtSet where E: 'static + Send + Sync + Clone {}
385);
386
387impl HamtSet {
388    pub fn new() -> Self {
389        Self::default()
390    }
391
392    pub async fn insert(&mut self, hash: Hash) -> object_rainbow::Result<bool> {
393        Ok(self.0.insert(hash, ()).await?.is_none())
394    }
395
396    pub async fn contains(&self, hash: Hash) -> object_rainbow::Result<bool> {
397        Ok(self.0.get(hash).await?.is_some())
398    }
399}
object_rainbow_hamt/lib.rs

object_rainbow_hamt/
lib.rs
