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: MaybeUninit<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     pub(crate) const NO_OP: Self = {
no_op_call(_raw: *mut u8)33         fn no_op_call(_raw: *mut u8) {}
34         Self {
35             call: no_op_call,
36             data: MaybeUninit::uninit(),
37             _marker: PhantomData,
38         }
39     };
40 
41     /// Constructs a new `Deferred` from a `FnOnce()`.
new<F: FnOnce()>(f: F) -> Self42     pub(crate) fn new<F: FnOnce()>(f: F) -> Self {
43         let size = mem::size_of::<F>();
44         let align = mem::align_of::<F>();
45 
46         unsafe {
47             if size <= mem::size_of::<Data>() && align <= mem::align_of::<Data>() {
48                 let mut data = MaybeUninit::<Data>::uninit();
49                 ptr::write(data.as_mut_ptr().cast::<F>(), f);
50 
51                 unsafe fn call<F: FnOnce()>(raw: *mut u8) {
52                     let f: F = ptr::read(raw.cast::<F>());
53                     f();
54                 }
55 
56                 Deferred {
57                     call: call::<F>,
58                     data,
59                     _marker: PhantomData,
60                 }
61             } else {
62                 let b: Box<F> = Box::new(f);
63                 let mut data = MaybeUninit::<Data>::uninit();
64                 ptr::write(data.as_mut_ptr().cast::<Box<F>>(), b);
65 
66                 unsafe fn call<F: FnOnce()>(raw: *mut u8) {
67                     // It's safe to cast `raw` from `*mut u8` to `*mut Box<F>`, because `raw` is
68                     // originally derived from `*mut Box<F>`.
69                     let b: Box<F> = ptr::read(raw.cast::<Box<F>>());
70                     (*b)();
71                 }
72 
73                 Deferred {
74                     call: call::<F>,
75                     data,
76                     _marker: PhantomData,
77                 }
78             }
79         }
80     }
81 
82     /// Calls the function.
83     #[inline]
call(mut self)84     pub(crate) fn call(mut self) {
85         let call = self.call;
86         unsafe { call(self.data.as_mut_ptr().cast::<u8>()) };
87     }
88 }
89 
90 #[cfg(all(test, not(crossbeam_loom)))]
91 mod tests {
92     use super::Deferred;
93     use std::cell::Cell;
94     use std::convert::identity;
95 
96     #[test]
on_stack()97     fn on_stack() {
98         let fired = &Cell::new(false);
99         let a = [0usize; 1];
100 
101         let d = Deferred::new(move || {
102             let _ = identity(a);
103             fired.set(true);
104         });
105 
106         assert!(!fired.get());
107         d.call();
108         assert!(fired.get());
109     }
110 
111     #[test]
on_heap()112     fn on_heap() {
113         let fired = &Cell::new(false);
114         let a = [0usize; 10];
115 
116         let d = Deferred::new(move || {
117             let _ = identity(a);
118             fired.set(true);
119         });
120 
121         assert!(!fired.get());
122         d.call();
123         assert!(fired.get());
124     }
125 
126     #[test]
string()127     fn string() {
128         let a = "hello".to_string();
129         let d = Deferred::new(move || assert_eq!(a, "hello"));
130         d.call();
131     }
132 
133     #[test]
boxed_slice_i32()134     fn boxed_slice_i32() {
135         let a: Box<[i32]> = vec![2, 3, 5, 7].into_boxed_slice();
136         let d = Deferred::new(move || assert_eq!(*a, [2, 3, 5, 7]));
137         d.call();
138     }
139 
140     #[test]
long_slice_usize()141     fn long_slice_usize() {
142         let a: [usize; 5] = [2, 3, 5, 7, 11];
143         let d = Deferred::new(move || assert_eq!(a, [2, 3, 5, 7, 11]));
144         d.call();
145     }
146 }
147