turbo_tasks/
once_map.rs

1use std::{
2    hash::Hash,
3    sync::{Arc, Mutex},
4};
5
6use dashmap::mapref::entry::Entry;
7
8use crate::FxDashMap;
9
10pub struct OnceConcurrentlyMap<K: Hash + Eq + Ord + Send + Sync + 'static, V: Clone + Send + Sync> {
11    inner: FxDashMap<&'static K, Arc<Mutex<Option<V>>>>,
12}
13
14impl<K: Hash + Eq + Ord + Send + Sync, V: Clone + Send + Sync> Default
15    for OnceConcurrentlyMap<K, V>
16{
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl<K: Hash + Eq + Ord + Send + Sync, V: Clone + Send + Sync> OnceConcurrentlyMap<K, V> {
23    pub fn new() -> Self {
24        Self {
25            inner: FxDashMap::default(),
26        }
27    }
28
29    pub fn action(&self, key: &K, func: impl FnOnce() -> V) -> V {
30        let temp = TemporarilyInserted {
31            inner: &self.inner,
32            key,
33        };
34        let mutex = match temp.entry() {
35            Entry::Occupied(e) => e.get().clone(),
36            Entry::Vacant(e) => e.insert(Arc::new(Mutex::new(None))).clone(),
37        };
38        let mut guard = mutex.lock().unwrap();
39        if let Some(value) = &*guard {
40            // Yeah, somebody else already did it for us
41            return value.clone();
42        }
43        // We are the one responsible for computing
44        let value = func();
45        *guard = Some(value.clone());
46        drop(guard);
47        drop(temp);
48        value
49    }
50}
51
52struct TemporarilyInserted<'a, K: 'static + Hash + Eq + Ord + Send + Sync, V: Send + Sync> {
53    inner: &'a FxDashMap<&'static K, V>,
54    key: &'a K,
55}
56
57impl<'a, K: Hash + Eq + Ord + Send + Sync, V: Send + Sync> TemporarilyInserted<'a, K, V> {
58    fn entry(&self) -> Entry<'a, &'static K, V> {
59        // SAFETY: We remove the value again after this function is done
60        let static_key: &'static K = unsafe { std::mem::transmute(self.key) };
61        self.inner.entry(static_key)
62    }
63}
64
65impl<K: Hash + Eq + Ord + Send + Sync, V: Send + Sync> Drop for TemporarilyInserted<'_, K, V> {
66    fn drop(&mut self) {
67        let static_key: &'static K = unsafe { std::mem::transmute(self.key) };
68        self.inner.remove(&static_key);
69    }
70}
71
72pub struct SafeOnceConcurrentlyMap<
73    K: Clone + Hash + Eq + Ord + Send + Sync + 'static,
74    V: Clone + Send + Sync,
75> {
76    inner: FxDashMap<K, Arc<Mutex<Option<V>>>>,
77}
78
79impl<K: Clone + Hash + Eq + Ord + Send + Sync, V: Clone + Send + Sync> Default
80    for SafeOnceConcurrentlyMap<K, V>
81{
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl<K: Clone + Hash + Eq + Ord + Send + Sync, V: Clone + Send + Sync>
88    SafeOnceConcurrentlyMap<K, V>
89{
90    pub fn new() -> Self {
91        Self {
92            inner: FxDashMap::default(),
93        }
94    }
95
96    pub fn action(&self, key: &K, func: impl FnOnce() -> V) -> V {
97        let temp = SafeTemporarilyInserted {
98            inner: &self.inner,
99            key,
100        };
101        let mutex = match temp.entry() {
102            Entry::Occupied(e) => e.get().clone(),
103            Entry::Vacant(e) => e.insert(Arc::new(Mutex::new(None))).clone(),
104        };
105        let mut guard = mutex.lock().unwrap();
106        if let Some(value) = &*guard {
107            // Yeah, somebody else already did it for us
108            return value.clone();
109        }
110        // We are the one responsible for computing
111        let value = func();
112        *guard = Some(value.clone());
113        drop(guard);
114        drop(temp);
115        value
116    }
117}
118
119struct SafeTemporarilyInserted<
120    'a,
121    K: 'static + Clone + Hash + Eq + Ord + Send + Sync,
122    V: Send + Sync,
123> {
124    inner: &'a FxDashMap<K, V>,
125    key: &'a K,
126}
127
128impl<K: Clone + Hash + Eq + Ord + Send + Sync, V: Send + Sync> SafeTemporarilyInserted<'_, K, V> {
129    fn entry(&self) -> Entry<'_, K, V> {
130        // SAFETY: We remove the value again after this function is done
131        self.inner.entry(self.key.clone())
132    }
133}
134
135impl<K: Clone + Hash + Eq + Ord + Send + Sync, V: Send + Sync> Drop
136    for SafeTemporarilyInserted<'_, K, V>
137{
138    fn drop(&mut self) {
139        self.inner.remove(self.key);
140    }
141}
turbo_tasks/once_map.rs

turbo_tasks/
once_map.rs
