library/std/src/io/lazy.rs

   1 use crate::cell::Cell;
   2 use crate::ptr;
   3 use crate::sync::Arc;
   4 use crate::sys_common;
   5 use crate::sys_common::mutex::Mutex;
   6
   7 pub struct Lazy<T> {
   8     // We never call `lock.init()`, so it is UB to attempt to acquire this mutex reentrantly!
   9     lock: Mutex,
  10     ptr: Cell<*mut Arc<T>>,
  11 }
  12
  13 #[inline]
  14 const fn done<T>() -> *mut Arc<T> {
  15     1_usize as *mut _
  16 }
  17
  18 unsafe impl<T> Sync for Lazy<T> {}
  19
  20 impl<T> Lazy<T> {
  21     pub const fn new() -> Lazy<T> {
  22         Lazy { lock: Mutex::new(), ptr: Cell::new(ptr::null_mut()) }
  23     }
  24 }
  25
  26 impl<T: Send + Sync + 'static> Lazy<T> {
  27     /// Safety: `init` must not call `get` on the variable that is being
  28     /// initialized.
  29     pub unsafe fn get(&'static self, init: fn() -> Arc<T>) -> Option<Arc<T>> {
  30         let _guard = self.lock.lock();
  31         let ptr = self.ptr.get();
  32         if ptr.is_null() {
  33             Some(self.init(init))
  34         } else if ptr == done() {
  35             None
  36         } else {
  37             Some((*ptr).clone())
  38         }
  39     }
  40
  41     // Must only be called with `lock` held
  42     unsafe fn init(&'static self, init: fn() -> Arc<T>) -> Arc<T> {
  43         // If we successfully register an at exit handler, then we cache the
  44         // `Arc` allocation in our own internal box (it will get deallocated by
  45         // the at exit handler). Otherwise we just return the freshly allocated
  46         // `Arc`.
  47         let registered = sys_common::at_exit(move || {
  48             let ptr = {
  49                 let _guard = self.lock.lock();
  50                 self.ptr.replace(done())
  51             };
  52             drop(Box::from_raw(ptr))
  53         });
  54         // This could reentrantly call `init` again, which is a problem
  55         // because our `lock` allows reentrancy!
  56         // That's why `get` is unsafe and requires the caller to ensure no reentrancy happens.
  57         let ret = init();
  58         if registered.is_ok() {
  59             self.ptr.set(Box::into_raw(Box::new(ret.clone())));
  60         }
  61         ret
  62     }
  63 }