cargo/util/
once.rs

1//! Extension functions for [`std::sync::OnceLock`] / [`std::cell::OnceCell`]
2//!
3//! This adds polyfills for functionality in `lazycell` that is not stable within `std`.
4
5pub trait OnceExt {
6    type T;
7
8    /// This might run `f` multiple times if different threads start initializing at once.
9    fn try_borrow_with<F, E>(&self, f: F) -> Result<&Self::T, E>
10    where
11        F: FnOnce() -> Result<Self::T, E>;
12
13    /// This might run `f` multiple times if different threads start initializing at once.
14    fn try_borrow_mut_with<F, E>(&mut self, f: F) -> Result<&mut Self::T, E>
15    where
16        F: FnOnce() -> Result<Self::T, E>;
17
18    fn replace(&mut self, new_value: Self::T) -> Option<Self::T>;
19
20    fn filled(&self) -> bool;
21}
22
23impl<T> OnceExt for std::sync::OnceLock<T> {
24    type T = T;
25
26    fn try_borrow_with<F, E>(&self, f: F) -> Result<&T, E>
27    where
28        F: FnOnce() -> Result<T, E>,
29    {
30        if let Some(value) = self.get() {
31            return Ok(value);
32        }
33
34        // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
35        // no other `f` is executing and the value is not initialized. However, correctly implementing that is
36        // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
37        let value = f()?;
38        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
39        // returns an error. We ignore it and return the value set by the other
40        // thread.
41        let _ = self.set(value);
42        Ok(self.get().unwrap())
43    }
44
45    fn try_borrow_mut_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
46    where
47        F: FnOnce() -> Result<T, E>,
48    {
49        let value = if let Some(value) = self.take() {
50            value
51        } else {
52            // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
53            // no other `f` is executing and the value is not initialized. However, correctly implementing that is
54            // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
55            f()?
56        };
57        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
58        // returns an error. We ignore it and return the value set by the other
59        // thread.
60        let _ = self.set(value);
61        Ok(self.get_mut().unwrap())
62    }
63
64    fn replace(&mut self, new_value: T) -> Option<T> {
65        if let Some(value) = self.get_mut() {
66            Some(std::mem::replace(value, new_value))
67        } else {
68            let result = self.set(new_value);
69            assert!(result.is_ok());
70            None
71        }
72    }
73
74    fn filled(&self) -> bool {
75        self.get().is_some()
76    }
77}
78
79impl<T> OnceExt for std::cell::OnceCell<T> {
80    type T = T;
81
82    fn try_borrow_with<F, E>(&self, f: F) -> Result<&T, E>
83    where
84        F: FnOnce() -> Result<T, E>,
85    {
86        if let Some(value) = self.get() {
87            return Ok(value);
88        }
89
90        // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
91        // no other `f` is executing and the value is not initialized. However, correctly implementing that is
92        // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
93        let value = f()?;
94        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
95        // returns an error. We ignore it and return the value set by the other
96        // thread.
97        let _ = self.set(value);
98        Ok(self.get().unwrap())
99    }
100
101    fn try_borrow_mut_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
102    where
103        F: FnOnce() -> Result<T, E>,
104    {
105        let value = if let Some(value) = self.take() {
106            value
107        } else {
108            // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
109            // no other `f` is executing and the value is not initialized. However, correctly implementing that is
110            // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
111            f()?
112        };
113        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
114        // returns an error. We ignore it and return the value set by the other
115        // thread.
116        let _ = self.set(value);
117        Ok(self.get_mut().unwrap())
118    }
119
120    fn replace(&mut self, new_value: T) -> Option<T> {
121        if let Some(value) = self.get_mut() {
122            Some(std::mem::replace(value, new_value))
123        } else {
124            let result = self.set(new_value);
125            assert!(result.is_ok());
126            None
127        }
128    }
129
130    fn filled(&self) -> bool {
131        self.get().is_some()
132    }
133}