pen_ffi/
closure.rs

1use crate::Arc;
2use core::{
3    mem::{transmute, ManuallyDrop},
4    ops::Deref,
5    sync::atomic::{AtomicPtr, Ordering},
6};
7
8struct ClosureMetadata<T> {
9    drop: extern "C" fn(&mut ClosureInner<T>),
10    #[allow(dead_code)]
11    synchronize: extern "C" fn(&mut ClosureInner<T>),
12}
13
14#[repr(C)]
15#[derive(Clone)]
16pub struct Closure<T = ()>(Arc<ClosureInner<T>>);
17
18#[repr(C)]
19struct ClosureInner<T> {
20    entry_function: AtomicPtr<u8>,
21    metadata: AtomicPtr<ClosureMetadata<T>>,
22    payload: ManuallyDrop<T>,
23}
24
25impl<T> Closure<T> {
26    const METADATA: ClosureMetadata<T> = ClosureMetadata {
27        drop: drop_closure::<T>,
28        synchronize: synchronize_closure::<T>,
29    };
30
31    pub fn new(entry_function: *const u8, payload: T) -> Self {
32        Self(
33            ClosureInner {
34                entry_function: AtomicPtr::new(entry_function as *mut u8),
35                metadata: AtomicPtr::new(&Self::METADATA as *const _ as *mut _),
36                payload: ManuallyDrop::new(payload),
37            }
38            .into(),
39        )
40    }
41
42    pub fn entry_function(&self) -> *const u8 {
43        self.0.entry_function.load(Ordering::Relaxed)
44    }
45
46    pub fn payload(&self) -> *const T {
47        self.0.payload.deref()
48    }
49
50    pub fn into_opaque(self) -> Closure<()> {
51        unsafe { transmute(self) }
52    }
53}
54
55extern "C" fn drop_closure<T>(closure: &mut ClosureInner<T>) {
56    unsafe { ManuallyDrop::drop(&mut closure.payload) }
57}
58
59// All closures created in Rust should implement Sync already.
60extern "C" fn synchronize_closure<T>(_: &mut ClosureInner<T>) {}
61
62impl<T> Drop for ClosureInner<T> {
63    fn drop(&mut self) {
64        let metadata = unsafe { &*self.metadata.load(Ordering::Relaxed) };
65
66        (metadata.drop)(self);
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use alloc::boxed::Box;
74    use core::{ptr::null, sync::atomic::AtomicBool};
75
76    fn spawn<T: Send + 'static>(_: impl (FnOnce() -> T) + Send + 'static) {}
77
78    #[test]
79    fn send() {
80        let closure = Closure::new(null(), ());
81
82        spawn(move || {
83            closure.entry_function();
84        });
85    }
86
87    #[test]
88    fn drop_payload() {
89        struct Foo {}
90
91        static FLAG: AtomicBool = AtomicBool::new(false);
92
93        impl Drop for Foo {
94            fn drop(&mut self) {
95                FLAG.store(true, Ordering::SeqCst);
96            }
97        }
98
99        Closure::new(null(), Foo {});
100
101        assert!(FLAG.load(Ordering::SeqCst));
102    }
103
104    #[test]
105    fn drop_boxed_payload() {
106        Closure::new(null(), Box::new(42.0));
107    }
108}
pen_ffi/closure.rs

pen_ffi/
closure.rs
