polars_python/gil_once_cell.rs
1use std::cell::UnsafeCell;
2
3use pyo3::{PyResult, Python};
4
5// Adapted from PYO3 with the only change that
6// we allow mutable access with when the GIL is held
7
8pub struct GILOnceCell<T>(UnsafeCell<Option<T>>);
9
10// T: Send is needed for Sync because the thread which drops the GILOnceCell can be different
11// to the thread which fills it.
12unsafe impl<T: Send + Sync> Sync for GILOnceCell<T> {}
13unsafe impl<T: Send> Send for GILOnceCell<T> {}
14
15impl<T> GILOnceCell<T> {
16 /// Create a `GILOnceCell` which does not yet contain a value.
17 #[allow(clippy::new_without_default)]
18 pub const fn new() -> Self {
19 Self(UnsafeCell::new(None))
20 }
21
22 /// as long as we have the GIL we can mutate
23 /// this creates a context that checks that.
24 pub fn with_gil<F, O>(&self, _py: Python<'_>, mut op: F) -> PyResult<O>
25 where
26 F: FnMut(&mut T) -> PyResult<O>,
27 {
28 // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
29 let inner = unsafe { &mut *self.0.get() }
30 .as_mut()
31 .expect("not yet initialized");
32
33 op(inner)
34 }
35
36 /// Set the value in the cell.
37 ///
38 /// If the cell has already been written, `Err(value)` will be returned containing the new
39 /// value which was not written.
40 pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
41 // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
42 let inner = unsafe { &mut *self.0.get() };
43 if inner.is_some() {
44 return Err(value);
45 }
46
47 *inner = Some(value);
48 Ok(())
49 }
50}