1use alloc::boxed::Box;
2use core::fmt;
3use core::marker::PhantomData;
4use core::mem::{self, MaybeUninit};
5use 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.
11const DATA_WORDS: usize = 3;
12
13/// Some space to keep a `FnOnce()` object on the stack.
14type 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.
19pub(crate) struct Deferred {
20 call: unsafe fn(*mut u8),
21 data: MaybeUninit<Data>,
22 _marker: PhantomData<*mut ()>, // !Send + !Sync
23}
24
25impl fmt::Debug for Deferred {
26 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
27 f.pad("Deferred { .. }")
28 }
29}
30
31impl Deferred {
32 pub(crate) const NO_OP: Self = {
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()`.
42 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]
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)))]
91mod tests {
92 use super::Deferred;
93 use std::cell::Cell;
94 use std::convert::identity;
95
96 #[test]
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]
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]
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]
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]
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