polars_utils/
arc.rs

1use std::mem::MaybeUninit;
2use std::sync::Arc;
3
4pub fn arc_map<T: Clone, F: FnMut(T) -> T>(mut arc: Arc<T>, mut f: F) -> Arc<T> {
5    unsafe {
6        // Make the Arc unique (cloning if necessary).
7        Arc::make_mut(&mut arc);
8
9        // If f panics we must be able to drop the Arc without assuming it is initialized.
10        let mut uninit_arc = Arc::from_raw(Arc::into_raw(arc).cast::<MaybeUninit<T>>());
11
12        // Replace the value inside the arc.
13        let ptr = Arc::get_mut(&mut uninit_arc).unwrap_unchecked() as *mut MaybeUninit<T>;
14        *ptr = MaybeUninit::new(f(ptr.read().assume_init()));
15
16        // Now the Arc is properly initialized again.
17        Arc::from_raw(Arc::into_raw(uninit_arc).cast::<T>())
18    }
19}
20
21pub fn try_arc_map<T: Clone, E, F: FnMut(T) -> Result<T, E>>(
22    mut arc: Arc<T>,
23    mut f: F,
24) -> Result<Arc<T>, E> {
25    unsafe {
26        // Make the Arc unique (cloning if necessary).
27        Arc::make_mut(&mut arc);
28
29        // If f panics we must be able to drop the Arc without assuming it is initialized.
30        let mut uninit_arc = Arc::from_raw(Arc::into_raw(arc).cast::<MaybeUninit<T>>());
31
32        // Replace the value inside the arc.
33        let ptr = Arc::get_mut(&mut uninit_arc).unwrap_unchecked() as *mut MaybeUninit<T>;
34        *ptr = MaybeUninit::new(f(ptr.read().assume_init())?);
35
36        // Now the Arc is properly initialized again.
37        Ok(Arc::from_raw(Arc::into_raw(uninit_arc).cast::<T>()))
38    }
39}
polars_utils/arc.rs

polars_utils/
arc.rs
