1 | pub(crate) use imp::*; |
2 | |
3 | #[cfg (not(feature = "std" ))] |
4 | mod imp { |
5 | use alloc::sync::Arc; |
6 | use core::sync::atomic::{AtomicPtr, Ordering}; |
7 | use core::{mem, ptr}; |
8 | |
9 | #[derive(Debug, Default)] |
10 | pub(crate) struct LazyArc<T> { |
11 | // Only written once with a value obtained from `Arc<T>::into_raw`. |
12 | // This holds a ref count for the `Arc`, so it is always safe to |
13 | // clone the `Arc` given a reference to the `LazyArc`. |
14 | value: AtomicPtr<T>, |
15 | } |
16 | |
17 | impl<T> Drop for LazyArc<T> { |
18 | fn drop(&mut self) { |
19 | let value_ptr = self.value.load(Ordering::Acquire); |
20 | if !value_ptr.is_null() { |
21 | // SAFETY: all writes to `self.value` are pointers obtained from `Arc::into_raw`. |
22 | drop(unsafe { Arc::from_raw(value_ptr) }); |
23 | } |
24 | } |
25 | } |
26 | |
27 | impl<T> LazyArc<T> { |
28 | pub(crate) fn get<E, F: FnOnce() -> Result<T, E>>(&self, f: F) -> Result<Arc<T>, E> { |
29 | // Clone an `Arc` given a pointer obtained from `Arc::into_raw`. |
30 | // SAFETY: `value_ptr` must be a valid pointer obtained from `Arc<T>::into_raw`. |
31 | unsafe fn clone_arc_ptr<T>(value_ptr: *const T) -> Arc<T> { |
32 | let value = Arc::from_raw(value_ptr); |
33 | let clone = Arc::clone(&value); |
34 | mem::forget(value); |
35 | clone |
36 | } |
37 | |
38 | // Return the existing value if already computed. |
39 | // `Ordering::Acquire` is needed so that the content of the loaded `Arc` is |
40 | // visible to this thread. |
41 | let value_ptr = self.value.load(Ordering::Acquire); |
42 | if !value_ptr.is_null() { |
43 | // SAFETY: all writes to `self.value` are pointers obtained from `Arc::into_raw`. |
44 | return Ok(unsafe { clone_arc_ptr(value_ptr) }); |
45 | } |
46 | |
47 | // Race to compute and set the value. |
48 | let value = f().map(Arc::new)?; |
49 | let value_ptr = Arc::into_raw(value); |
50 | match self.value.compare_exchange( |
51 | ptr::null_mut(), |
52 | value_ptr as *mut T, |
53 | // Success: `Ordering::Release` is needed so that the content of the stored `Arc` |
54 | // is visible to other threads. No ordering is required for the null ptr that is |
55 | // loaded, but older rust versions (< 1.64) require that its ordering must not |
56 | // be weaker than the failure ordering, so we use `Ordering::AcqRel`. |
57 | Ordering::AcqRel, |
58 | // Failure: `Ordering::Acquire` is needed so that the content of the loaded `Arc` |
59 | // is visible to this thread. |
60 | Ordering::Acquire, |
61 | ) { |
62 | Ok(_) => { |
63 | // Return the value we computed. |
64 | // SAFETY: `value_ptr` was obtained from `Arc::into_raw`. |
65 | Ok(unsafe { clone_arc_ptr(value_ptr) }) |
66 | } |
67 | Err(existing_value_ptr) => { |
68 | // We lost the race, drop unneeded `value_ptr`. |
69 | // SAFETY: `value_ptr` was obtained from `Arc::into_raw`. |
70 | drop(unsafe { Arc::from_raw(value_ptr) }); |
71 | // Return the existing value. |
72 | // SAFETY: all writes to `self.value` are pointers obtained from `Arc::into_raw`. |
73 | Ok(unsafe { clone_arc_ptr(existing_value_ptr) }) |
74 | } |
75 | } |
76 | } |
77 | } |
78 | } |
79 | |
80 | #[cfg (feature = "std" )] |
81 | mod imp { |
82 | use std::sync::{Arc, Mutex}; |
83 | |
84 | #[derive(Debug, Default)] |
85 | pub(crate) struct LazyArc<T> { |
86 | value: Mutex<Option<Arc<T>>>, |
87 | } |
88 | |
89 | impl<T> LazyArc<T> { |
90 | pub(crate) fn get<E, F: FnOnce() -> Result<T, E>>(&self, f: F) -> Result<Arc<T>, E> { |
91 | let mut lock: MutexGuard<'_, {unknown}> = self.value.lock().unwrap(); |
92 | if let Some(value: &{unknown}) = &*lock { |
93 | return Ok(value.clone()); |
94 | } |
95 | let value = f().map(Arc::new)?; |
96 | *lock = Some(value.clone()); |
97 | Ok(value) |
98 | } |
99 | } |
100 | } |
101 | |
102 | #[cfg (test)] |
103 | mod tests { |
104 | use super::*; |
105 | |
106 | #[test] |
107 | fn lazy_arc() { |
108 | let lazy = LazyArc::default(); |
109 | let value = lazy.get(|| Err(())); |
110 | assert_eq!(value, Err(())); |
111 | let value = lazy.get(|| Ok::<i32, ()>(3)).unwrap(); |
112 | assert_eq!(*value, 3); |
113 | let value = lazy.get(|| Err(())).unwrap(); |
114 | assert_eq!(*value, 3); |
115 | } |
116 | } |
117 | |