1#![allow(missing_docs)]
2//! Crate-private implementation of pycell
3
4use std::cell::Cell;
5use std::marker::PhantomData;
6
7use crate::impl_::pyclass::{PyClassBaseType, PyClassImpl};
8use crate::PyCell;
9
10use super::{PyBorrowError, PyBorrowMutError};
11
12pub trait PyClassMutability {
13 // The storage for this inheritance layer. Only the first mutable class in
14 // an inheritance hierarchy needs to store the borrow flag.
15 type Storage: PyClassBorrowChecker;
16 // The borrow flag needed to implement this class' mutability. Empty until
17 // the first mutable class, at which point it is BorrowChecker and will be
18 // for all subclasses.
19 type Checker: PyClassBorrowChecker;
20 type ImmutableChild: PyClassMutability;
21 type MutableChild: PyClassMutability;
22}
23
24pub struct ImmutableClass(());
25pub struct MutableClass(());
26pub struct ExtendsMutableAncestor<M: PyClassMutability>(PhantomData<M>);
27
28impl PyClassMutability for ImmutableClass {
29 type Storage = EmptySlot;
30 type Checker = EmptySlot;
31 type ImmutableChild = ImmutableClass;
32 type MutableChild = MutableClass;
33}
34
35impl PyClassMutability for MutableClass {
36 type Storage = BorrowChecker;
37 type Checker = BorrowChecker;
38 type ImmutableChild = ExtendsMutableAncestor<ImmutableClass>;
39 type MutableChild = ExtendsMutableAncestor<MutableClass>;
40}
41
42impl<M: PyClassMutability> PyClassMutability for ExtendsMutableAncestor<M> {
43 type Storage = EmptySlot;
44 type Checker = BorrowChecker;
45 type ImmutableChild = ExtendsMutableAncestor<ImmutableClass>;
46 type MutableChild = ExtendsMutableAncestor<MutableClass>;
47}
48
49#[derive(Debug, Copy, Clone, Eq, PartialEq)]
50struct BorrowFlag(usize);
51
52impl BorrowFlag {
53 pub(crate) const UNUSED: BorrowFlag = BorrowFlag(0);
54 const HAS_MUTABLE_BORROW: BorrowFlag = BorrowFlag(usize::max_value());
55 const fn increment(self) -> Self {
56 Self(self.0 + 1)
57 }
58 const fn decrement(self) -> Self {
59 Self(self.0 - 1)
60 }
61}
62
63pub struct EmptySlot(());
64pub struct BorrowChecker(Cell<BorrowFlag>);
65
66pub trait PyClassBorrowChecker {
67 /// Initial value for self
68 fn new() -> Self;
69
70 /// Increments immutable borrow count, if possible
71 fn try_borrow(&self) -> Result<(), PyBorrowError>;
72
73 fn try_borrow_unguarded(&self) -> Result<(), PyBorrowError>;
74
75 /// Decrements immutable borrow count
76 fn release_borrow(&self);
77 /// Increments mutable borrow count, if possible
78 fn try_borrow_mut(&self) -> Result<(), PyBorrowMutError>;
79 /// Decremements mutable borrow count
80 fn release_borrow_mut(&self);
81}
82
83impl PyClassBorrowChecker for EmptySlot {
84 #[inline]
85 fn new() -> Self {
86 EmptySlot(())
87 }
88
89 #[inline]
90 fn try_borrow(&self) -> Result<(), PyBorrowError> {
91 Ok(())
92 }
93
94 #[inline]
95 fn try_borrow_unguarded(&self) -> Result<(), PyBorrowError> {
96 Ok(())
97 }
98
99 #[inline]
100 fn release_borrow(&self) {}
101
102 #[inline]
103 fn try_borrow_mut(&self) -> Result<(), PyBorrowMutError> {
104 unreachable!()
105 }
106
107 #[inline]
108 fn release_borrow_mut(&self) {
109 unreachable!()
110 }
111}
112
113impl PyClassBorrowChecker for BorrowChecker {
114 #[inline]
115 fn new() -> Self {
116 Self(Cell::new(BorrowFlag::UNUSED))
117 }
118
119 fn try_borrow(&self) -> Result<(), PyBorrowError> {
120 let flag = self.0.get();
121 if flag != BorrowFlag::HAS_MUTABLE_BORROW {
122 self.0.set(flag.increment());
123 Ok(())
124 } else {
125 Err(PyBorrowError { _private: () })
126 }
127 }
128
129 fn try_borrow_unguarded(&self) -> Result<(), PyBorrowError> {
130 let flag = self.0.get();
131 if flag != BorrowFlag::HAS_MUTABLE_BORROW {
132 Ok(())
133 } else {
134 Err(PyBorrowError { _private: () })
135 }
136 }
137
138 fn release_borrow(&self) {
139 let flag = self.0.get();
140 self.0.set(flag.decrement())
141 }
142
143 fn try_borrow_mut(&self) -> Result<(), PyBorrowMutError> {
144 let flag = self.0.get();
145 if flag == BorrowFlag::UNUSED {
146 self.0.set(BorrowFlag::HAS_MUTABLE_BORROW);
147 Ok(())
148 } else {
149 Err(PyBorrowMutError { _private: () })
150 }
151 }
152
153 fn release_borrow_mut(&self) {
154 self.0.set(BorrowFlag::UNUSED)
155 }
156}
157
158pub trait GetBorrowChecker<T: PyClassImpl> {
159 fn borrow_checker(cell: &PyCell<T>) -> &<T::PyClassMutability as PyClassMutability>::Checker;
160}
161
162impl<T: PyClassImpl<PyClassMutability = Self>> GetBorrowChecker<T> for MutableClass {
163 fn borrow_checker(cell: &PyCell<T>) -> &BorrowChecker {
164 &cell.contents.borrow_checker
165 }
166}
167
168impl<T: PyClassImpl<PyClassMutability = Self>> GetBorrowChecker<T> for ImmutableClass {
169 fn borrow_checker(cell: &PyCell<T>) -> &EmptySlot {
170 &cell.contents.borrow_checker
171 }
172}
173
174impl<T: PyClassImpl<PyClassMutability = Self>, M: PyClassMutability> GetBorrowChecker<T>
175 for ExtendsMutableAncestor<M>
176where
177 T::BaseType: PyClassImpl + PyClassBaseType<LayoutAsBase = PyCell<T::BaseType>>,
178 <T::BaseType as PyClassImpl>::PyClassMutability: PyClassMutability<Checker = BorrowChecker>,
179{
180 fn borrow_checker(cell: &PyCell<T>) -> &BorrowChecker {
181 <<T::BaseType as PyClassImpl>::PyClassMutability as GetBorrowChecker<T::BaseType>>::borrow_checker(&cell.ob_base)
182 }
183}
184
185#[cfg(test)]
186#[cfg(feature = "macros")]
187mod tests {
188 use super::*;
189
190 use crate::prelude::*;
191 use crate::pyclass::boolean_struct::{False, True};
192 use crate::PyClass;
193
194 #[pyclass(crate = "crate", subclass)]
195 struct MutableBase;
196
197 #[pyclass(crate = "crate", extends = MutableBase, subclass)]
198 struct MutableChildOfMutableBase;
199
200 #[pyclass(crate = "crate", extends = MutableBase, frozen, subclass)]
201 struct ImmutableChildOfMutableBase;
202
203 #[pyclass(crate = "crate", extends = MutableChildOfMutableBase)]
204 struct MutableChildOfMutableChildOfMutableBase;
205
206 #[pyclass(crate = "crate", extends = ImmutableChildOfMutableBase)]
207 struct MutableChildOfImmutableChildOfMutableBase;
208
209 #[pyclass(crate = "crate", extends = MutableChildOfMutableBase, frozen)]
210 struct ImmutableChildOfMutableChildOfMutableBase;
211
212 #[pyclass(crate = "crate", extends = ImmutableChildOfMutableBase, frozen)]
213 struct ImmutableChildOfImmutableChildOfMutableBase;
214
215 #[pyclass(crate = "crate", frozen, subclass)]
216 struct ImmutableBase;
217
218 #[pyclass(crate = "crate", extends = ImmutableBase, subclass)]
219 struct MutableChildOfImmutableBase;
220
221 #[pyclass(crate = "crate", extends = ImmutableBase, frozen, subclass)]
222 struct ImmutableChildOfImmutableBase;
223
224 #[pyclass(crate = "crate", extends = MutableChildOfImmutableBase)]
225 struct MutableChildOfMutableChildOfImmutableBase;
226
227 #[pyclass(crate = "crate", extends = ImmutableChildOfImmutableBase)]
228 struct MutableChildOfImmutableChildOfImmutableBase;
229
230 #[pyclass(crate = "crate", extends = MutableChildOfImmutableBase, frozen)]
231 struct ImmutableChildOfMutableChildOfImmutableBase;
232
233 #[pyclass(crate = "crate", extends = ImmutableChildOfImmutableBase, frozen)]
234 struct ImmutableChildOfImmutableChildOfImmutableBase;
235
236 fn assert_mutable<T: PyClass<Frozen = False, PyClassMutability = MutableClass>>() {}
237 fn assert_immutable<T: PyClass<Frozen = True, PyClassMutability = ImmutableClass>>() {}
238 fn assert_mutable_with_mutable_ancestor<
239 T: PyClass<Frozen = False, PyClassMutability = ExtendsMutableAncestor<MutableClass>>,
240 >() {
241 }
242 fn assert_immutable_with_mutable_ancestor<
243 T: PyClass<Frozen = True, PyClassMutability = ExtendsMutableAncestor<ImmutableClass>>,
244 >() {
245 }
246
247 #[test]
248 fn test_inherited_mutability() {
249 // mutable base
250 assert_mutable::<MutableBase>();
251
252 // children of mutable base have a mutable ancestor
253 assert_mutable_with_mutable_ancestor::<MutableChildOfMutableBase>();
254 assert_immutable_with_mutable_ancestor::<ImmutableChildOfMutableBase>();
255
256 // grandchildren of mutable base have a mutable ancestor
257 assert_mutable_with_mutable_ancestor::<MutableChildOfMutableChildOfMutableBase>();
258 assert_mutable_with_mutable_ancestor::<MutableChildOfImmutableChildOfMutableBase>();
259 assert_immutable_with_mutable_ancestor::<ImmutableChildOfMutableChildOfMutableBase>();
260 assert_immutable_with_mutable_ancestor::<ImmutableChildOfImmutableChildOfMutableBase>();
261
262 // immutable base and children
263 assert_immutable::<ImmutableBase>();
264 assert_immutable::<ImmutableChildOfImmutableBase>();
265 assert_immutable::<ImmutableChildOfImmutableChildOfImmutableBase>();
266
267 // mutable children of immutable at any level are simply mutable
268 assert_mutable::<MutableChildOfImmutableBase>();
269 assert_mutable::<MutableChildOfImmutableChildOfImmutableBase>();
270
271 // children of the mutable child display this property
272 assert_mutable_with_mutable_ancestor::<MutableChildOfMutableChildOfImmutableBase>();
273 assert_immutable_with_mutable_ancestor::<ImmutableChildOfMutableChildOfImmutableBase>();
274 }
275
276 #[test]
277 fn test_mutable_borrow_prevents_further_borrows() {
278 Python::with_gil(|py| {
279 let mmm = Py::new(
280 py,
281 PyClassInitializer::from(MutableBase)
282 .add_subclass(MutableChildOfMutableBase)
283 .add_subclass(MutableChildOfMutableChildOfMutableBase),
284 )
285 .unwrap();
286
287 let mmm_cell: &PyCell<MutableChildOfMutableChildOfMutableBase> = mmm.as_ref(py);
288
289 let mmm_refmut = mmm_cell.borrow_mut();
290
291 // Cannot take any other mutable or immutable borrows whilst the object is borrowed mutably
292 assert!(mmm_cell
293 .extract::<PyRef<'_, MutableChildOfMutableChildOfMutableBase>>()
294 .is_err());
295 assert!(mmm_cell
296 .extract::<PyRef<'_, MutableChildOfMutableBase>>()
297 .is_err());
298 assert!(mmm_cell.extract::<PyRef<'_, MutableBase>>().is_err());
299 assert!(mmm_cell
300 .extract::<PyRefMut<'_, MutableChildOfMutableChildOfMutableBase>>()
301 .is_err());
302 assert!(mmm_cell
303 .extract::<PyRefMut<'_, MutableChildOfMutableBase>>()
304 .is_err());
305 assert!(mmm_cell.extract::<PyRefMut<'_, MutableBase>>().is_err());
306
307 // With the borrow dropped, all other borrow attempts will succeed
308 drop(mmm_refmut);
309
310 assert!(mmm_cell
311 .extract::<PyRef<'_, MutableChildOfMutableChildOfMutableBase>>()
312 .is_ok());
313 assert!(mmm_cell
314 .extract::<PyRef<'_, MutableChildOfMutableBase>>()
315 .is_ok());
316 assert!(mmm_cell.extract::<PyRef<'_, MutableBase>>().is_ok());
317 assert!(mmm_cell
318 .extract::<PyRefMut<'_, MutableChildOfMutableChildOfMutableBase>>()
319 .is_ok());
320 assert!(mmm_cell
321 .extract::<PyRefMut<'_, MutableChildOfMutableBase>>()
322 .is_ok());
323 assert!(mmm_cell.extract::<PyRefMut<'_, MutableBase>>().is_ok());
324 })
325 }
326
327 #[test]
328 fn test_immutable_borrows_prevent_mutable_borrows() {
329 Python::with_gil(|py| {
330 let mmm = Py::new(
331 py,
332 PyClassInitializer::from(MutableBase)
333 .add_subclass(MutableChildOfMutableBase)
334 .add_subclass(MutableChildOfMutableChildOfMutableBase),
335 )
336 .unwrap();
337
338 let mmm_cell: &PyCell<MutableChildOfMutableChildOfMutableBase> = mmm.as_ref(py);
339
340 let mmm_refmut = mmm_cell.borrow();
341
342 // Further immutable borrows are ok
343 assert!(mmm_cell
344 .extract::<PyRef<'_, MutableChildOfMutableChildOfMutableBase>>()
345 .is_ok());
346 assert!(mmm_cell
347 .extract::<PyRef<'_, MutableChildOfMutableBase>>()
348 .is_ok());
349 assert!(mmm_cell.extract::<PyRef<'_, MutableBase>>().is_ok());
350
351 // Further mutable borrows are not ok
352 assert!(mmm_cell
353 .extract::<PyRefMut<'_, MutableChildOfMutableChildOfMutableBase>>()
354 .is_err());
355 assert!(mmm_cell
356 .extract::<PyRefMut<'_, MutableChildOfMutableBase>>()
357 .is_err());
358 assert!(mmm_cell.extract::<PyRefMut<'_, MutableBase>>().is_err());
359
360 // With the borrow dropped, all mutable borrow attempts will succeed
361 drop(mmm_refmut);
362
363 assert!(mmm_cell
364 .extract::<PyRefMut<'_, MutableChildOfMutableChildOfMutableBase>>()
365 .is_ok());
366 assert!(mmm_cell
367 .extract::<PyRefMut<'_, MutableChildOfMutableBase>>()
368 .is_ok());
369 assert!(mmm_cell.extract::<PyRefMut<'_, MutableBase>>().is_ok());
370 })
371 }
372}
373