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
109
110
111
112
113
114
115
116
117
118
119
120
121

//! Possible [`std::thread`] additions. Contains examples.
//!
//! # Examples
//!
//! Static JoinGuard self ownership
//!
//! ```
//! use leak_playground_std::*;
//! let (tx, rx) = sync::mpsc::rendezvous_channel();
//! let thrd = thread::spawn_scoped(move || {
//!     let _this_thread = rx.recv().unwrap();
//! });
//! tx.send(thrd).unwrap();
//! ```

use std::thread::JoinHandle;
use std::{marker::PhantomData, thread};

use crate::marker::{Forget, Unforget};
use crate::mem::{self, ManuallyDrop};
use crate::rc::Rc;
use crate::sync::Arc;

/// Spawn borrowing thread handles.
pub fn spawn_scoped<'a, F, T>(f: F) -> JoinGuard<'a, T>
where
    F: FnOnce() -> T + Send + 'a,
    T: Send + 'a,
{
    JoinGuard {
        // SAFETY: destruction guarantee from `Unforget<&'a ()>` and `T: 'a`
        child: unsafe {
            ManuallyDrop::new_unchecked(thread::Builder::new().spawn_unchecked(f).unwrap())
        },
        _borrow: Unforget::new(PhantomData),
        _unsend: PhantomData,
    }
}

/// Handle to a thread, which joins on drop.
///
/// Cannot be sent across threads.
/// This is made to ensure we won't put this into itself, thus forgetting it.
///
/// To spawn use [`spawn_scoped`].
pub struct JoinGuard<'a, T> {
    child: ManuallyDrop<thread::JoinHandle<T>>,

    /// Not sure about covariance there.
    _borrow: Unforget<'static, PhantomData<&'a ()>>,
    _unsend: PhantomData<*mut ()>,
}

unsafe impl<T> Send for JoinGuard<'_, T> where Self: Forget {}
unsafe impl<T> Sync for JoinGuard<'_, T> {}

impl<T> JoinGuard<'_, T> {
    pub fn join(mut self) -> std::thread::Result<T> {
        let join_handle;
        // SAFETY: we immediately, join after
        unsafe {
            join_handle = ManuallyDrop::take(&mut self.child);
            // need this to avoid calling `JoinGuard::drop`
            mem::forget_unchecked(self);
        }
        join_handle.join()
    }

    pub fn thread(&self) -> &std::thread::Thread {
        self.child.thread()
    }

    pub fn is_finished(&self) -> bool {
        self.child.is_finished()
    }

    pub fn into_rc(self) -> Rc<Self> {
        // SAFETY: we cannot move Rc<JoinGuard> into it's closure because
        //   impl !Send for Rc<JoinGuard>
        unsafe { Rc::new_unchecked(self) }
    }

    pub fn into_arc(self) -> Arc<Self> {
        // SAFETY: we cannot move Arc<JoinGuard> into it's closure because
        //   Arc<JoinGuard>: !Send, or otherwise JoinGuard: Forget
        unsafe { Arc::new_unchecked(self) }
    }
}

impl<T> JoinGuard<'static, T> {
    pub fn into_handle(self) -> JoinHandle<T> {
        self.into()
    }

    pub fn detach(self) {
        let _ = self.into_handle();
    }
}

impl<T> From<JoinGuard<'static, T>> for JoinHandle<T> {
    fn from(mut value: JoinGuard<'static, T>) -> Self {
        unsafe { ManuallyDrop::take(&mut value.child) }
    }
}

impl<'a, T> Drop for JoinGuard<'a, T> {
    fn drop(&mut self) {
        let join_handle = unsafe { ManuallyDrop::take(&mut self.child) };
        // Shouldn't panic
        let child = join_handle.thread().clone();
        // No panic since we guarantee that we would never join on ourselves,
        // except when `Self: Forget`, then we don't care.
        let res = join_handle.join();
        // Propagating panic there since structured parallelism, but ignoring
        // during panic. Anyway child thread is joined thus either would
        // be fine.
        if res.is_err() && !std::thread::panicking() {
            panic!("child thread {child:?} panicked");
        }
    }
}