panic_room/
lib.rs

1// Need this to actually recover propertly, I don't know how to gain access.
2// #![feature(update_panic_count)]
3
4use setjmp::{jmp_buf, longjmp, setjmp};
5use std::{
6    any::Any,
7    cell::{Ref, RefCell, RefMut},
8    panic::PanicInfo,
9};
10
11thread_local! {
12    static NEXT_ROOM_ID: std::cell::RefCell<u64> = std::cell::RefCell::new(0);
13    static ROOMS: std::cell::RefCell<Vec<Room>> = std::cell::RefCell::new(Vec::new());
14}
15
16pub struct Room {
17    id: u64,
18    data: Vec<RefCell<Option<Box<dyn Any>>>>,
19    jmp_buf: jmp_buf,
20    #[allow(dead_code)]
21    old_panic_hook: Box<dyn Fn(&std::panic::PanicInfo<'_>) + Sync + Send>,
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
25pub struct Handle<T> {
26    id: u64,
27    index: usize,
28    _phantom: std::marker::PhantomData<T>,
29}
30
31// private helpers
32impl Room {
33    fn push_room() -> Option<()> {
34        let id = NEXT_ROOM_ID.with(|next_room_id| {
35            let mut next_room_id = next_room_id.borrow_mut();
36            let id = *next_room_id;
37            *next_room_id += 1;
38            id
39        });
40        ROOMS.with(|rooms| {
41            let mut rooms = rooms.try_borrow_mut().ok()?;
42            rooms.reserve(1);
43            let new_panic_hook: Box<dyn Fn(&PanicInfo<'_>) + 'static + Sync + Send> =
44                Box::new(|_panic_info| {
45                    let jmp_buf = Self::current_jmp_buf();
46                    if let Some(jmp_buf) = jmp_buf {
47                        unsafe { longjmp(jmp_buf, 1) };
48                    }
49                });
50            let data = Vec::new();
51            let jmp_buf = unsafe { std::mem::zeroed() };
52            let old_panic_hook = std::panic::take_hook();
53            let room = Room {
54                id,
55                data,
56                jmp_buf,
57                old_panic_hook,
58            };
59            std::panic::set_hook(new_panic_hook);
60            rooms.push(room);
61            Some(())
62        })
63    }
64
65    fn pop_room() -> Option<()> {
66        ROOMS.with(|rooms| {
67            let mut rooms = rooms.try_borrow_mut().ok()?;
68            if let Some(_room) = rooms.pop() {
69                // Can't quite recover old state here since we have to tell
70                // the runtime we're not panicking, which needs access to the
71                // internal panic count. feature(update_panic_count) is not
72                // available to the public however.
73                // std::panic::set_hook(room.old_panic_hook);
74                Some(())
75            } else {
76                None
77            }
78        })
79    }
80
81    fn current_jmp_buf() -> Option<*mut jmp_buf> {
82        Room::with_current_mut(|room| Some(&mut room?.jmp_buf as *mut jmp_buf))
83    }
84}
85
86// public API
87impl Room {
88    pub fn with_current_mut<T, F>(f: F) -> Option<T>
89    where
90        F: FnOnce(Option<&mut Room>) -> Option<T>,
91    {
92        ROOMS.with(|rooms| {
93            let mut rooms = rooms.try_borrow_mut().ok()?;
94            let room = rooms.last_mut();
95            f(room)
96        })
97    }
98
99    pub fn with_current<T, F>(f: F) -> Option<T>
100    where
101        F: FnOnce(Option<&Room>) -> Option<T>,
102    {
103        ROOMS.with(|rooms| {
104            let rooms = rooms.try_borrow().ok()?;
105            let room = rooms.last();
106            f(room)
107        })
108    }
109
110    pub fn contain_panics<T, F>(f: F) -> Option<T>
111    where
112        F: FnOnce() -> T,
113    {
114        Room::push_room()?;
115        let Some(jmp_buf) = Room::current_jmp_buf() else {
116            Room::pop_room();
117            return None;
118        };
119        let res = if unsafe { setjmp(jmp_buf) } == 0 {
120            Some(f())
121        } else {
122            // Need access to this internal API to recover. I don't know
123            // the right compiler magic to get it.
124            // std::panicking::panic_count::decrease();
125            None
126        };
127        Room::pop_room();
128        res
129    }
130
131    pub fn alloc<T: 'static>(&mut self, value: T) -> Handle<T> {
132        let index = self.data.len();
133        self.data.push(RefCell::new(Some(Box::new(value))));
134        Handle {
135            id: self.id,
136            index,
137            _phantom: std::marker::PhantomData,
138        }
139    }
140
141    pub fn get<T: 'static>(&self, handle: Handle<T>) -> Option<Ref<T>> {
142        if handle.id == self.id {
143            if let Some(r) = self.data.get(handle.index)?.try_borrow().ok() {
144                Ref::filter_map(r, |any_opt| {
145                    if let Some(any) = any_opt {
146                        any.downcast_ref::<T>()
147                    } else {
148                        None
149                    }
150                })
151                .ok()
152            } else {
153                None
154            }
155        } else {
156            None
157        }
158    }
159
160    pub fn get_mut<T: 'static>(&self, handle: Handle<T>) -> Option<RefMut<T>> {
161        if handle.id == self.id {
162            if let Some(r) = self.data.get(handle.index)?.try_borrow_mut().ok() {
163                RefMut::filter_map(r, |any_opt| {
164                    if let Some(any) = any_opt {
165                        any.downcast_mut::<T>()
166                    } else {
167                        None
168                    }
169                })
170                .ok()
171            } else {
172                None
173            }
174        } else {
175            None
176        }
177    }
178
179    pub fn take<T: 'static>(&mut self, handle: Handle<T>) -> Option<T> {
180        if handle.id == self.id {
181            self.data
182                .get(handle.index)?
183                .borrow_mut()
184                .take()?
185                .downcast()
186                .ok()
187                .map(|x| *x)
188        } else {
189            None
190        }
191    }
192}
panic_room/lib.rs

panic_room/
lib.rs
