hooks_gen/
lib.rs

1#![forbid(unsafe_code)]
2
3pub trait Key: Copy + PartialEq {
4    type Value;
5    type Owner: Owner<Key = Self>;
6
7    fn owner(self) -> Self::Owner;
8
9    fn map<R>(self, f: impl FnOnce(&Self::Value) -> R) -> R;
10    fn map_mut<R>(self, f: impl FnOnce(&mut Self::Value) -> R) -> R;
11}
12
13pub trait Owner: Clone + PartialEq {
14    type Key: Key;
15    fn key(&self) -> Self::Key;
16
17    fn shared_count(&self) -> usize;
18}
19
20pub trait Store<T> {
21    type Key: Key<Value = T>;
22
23    fn insert(&self, value: T) -> <Self::Key as Key>::Owner;
24}
25
26pub mod local {
27    use std::marker::PhantomData;
28
29    /// Key is !Send + !Sync.
30    /// ```compile_fail
31    /// # fn test(v: hooks_gen::local::Key<()>) -> impl Send { v };
32    /// ```
33    /// ```compile_fail
34    /// # fn test(v: hooks_gen::local::Key<()>) -> impl Sync { v };
35    /// ```
36    pub struct Key<T: 'static> {
37        index: store::UntypedKey,
38        _phantom: PhantomData<*const T>,
39    }
40
41    impl<T: std::fmt::Debug> std::fmt::Debug for Key<T> {
42        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43            let mut d = f.debug_struct("Key");
44            let d = d.field("index", &self.index);
45            match self.try_map(|v| {
46                if let Some(v) = v {
47                    d.field("value", v)
48                } else {
49                    d.field("<value not present>", &())
50                }
51            }) {
52                Ok(d) => d,
53                Err(error) => d.field("store", &error),
54            }
55            .finish()
56        }
57    }
58
59    impl<T> Eq for Key<T> {}
60    impl<T> PartialEq for Key<T> {
61        fn eq(&self, other: &Self) -> bool {
62            self.index == other.index
63        }
64    }
65
66    impl<T> Key<T> {
67        const fn new(index: store::UntypedKey) -> Self {
68            Self {
69                index,
70                _phantom: PhantomData,
71            }
72        }
73
74        fn try_map<R>(self, f: impl FnOnce(Option<&T>) -> R) -> Result<R, std::cell::BorrowError> {
75            <Store as StoreImpl>::try_map_at(self, f)
76        }
77    }
78
79    impl<T> Copy for Key<T> {}
80    impl<T> Clone for Key<T> {
81        fn clone(&self) -> Self {
82            *self
83        }
84    }
85
86    impl<T: 'static> super::Key for Key<T> {
87        type Value = T;
88
89        type Owner = Owner<T>;
90
91        fn owner(self) -> Self::Owner {
92            Store::register_owner_at(self)
93        }
94
95        fn map<R>(self, f: impl FnOnce(&Self::Value) -> R) -> R {
96            Store::map_at(self, f)
97        }
98
99        /// Note that this locks the whole store.
100        fn map_mut<R>(self, f: impl FnOnce(&mut Self::Value) -> R) -> R {
101            Store::map_mut_at(self, f)
102        }
103    }
104
105    pub struct Owner<T: 'static> {
106        key: Key<T>,
107    }
108
109    impl<T: std::fmt::Debug> std::fmt::Debug for Owner<T> {
110        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111            f.debug_tuple("Owner").field(&self.key).finish()
112        }
113    }
114
115    impl<T: 'static> PartialEq for Owner<T> {
116        fn eq(&self, other: &Self) -> bool {
117            self.key == other.key
118        }
119    }
120
121    pub struct Store;
122
123    trait StoreImpl {
124        fn insert<T: 'static>(value: T) -> Owner<T>;
125        fn unregister_owner<T>(owner: &mut Owner<T>);
126        fn register_owner_at<T>(key: Key<T>) -> Owner<T>;
127
128        fn owners_count_at<T: 'static>(key: Key<T>) -> usize;
129
130        fn map_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&T) -> R) -> R;
131        fn try_map_at<T: 'static, R>(
132            key: Key<T>,
133            f: impl FnOnce(Option<&T>) -> R,
134        ) -> Result<R, std::cell::BorrowError>;
135
136        fn map_mut_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&mut T) -> R) -> R;
137    }
138
139    mod store {
140        use std::{
141            any::{Any, TypeId},
142            cell::{Cell, RefCell},
143            collections::HashMap,
144        };
145
146        use super::{Key, Owner};
147
148        pub(super) type UntypedKey = generational_slab::Index;
149
150        mod generational_slab {
151            #[derive(Debug, Clone, Copy, PartialEq, Eq)]
152            pub(crate) struct Index {
153                index: usize,
154                generation: u64,
155            }
156
157            pub struct ValueWithGeneration<T> {
158                value: T,
159                generation: u64,
160            }
161
162            pub(super) struct GenerationalSlab<T> {
163                slab: slab::Slab<ValueWithGeneration<T>>,
164                generation: u64,
165            }
166
167            impl<T> GenerationalSlab<T> {
168                pub(super) const fn new() -> Self {
169                    Self {
170                        slab: slab::Slab::new(),
171                        generation: 0,
172                    }
173                }
174
175                pub(super) fn insert(&mut self, value: T) -> Index {
176                    let generation = self.generation;
177                    let index = self.slab.insert(ValueWithGeneration { value, generation });
178                    self.generation = self.generation.wrapping_add(1);
179                    Index { index, generation }
180                }
181
182                pub(super) fn get(&self, index: Index) -> Option<&T> {
183                    match self.slab.get(index.index) {
184                        Some(vg) if vg.generation == index.generation => Some(&vg.value),
185                        _ => None,
186                    }
187                }
188
189                pub(super) fn get_mut(&mut self, index: Index) -> Option<&mut T> {
190                    match self.slab.get_mut(index.index) {
191                        Some(vg) if vg.generation == index.generation => Some(&mut vg.value),
192                        _ => None,
193                    }
194                }
195
196                pub(super) fn remove(&mut self, index: Index) -> T {
197                    let vg = self.slab.remove(index.index);
198                    assert_eq!(vg.generation, index.generation, "wrong generation");
199                    vg.value
200                }
201            }
202        }
203
204        struct Item<T> {
205            value: T,
206            owners_count: Cell<usize>,
207        }
208
209        type TypedStore<T> = generational_slab::GenerationalSlab<Item<T>>;
210
211        thread_local!(
212            static STORES: RefCell<
213                HashMap<
214                    TypeId,
215                    // Box<TypedStore<T>>
216                    Box<dyn Any>,
217                >,
218            > = Default::default();
219        );
220
221        fn store_mut<T: 'static, R>(f: impl FnOnce(&mut TypedStore<T>) -> R) -> R {
222            STORES.with_borrow_mut(|stores| {
223                let store = stores
224                    .entry(TypeId::of::<T>())
225                    .or_insert_with(|| Box::new(TypedStore::<T>::new()))
226                    .downcast_mut::<TypedStore<T>>()
227                    .unwrap();
228                f(store)
229            })
230        }
231
232        fn store<T: 'static, R>(f: impl FnOnce(&TypedStore<T>) -> R) -> R {
233            STORES.with_borrow(|stores| {
234                let store = stores
235                    .get(&TypeId::of::<T>())
236                    .unwrap()
237                    .downcast_ref::<TypedStore<T>>()
238                    .unwrap();
239                f(store)
240            })
241        }
242
243        fn try_store<T: 'static, R>(
244            f: impl FnOnce(Option<&TypedStore<T>>) -> R,
245        ) -> Result<R, std::cell::BorrowError> {
246            STORES.with(|stores| {
247                let stores = stores.try_borrow()?;
248                let store = stores
249                    .get(&TypeId::of::<T>())
250                    .map(|store| store.downcast_ref::<TypedStore<T>>().unwrap());
251                Ok(f(store))
252            })
253        }
254
255        impl super::StoreImpl for super::Store {
256            fn insert<T: 'static>(value: T) -> Owner<T> {
257                store_mut::<T, Owner<T>>(|s| {
258                    let index = s.insert(Item {
259                        value,
260                        owners_count: Cell::new(1),
261                    });
262                    let key = Key::new(index);
263                    Owner { key }
264                })
265            }
266
267            fn unregister_owner<T>(owner: &mut Owner<T>) {
268                let index = owner.key.index;
269
270                let should_remove = store::<T, bool>(|items| {
271                    let item = items.get(index).unwrap();
272
273                    let new_count = item.owners_count.get() - 1;
274
275                    if new_count == 0 {
276                        true
277                    } else {
278                        item.owners_count.set(new_count);
279                        false
280                    }
281                });
282
283                if should_remove {
284                    store_mut::<T, ()>(|items| {
285                        let _: T = items.remove(index).value;
286                    })
287                }
288            }
289
290            fn register_owner_at<T>(key: Key<T>) -> Owner<T> {
291                store::<T, ()>(|items| {
292                    let item = items.get(key.index).unwrap();
293
294                    item.owners_count.set(item.owners_count.get() + 1);
295                });
296
297                Owner { key }
298            }
299
300            fn owners_count_at<T: 'static>(key: Key<T>) -> usize {
301                store::<T, usize>(|items| items.get(key.index).unwrap().owners_count.get())
302            }
303
304            fn map_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&T) -> R) -> R {
305                store::<T, R>(|items| f(&items.get(key.index).unwrap().value))
306            }
307
308            fn try_map_at<T: 'static, R>(
309                key: Key<T>,
310                f: impl FnOnce(Option<&T>) -> R,
311            ) -> Result<R, std::cell::BorrowError> {
312                try_store::<T, R>(|items| {
313                    let item = if let Some(items) = items {
314                        items.get(key.index).map(|item: &Item<T>| &item.value)
315                    } else {
316                        None
317                    };
318                    f(item)
319                })
320            }
321
322            fn map_mut_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&mut T) -> R) -> R {
323                store_mut::<T, R>(|items| f(&mut items.get_mut(key.index).unwrap().value))
324            }
325        }
326    }
327
328    impl<T: 'static> super::Store<T> for Store {
329        type Key = Key<T>;
330
331        fn insert(&self, value: T) -> Owner<T> {
332            <Store as StoreImpl>::insert::<T>(value)
333        }
334    }
335
336    impl<T: 'static> super::Owner for Owner<T> {
337        type Key = Key<T>;
338
339        fn key(&self) -> Self::Key {
340            self.key
341        }
342
343        fn shared_count(&self) -> usize {
344            Store::owners_count_at(self.key)
345        }
346    }
347
348    impl<T: 'static> Clone for Owner<T> {
349        fn clone(&self) -> Self {
350            use super::Key;
351            self.key.owner()
352        }
353    }
354
355    impl<T: 'static> Drop for Owner<T> {
356        fn drop(&mut self) {
357            Store::unregister_owner::<T>(self)
358        }
359    }
360}
hooks_gen/lib.rs

hooks_gen/
lib.rs
