1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use crate::Arc;
use core::{
    mem::{transmute, ManuallyDrop},
    ops::Deref,
    sync::atomic::{AtomicPtr, Ordering},
};

struct ClosureMetadata<T> {
    drop: extern "C" fn(&mut ClosureInner<T>),
    #[allow(dead_code)]
    synchronize: extern "C" fn(&mut ClosureInner<T>),
}

#[repr(C)]
#[derive(Clone)]
pub struct Closure<T = ()>(Arc<ClosureInner<T>>);

#[repr(C)]
struct ClosureInner<T> {
    entry_function: AtomicPtr<u8>,
    metadata: AtomicPtr<ClosureMetadata<T>>,
    payload: ManuallyDrop<T>,
}

impl<T> Closure<T> {
    const METADATA: ClosureMetadata<T> = ClosureMetadata {
        drop: drop_closure::<T>,
        synchronize: synchronize_closure::<T>,
    };

    pub fn new(entry_function: *const u8, payload: T) -> Self {
        Self(
            ClosureInner {
                entry_function: AtomicPtr::new(entry_function as *mut u8),
                metadata: AtomicPtr::new(&Self::METADATA as *const _ as *mut _),
                payload: ManuallyDrop::new(payload),
            }
            .into(),
        )
    }

    pub fn entry_function(&self) -> *const u8 {
        self.0.entry_function.load(Ordering::Relaxed)
    }

    pub fn payload(&self) -> *const T {
        self.0.payload.deref()
    }

    pub fn into_opaque(self) -> Closure<()> {
        unsafe { transmute(self) }
    }
}

extern "C" fn drop_closure<T>(closure: &mut ClosureInner<T>) {
    unsafe { ManuallyDrop::drop(&mut closure.payload) }
}

// All closures created in Rust should implement Sync already.
extern "C" fn synchronize_closure<T>(_: &mut ClosureInner<T>) {}

impl<T> Drop for ClosureInner<T> {
    fn drop(&mut self) {
        let metadata = unsafe { &*self.metadata.load(Ordering::Relaxed) };

        (metadata.drop)(self);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::boxed::Box;
    use core::{ptr::null, sync::atomic::AtomicBool};

    fn spawn<T: Send + 'static>(_: impl (FnOnce() -> T) + Send + 'static) {}

    #[test]
    fn send() {
        let closure = Closure::new(null(), ());

        spawn(move || {
            closure.entry_function();
        });
    }

    #[test]
    fn drop_payload() {
        struct Foo {}

        static FLAG: AtomicBool = AtomicBool::new(false);

        impl Drop for Foo {
            fn drop(&mut self) {
                FLAG.store(true, Ordering::SeqCst);
            }
        }

        Closure::new(null(), Foo {});

        assert!(FLAG.load(Ordering::SeqCst));
    }

    #[test]
    fn drop_boxed_payload() {
        Closure::new(null(), Box::new(42.0));
    }
}