1 use alloc::boxed::Box; 2 use core::fmt; 3 use core::marker::PhantomData; 4 use core::mem::{self, MaybeUninit}; 5 use core::ptr; 6 7 /// Number of words a piece of `Data` can hold. 8 /// 9 /// Three words should be enough for the majority of cases. For example, you can fit inside it the 10 /// function pointer together with a fat pointer representing an object that needs to be destroyed. 11 const DATA_WORDS: usize = 3; 12 13 /// Some space to keep a `FnOnce()` object on the stack. 14 type Data = [usize; DATA_WORDS]; 15 16 /// A `FnOnce()` that is stored inline if small, or otherwise boxed on the heap. 17 /// 18 /// This is a handy way of keeping an unsized `FnOnce()` within a sized structure. 19 pub(crate) struct Deferred { 20 call: unsafe fn(*mut u8), 21 data: Data, 22 _marker: PhantomData<*mut ()>, // !Send + !Sync 23 } 24 25 impl fmt::Debug for Deferred { fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error>26 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { 27 f.pad("Deferred { .. }") 28 } 29 } 30 31 impl Deferred { 32 /// Constructs a new `Deferred` from a `FnOnce()`. new<F: FnOnce()>(f: F) -> Self33 pub(crate) fn new<F: FnOnce()>(f: F) -> Self { 34 let size = mem::size_of::<F>(); 35 let align = mem::align_of::<F>(); 36 37 unsafe { 38 if size <= mem::size_of::<Data>() && align <= mem::align_of::<Data>() { 39 let mut data = MaybeUninit::<Data>::uninit(); 40 ptr::write(data.as_mut_ptr() as *mut F, f); 41 42 unsafe fn call<F: FnOnce()>(raw: *mut u8) { 43 let f: F = ptr::read(raw as *mut F); 44 f(); 45 } 46 47 Deferred { 48 call: call::<F>, 49 data: data.assume_init(), 50 _marker: PhantomData, 51 } 52 } else { 53 let b: Box<F> = Box::new(f); 54 let mut data = MaybeUninit::<Data>::uninit(); 55 ptr::write(data.as_mut_ptr() as *mut Box<F>, b); 56 57 unsafe fn call<F: FnOnce()>(raw: *mut u8) { 58 // It's safe to cast `raw` from `*mut u8` to `*mut Box<F>`, because `raw` is 59 // originally derived from `*mut Box<F>`. 60 #[allow(clippy::cast_ptr_alignment)] 61 let b: Box<F> = ptr::read(raw as *mut Box<F>); 62 (*b)(); 63 } 64 65 Deferred { 66 call: call::<F>, 67 data: data.assume_init(), 68 _marker: PhantomData, 69 } 70 } 71 } 72 } 73 74 /// Calls the function. 75 #[inline] call(mut self)76 pub(crate) fn call(mut self) { 77 let call = self.call; 78 unsafe { call(&mut self.data as *mut Data as *mut u8) }; 79 } 80 } 81 82 #[cfg(all(test, not(crossbeam_loom)))] 83 mod tests { 84 use super::Deferred; 85 use std::cell::Cell; 86 87 #[test] on_stack()88 fn on_stack() { 89 let fired = &Cell::new(false); 90 let a = [0usize; 1]; 91 92 let d = Deferred::new(move || { 93 drop(a); 94 fired.set(true); 95 }); 96 97 assert!(!fired.get()); 98 d.call(); 99 assert!(fired.get()); 100 } 101 102 #[test] on_heap()103 fn on_heap() { 104 let fired = &Cell::new(false); 105 let a = [0usize; 10]; 106 107 let d = Deferred::new(move || { 108 drop(a); 109 fired.set(true); 110 }); 111 112 assert!(!fired.get()); 113 d.call(); 114 assert!(fired.get()); 115 } 116 117 #[test] string()118 fn string() { 119 let a = "hello".to_string(); 120 let d = Deferred::new(move || assert_eq!(a, "hello")); 121 d.call(); 122 } 123 124 #[test] boxed_slice_i32()125 fn boxed_slice_i32() { 126 let a: Box<[i32]> = vec![2, 3, 5, 7].into_boxed_slice(); 127 let d = Deferred::new(move || assert_eq!(*a, [2, 3, 5, 7])); 128 d.call(); 129 } 130 131 #[test] long_slice_usize()132 fn long_slice_usize() { 133 let a: [usize; 5] = [2, 3, 5, 7, 11]; 134 let d = Deferred::new(move || assert_eq!(a, [2, 3, 5, 7, 11])); 135 d.call(); 136 } 137 } 138