fn_store/
lib.rs

1/*
2 * Created on Thu Jun 29 2023
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7#![doc = include_str!("../README.md")]
8
9pub mod raw;
10
11use parking_lot::RwLock;
12use type_key::TypeKey;
13use std::cell::UnsafeCell;
14
15use crate::raw::RawStore;
16
17#[derive(Debug, Default)]
18/// Single thread only FnStore implementation.
19///
20/// This implementation is zero cost.
21pub struct LocalFnStore<'a>(UnsafeCell<RawStore<'a>>);
22
23impl<'a> LocalFnStore<'a> {
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    pub fn get_ptr<T: 'a + Send>(&self, key_fn: impl FnOnce() -> T) -> *const T {
29        let key = TypeKey::of_val(&key_fn);
30
31        // SAFETY: safe to borrow shared because self is borrowed shared
32        if let Some(ptr) = unsafe { &*self.0.get().cast_const() }.get(&key) {
33            return ptr;
34        }
35
36        // accuire value first before borrowing exclusively
37        let value = key_fn();
38
39        // SAFETY: safe to borrow exclusively since no one can borrow more
40        unsafe { &mut *self.0.get() }.insert(key, value)
41    }
42
43    /// Get or compute value using key
44    pub fn get<T: 'a + Send>(&self, key: impl FnOnce() -> T) -> &T {
45        // SAFETY: pointer is valid and reference cannot outlive more than Self
46        unsafe { &*self.get_ptr(key) }
47    }
48
49    /// Get or compute value using key
50    pub fn get_mut<T: 'a + Send>(&mut self, key: impl FnOnce() -> T) -> &mut T {
51        // SAFETY: pointer is valid and reference cannot outlive more than Self
52        unsafe { &mut *self.get_ptr(key).cast_mut() }
53    }
54
55    /// Reset stored values
56    pub fn reset(&mut self) {
57        self.0.get_mut().reset();
58    }
59}
60
61unsafe impl Send for LocalFnStore<'_> {}
62
63#[derive(Debug, Default)]
64/// Single thread only and non-Send FnStore implementation
65///
66/// This implementation is zero cost.
67pub struct LocalOnlyFnStore<'a>(UnsafeCell<RawStore<'a>>);
68
69impl<'a> LocalOnlyFnStore<'a> {
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    pub fn get_ptr<T: 'a + Send>(&self, key_fn: impl FnOnce() -> T) -> *const T {
75        let key = TypeKey::of_val(&key_fn);
76
77        // SAFETY: safe to borrow shared because self is borrowed shared
78        if let Some(ptr) = unsafe { &*self.0.get().cast_const() }.get(&key) {
79            return ptr;
80        }
81
82        // accuire value first before borrowing exclusively
83        let value = key_fn();
84
85        // SAFETY: safe to borrow exclusively since no one can borrow more
86        unsafe { &mut *self.0.get() }.insert(key, value)
87    }
88
89    /// Get or compute value using key
90    pub fn get<T: 'a + Send>(&self, key: impl FnOnce() -> T) -> &T {
91        // SAFETY: pointer is valid and reference cannot outlive more than Self
92        unsafe { &*self.get_ptr(key) }
93    }
94
95    /// Get or compute value using key
96    pub fn get_mut<T: 'a + Send>(&mut self, key: impl FnOnce() -> T) -> &mut T {
97        // SAFETY: pointer is valid and reference cannot outlive more than Self
98        unsafe { &mut *self.get_ptr(key).cast_mut() }
99    }
100
101    /// Reset stored values
102    pub fn reset(&mut self) {
103        self.0.get_mut().reset();
104    }
105}
106
107#[derive(Debug, Default)]
108/// Thread safe FnStore implementation.
109///
110/// Uses parking_lot's [`RwLock`] to accuire mutable access to Map.
111pub struct AtomicFnStore<'a>(RwLock<RawStore<'a>>);
112
113impl<'a> AtomicFnStore<'a> {
114    pub fn new() -> Self {
115        Self::default()
116    }
117
118    pub fn get_ptr<T: 'a + Send + Sync>(&self, key_fn: impl FnOnce() -> T) -> *const T {
119        let key = TypeKey::of_val(&key_fn);
120
121        if let Some(ptr) = self.0.read().get(&key) {
122            return ptr;
123        }
124
125        let value = key_fn();
126
127        self.0.write().insert(key, value)
128    }
129
130    /// Get or compute value using key
131    pub fn get<T: 'a + Send + Sync>(&self, key_fn: impl FnOnce() -> T) -> &T {
132        // SAFETY: pointer is valid and reference cannot outlive more than Self
133        unsafe { &*self.get_ptr(key_fn) }
134    }
135
136    /// Get or compute value using key
137    pub fn get_mut<T: 'a + Send + Sync, F>(&mut self, key_fn: impl FnOnce() -> T) -> &mut T {
138        // SAFETY: pointer is valid and reference cannot outlive more than Self
139        unsafe { &mut *self.get_ptr(key_fn).cast_mut() }
140    }
141
142    /// Reset stored values
143    pub fn reset(&mut self) {
144        self.0.get_mut().reset();
145    }
146}
147
148unsafe impl Send for AtomicFnStore<'_> {}
149unsafe impl Sync for AtomicFnStore<'_> {}
150
151#[cfg(test)]
152mod tests {
153    use crate::LocalOnlyFnStore;
154
155    use super::{AtomicFnStore, LocalFnStore};
156
157    #[test]
158    fn test_trait() {
159        const fn is_send<T: Send>() {}
160        const fn is_sync<T: Sync>() {}
161
162        is_send::<LocalFnStore>();
163
164        is_send::<AtomicFnStore>();
165        is_sync::<AtomicFnStore>();
166    }
167
168    #[test]
169    fn test_local() {
170        let store = LocalFnStore::new();
171
172        fn one() -> i32 {
173            1
174        }
175
176        let b = store.get(|| store.get(one) + 1);
177        let a = store.get(one);
178
179        assert_eq!(*b, 2);
180        assert_eq!(*a, 1);
181    }
182
183    #[test]
184    fn test_local_only() {
185        let store = LocalOnlyFnStore::new();
186
187        fn one() -> i32 {
188            1
189        }
190
191        let b = store.get(|| store.get(one) + 1);
192        let a = store.get(one);
193
194        assert_eq!(*b, 2);
195        assert_eq!(*a, 1);
196    }
197
198    #[test]
199    fn test_atomic() {
200        let store = AtomicFnStore::new();
201
202        fn one() -> i32 {
203            1
204        }
205
206        let b = store.get(|| store.get(one) + 1);
207        let a = store.get(one);
208
209        assert_eq!(*b, 2);
210        assert_eq!(*a, 1);
211    }
212}
fn_store/lib.rs

fn_store/
lib.rs
