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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
//! A write-once cell mediated by the Python GIL.
use crate::Python;
use std::cell::UnsafeCell;

/// A write-once cell similar to [`once_cell::OnceCell`](https://docs.rs/once_cell/1.4.0/once_cell/).
///
/// Unlike `once_cell::sync` which blocks threads to achieve thread safety, this implementation
/// uses the Python GIL to mediate concurrent access. This helps in cases where `once_sync` or
/// `lazy_static`'s synchronization strategy can lead to deadlocks when interacting with the Python
/// GIL. For an example, see [the FAQ section](https://pyo3.rs/latest/faq.html) of the guide.
///
/// # Examples
///
/// The following example shows how to use `GILOnceCell` to share a reference to a Python list
/// between threads:
///
/// ```
/// use pyo3::once_cell::GILOnceCell;
/// use pyo3::prelude::*;
/// use pyo3::types::PyList;
///
/// static LIST_CELL: GILOnceCell<Py<PyList>> = GILOnceCell::new();
///
/// pub fn get_shared_list(py: Python<'_>) -> &PyList {
///     LIST_CELL
///         .get_or_init(py, || PyList::empty(py).into())
///         .as_ref(py)
/// }
/// # Python::with_gil(|py| assert_eq!(get_shared_list(py).len(), 0));
/// ```
pub struct GILOnceCell<T>(UnsafeCell<Option<T>>);

// T: Send is needed for Sync because the thread which drops the GILOnceCell can be different
// to the thread which fills it.
unsafe impl<T: Send + Sync> Sync for GILOnceCell<T> {}
unsafe impl<T: Send> Send for GILOnceCell<T> {}

impl<T> GILOnceCell<T> {
    /// Create a `GILOnceCell` which does not yet contain a value.
    pub const fn new() -> Self {
        Self(UnsafeCell::new(None))
    }

    /// Get a reference to the contained value, or `None` if the cell has not yet been written.
    #[inline]
    pub fn get(&self, _py: Python<'_>) -> Option<&T> {
        // Safe because if the cell has not yet been written, None is returned.
        unsafe { &*self.0.get() }.as_ref()
    }

    /// Get a reference to the contained value, initializing it if needed using the provided
    /// closure.
    ///
    /// Note that:
    ///  1) reentrant initialization can cause a stack overflow.
    ///  2) if f() temporarily releases the GIL (e.g. by calling `Python::import`) then it is
    ///     possible (and well-defined) that a second thread may also call get_or_init and begin
    ///     calling `f()`. Even when this happens `GILOnceCell` guarantees that only **one** write
    ///     to the cell ever occurs - other threads will simply discard the value they compute and
    ///     return the result of the first complete computation.
    ///  3) if f() does not release the GIL and does not panic, it is guaranteed to be called
    ///     exactly once, even if multiple threads attempt to call `get_or_init`
    ///  4) if f() can panic but still does not release the GIL, it may be called multiple times,
    ///     but it is guaranteed that f() will never be called concurrently
    #[inline]
    pub fn get_or_init<F>(&self, py: Python<'_>, f: F) -> &T
    where
        F: FnOnce() -> T,
    {
        if let Some(value) = self.get(py) {
            return value;
        }

        self.init(py, f)
    }

    #[cold]
    fn init<F>(&self, py: Python<'_>, f: F) -> &T
    where
        F: FnOnce() -> T,
    {
        // Note that f() could temporarily release the GIL, so it's possible that another thread
        // writes to this GILOnceCell before f() finishes. That's fine; we'll just have to discard
        // the value computed here and accept a bit of wasted computation.
        let value = f();
        let _ = self.set(py, value);

        self.get(py).unwrap()
    }

    /// Get the contents of the cell mutably. This is only possible if the reference to the cell is
    /// unique.
    pub fn get_mut(&mut self) -> Option<&mut T> {
        // Safe because we have &mut self
        unsafe { &mut *self.0.get() }.as_mut()
    }

    /// Set the value in the cell.
    ///
    /// If the cell has already been written, `Err(value)` will be returned containing the new
    /// value which was not written.
    pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
        // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
        let inner = unsafe { &mut *self.0.get() };
        if inner.is_some() {
            return Err(value);
        }

        *inner = Some(value);
        Ok(())
    }
}

/// Interns `text` as a Python string and stores a reference to it in static storage.
///
/// A reference to the same Python string is returned on each invocation.
///
/// # Example: Using `intern!` to avoid needlessly recreating the same Python string
///
/// ```
/// use pyo3::intern;
/// # use pyo3::{pyfunction, types::PyDict, wrap_pyfunction, PyResult, Python};
///
/// #[pyfunction]
/// fn create_dict(py: Python<'_>) -> PyResult<&PyDict> {
///    let dict = PyDict::new(py);
///    //             👇 A new `PyString` is created
///    //                for every call of this function.
///    dict.set_item("foo", 42)?;
///    Ok(dict)
/// }
///
/// #[pyfunction]
/// fn create_dict_faster(py: Python<'_>) -> PyResult<&PyDict> {
///    let dict = PyDict::new(py);
///    //               👇 A `PyString` is created once and reused
///    //                  for the lifetime of the program.
///    dict.set_item(intern!(py, "foo"), 42)?;
///    Ok(dict)
/// }
/// #
/// # Python::with_gil(|py| {
/// #     let fun = wrap_pyfunction!(create_dict_faster, py).unwrap();
/// #     let dict = fun.call0().unwrap();
/// #     assert!(dict.contains("foo").unwrap());
/// # });
/// ```
#[macro_export]
macro_rules! intern {
    ($py: expr, $text: expr) => {{
        fn isolate_from_dyn_env(py: $crate::Python<'_>) -> &$crate::types::PyString {
            static INTERNED: $crate::once_cell::GILOnceCell<$crate::Py<$crate::types::PyString>> =
                $crate::once_cell::GILOnceCell::new();

            INTERNED
                .get_or_init(py, || {
                    $crate::conversion::IntoPy::into_py(
                        $crate::types::PyString::intern(py, $text),
                        py,
                    )
                })
                .as_ref(py)
        }

        isolate_from_dyn_env($py)
    }};
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::types::PyDict;

    #[test]
    fn test_intern() {
        Python::with_gil(|py| {
            let foo1 = "foo";
            let foo2 = intern!(py, "foo");
            let foo3 = intern!(py, stringify!(foo));

            let dict = PyDict::new(py);
            dict.set_item(foo1, 42_usize).unwrap();
            assert!(dict.contains(foo2).unwrap());
            assert_eq!(dict.get_item(foo3).unwrap().extract::<usize>().unwrap(), 42);
        });
    }
}