1 use pyo3::prelude::*;

     2 use pyo3_sharedref::*;

     3 

     4 #[pyclass]

     5 struct Owner {

     6     string: PySharedRefCell<String>,

     7 }

     8 

     9 #[pymethods]

    10 impl Owner {

    11     #[new]

    12     fn new(s: String) -> Self {

    13         Self {

    14             string: PySharedRefCell::new(s),

    15         }

    16     }

    17 }

    18 

    19 fn with_setup(

    20     test: impl FnOnce(Python<'_>, &Bound<'_, Owner>) -> PyResult<()>,

    21 ) -> PyResult<()> {

    22     pyo3::prepare_freethreaded_python();

    23     Python::with_gil(|py| {

    24         let owner = Bound::new(py, Owner::new("new".to_owned()))?;

    25         test(py, &owner)

    26     })

    27 }

    28 

    29 /// "leak" in the sense of `UnsafePyLeaked` the `string` data field,

    30 /// taking care of all the boilerplate

    31 fn leak_string(owner: &Bound<'_, Owner>) -> UnsafePyLeaked<&'static String> {

    32     let cell = &owner.borrow().string;

    33     let shared_ref = unsafe { cell.borrow(owner) };

    34     shared_ref.leak_immutable()

    35 }

    36 

    37 fn try_leak_string(

    38     owner: &Bound<'_, Owner>,

    39 ) -> Result<UnsafePyLeaked<&'static String>, TryLeakError> {

    40     let cell = &owner.borrow().string;

    41     let shared_ref = unsafe { cell.borrow(owner) };

    42     shared_ref.try_leak_immutable()

    43 }

    44 

    45 /// Mutate the `string` field of `owner` as would be done from Python code

    46 ///

    47 /// This is to simulate normal mutation of the owner object from

    48 /// the Python interpreter. This could be replaced by methods of [`Owner`]

    49 /// (wih closure replaced by a small fixed operations)

    50 /// and perhaps will, once we are done converting the original tests

    51 /// from rust-cpython

    52 fn mutate_string<'py>(

    53     owner: &'py Bound<'py, Owner>,

    54     f: impl FnOnce(&mut String),

    55 ) -> () {

    56     let cell = &owner.borrow_mut().string;

    57     let shared_ref = unsafe { cell.borrow(owner) };

    58     f(&mut shared_ref.borrow_mut());

    59 }

    60 

    61 #[test]

    62 fn test_leaked_borrow() -> PyResult<()> {

    63     with_setup(|py, owner| {

    64         let leaked = leak_string(owner);

    65         let leaked_ref = unsafe { leaked.try_borrow(py) }.unwrap();

    66         assert_eq!(*leaked_ref, "new");

    67         Ok(())

    68     })

    69 }

    70 

    71 #[test]

    72 fn test_leaked_borrow_mut() -> PyResult<()> {

    73     with_setup(|py, owner| {

    74         let leaked = leak_string(owner);

    75         let mut leaked_iter = unsafe { leaked.map(py, |s| s.chars()) };

    76         let mut leaked_ref =

    77             unsafe { leaked_iter.try_borrow_mut(py) }.unwrap();

    78         assert_eq!(leaked_ref.next(), Some('n'));

    79         assert_eq!(leaked_ref.next(), Some('e'));

    80         assert_eq!(leaked_ref.next(), Some('w'));

    81         assert_eq!(leaked_ref.next(), None);

    82         Ok(())

    83     })

    84 }

    85 

    86 #[test]

    87 fn test_leaked_borrow_after_mut() -> PyResult<()> {

    88     with_setup(|py, owner| {

    89         let leaked = leak_string(owner);

    90         mutate_string(owner, String::clear);

    91         assert!(unsafe { leaked.try_borrow(py) }.is_err());

    92         Ok(())

    93     })

    94 }

    95 

    96 #[test]

    97 fn test_leaked_borrow_mut_after_mut() -> PyResult<()> {

    98     with_setup(|py, owner| {

    99         let leaked = leak_string(owner);

   100         let mut leaked_iter = unsafe { leaked.map(py, |s| s.chars()) };

   101 

   102         mutate_string(owner, String::clear);

   103         assert!(unsafe { leaked_iter.try_borrow_mut(py) }.is_err());

   104         Ok(())

   105     })

   106 }

   107 

   108 #[test]

   109 #[should_panic(expected = "map() over invalidated leaked reference")]

   110 fn test_leaked_map_after_mut() {

   111     with_setup(|py, owner| {

   112         let leaked = leak_string(owner);

   113         mutate_string(owner, String::clear);

   114         let _leaked_iter = unsafe { leaked.map(py, |s| s.chars()) };

   115         Ok(())

   116     })

   117     .expect("should already have panicked")

   118 }

   119 

   120 /// run `try_borrow_mut` on the `string` field and assert it is not an error

   121 ///

   122 /// Simply returning the `Result` is not possible, because that is

   123 /// returning a reference to data owned by the function

   124 fn assert_try_borrow_string_mut_ok(owner: &Bound<'_, Owner>) {

   125     let cell = &owner.borrow().string;

   126     let shared_ref = unsafe { cell.borrow(owner) };

   127     assert!(shared_ref.try_borrow_mut().is_ok());

   128 }

   129 

   130 fn assert_try_borrow_string_mut_err(owner: &Bound<'_, Owner>) {

   131     let cell = &owner.borrow().string;

   132     let shared_ref = unsafe { cell.borrow(owner) };

   133     assert!(shared_ref.try_borrow_mut().is_err());

   134 }

   135 

   136 fn assert_try_borrow_string_err(owner: &Bound<'_, Owner>) {

   137     let cell = &owner.borrow().string;

   138     let shared_ref = unsafe { cell.borrow(owner) };

   139     assert!(shared_ref.try_borrow().is_err());

   140 }

   141 

   142 #[test]

   143 fn test_try_borrow_mut_while_leaked_ref() -> PyResult<()> {

   144     with_setup(|py, owner| {

   145         assert_try_borrow_string_mut_ok(owner);

   146         let leaked = leak_string(owner);

   147         {

   148             let _leaked_ref = unsafe { leaked.try_borrow(py) }.unwrap();

   149             assert_try_borrow_string_mut_err(owner);

   150             {

   151                 let _leaked_ref2 = unsafe { leaked.try_borrow(py) }.unwrap();

   152                 assert_try_borrow_string_mut_err(owner);

   153             }

   154             assert_try_borrow_string_mut_err(owner);

   155         }

   156         assert_try_borrow_string_mut_ok(owner);

   157         Ok(())

   158     })

   159 }

   160 

   161 #[test]

   162 fn test_try_borrow_mut_while_leaked_ref_mut() -> PyResult<()> {

   163     with_setup(|py, owner| {

   164         assert_try_borrow_string_mut_ok(owner);

   165         let leaked = leak_string(owner);

   166         let mut leaked_iter = unsafe { leaked.map(py, |s| s.chars()) };

   167         {

   168             let _leaked_ref =

   169                 unsafe { leaked_iter.try_borrow_mut(py) }.unwrap();

   170             assert_try_borrow_string_mut_err(owner);

   171         }

   172         assert_try_borrow_string_mut_ok(owner);

   173         Ok(())

   174     })

   175 }

   176 

   177 #[test]

   178 fn test_try_leak_while_borrow_mut() -> PyResult<()> {

   179     with_setup(|_py, owner| {

   180         let cell = &owner.borrow().string;

   181         let shared_ref = unsafe { cell.borrow(owner) };

   182         let _mut_ref = shared_ref.borrow_mut();

   183 

   184         assert!(try_leak_string(owner).is_err());

   185         Ok(())

   186     })

   187 }

   188 

   189 #[test]

   190 #[should_panic(expected = "already mutably borrowed")]

   191 fn test_leak_while_borrow_mut() {

   192     with_setup(|_py, owner| {

   193         let cell = &owner.borrow().string;

   194         let shared_ref = unsafe { cell.borrow(owner) };

   195         let _mut_ref = shared_ref.borrow_mut();

   196 

   197         leak_string(owner);

   198         Ok(())

   199     })

   200     .expect("should already have panicked")

   201 }

   202 

   203 #[test]

   204 fn test_try_borrow_mut_while_borrow() -> PyResult<()> {

   205     with_setup(|_py, owner| {

   206         let cell = &owner.borrow().string;

   207         let shared_ref = unsafe { cell.borrow(owner) };

   208         let _ref = shared_ref.borrow();

   209 

   210         assert_try_borrow_string_mut_err(owner);

   211         Ok(())

   212     })

   213 }

   214 

   215 #[test]

   216 #[should_panic(expected = "already borrowed")]

   217 fn test_borrow_mut_while_borrow() {

   218     with_setup(|_py, owner| {

   219         let cell = &owner.borrow().string;

   220         let shared_ref = unsafe { cell.borrow(owner) };

   221         let _ref = shared_ref.borrow();

   222 

   223         let shared_ref2 = unsafe { cell.borrow(owner) };

   224         let _mut_ref = shared_ref2.borrow_mut();

   225         Ok(())

   226     })

   227     .expect("should already have panicked")

   228 }

   229 

   230 #[test]

   231 fn test_try_borrow_while_borrow_mut() -> PyResult<()> {

   232     with_setup(|_py, owner| {

   233         let cell = &owner.borrow().string;

   234         let shared_ref = unsafe { cell.borrow(owner) };

   235         let _mut_ref = shared_ref.borrow_mut();

   236 

   237         assert_try_borrow_string_err(owner);

   238         Ok(())

   239     })

   240 }

   241 

   242 #[test]

   243 #[should_panic(expected = "already mutably borrowed")]

   244 fn test_borrow_while_borrow_mut() {

   245     with_setup(|_py, owner| {

   246         let cell = &owner.borrow().string;

   247         let shared_ref = unsafe { cell.borrow(owner) };

   248         let _mut_ref = shared_ref.borrow_mut();

   249 

   250         let shared_ref2 = unsafe { cell.borrow(owner) };

   251         let _ref = shared_ref2.borrow();

   252         Ok(())

   253     })

   254     .expect("should already have panicked")

   255 }