Add Lua::try_app_data_ref and Lua::try_app_data_mut

mlua-rs · Oct 30, 2024 · a8d5f23 · a8d5f23
1 parent 5ec4e03
commit a8d5f23
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 18 deletions.
diff --git a/src/state.rs b/src/state.rs
@@ -1,5 +1,5 @@
 use std::any::TypeId;
-use std::cell::RefCell;
+use std::cell::{BorrowError, BorrowMutError, RefCell};
 use std::marker::PhantomData;
 use std::ops::Deref;
 use std::os::raw::c_int;
@@ -1854,6 +1854,14 @@ impl Lua {
         extra.app_data.borrow(Some(guard))
     }
 
+    /// Tries to get a reference to an application data object stored by [`Lua::set_app_data`] of
+    /// type `T`.
+    pub fn try_app_data_ref<T: 'static>(&self) -> StdResult<Option<AppDataRef<T>>, BorrowError> {
+        let guard = self.lock_arc();
+        let extra = unsafe { &*guard.extra.get() };
+        extra.app_data.try_borrow(Some(guard))
+    }
+
     /// Gets a mutable reference to an application data object stored by [`Lua::set_app_data`] of
     /// type `T`.
     ///
@@ -1867,6 +1875,14 @@ impl Lua {
         extra.app_data.borrow_mut(Some(guard))
     }
 
+    /// Tries to get a mutable reference to an application data object stored by
+    /// [`Lua::set_app_data`] of type `T`.
+    pub fn try_app_data_mut<T: 'static>(&self) -> StdResult<Option<AppDataRefMut<T>>, BorrowMutError> {
+        let guard = self.lock_arc();
+        let extra = unsafe { &*guard.extra.get() };
+        extra.app_data.try_borrow_mut(Some(guard))
+    }
+
     /// Removes an application data of type `T`.
     ///
     /// # Panics

diff --git a/src/types/app_data.rs b/src/types/app_data.rs
@@ -1,5 +1,5 @@
 use std::any::{Any, TypeId};
-use std::cell::{Cell, Ref, RefCell, RefMut, UnsafeCell};
+use std::cell::{BorrowError, BorrowMutError, Cell, Ref, RefCell, RefMut, UnsafeCell};
 use std::fmt;
 use std::ops::{Deref, DerefMut};
 use std::result::Result as StdResult;
@@ -41,30 +41,66 @@ impl AppData {
             .and_then(|data| data.into_inner().downcast::<T>().ok().map(|data| *data)))
     }
 
+    #[inline]
     #[track_caller]
     pub(crate) fn borrow<T: 'static>(&self, guard: Option<LuaGuard>) -> Option<AppDataRef<T>> {
+        match self.try_borrow(guard) {
+            Ok(data) => data,
+            Err(err) => panic!("already mutably borrowed: {err:?}"),
+        }
+    }
+
+    pub(crate) fn try_borrow<T: 'static>(
+        &self,
+        guard: Option<LuaGuard>,
+    ) -> Result<Option<AppDataRef<T>>, BorrowError> {
         let data = unsafe { &*self.container.get() }
-            .get(&TypeId::of::<T>())?
-            .borrow();
-        self.borrow.set(self.borrow.get() + 1);
-        Some(AppDataRef {
-            data: Ref::filter_map(data, |data| data.downcast_ref()).ok()?,
-            borrow: &self.borrow,
-            _guard: guard,
-        })
+            .get(&TypeId::of::<T>())
+            .map(|c| c.try_borrow())
+            .transpose()?
+            .and_then(|data| Ref::filter_map(data, |data| data.downcast_ref()).ok());
+        match data {
+            Some(data) => {
+                self.borrow.set(self.borrow.get() + 1);
+                Ok(Some(AppDataRef {
+                    data,
+                    borrow: &self.borrow,
+                    _guard: guard,
+                }))
+            }
+            None => Ok(None),
+        }
     }
 
+    #[inline]
     #[track_caller]
     pub(crate) fn borrow_mut<T: 'static>(&self, guard: Option<LuaGuard>) -> Option<AppDataRefMut<T>> {
+        match self.try_borrow_mut(guard) {
+            Ok(data) => data,
+            Err(err) => panic!("already borrowed: {err:?}"),
+        }
+    }
+
+    pub(crate) fn try_borrow_mut<T: 'static>(
+        &self,
+        guard: Option<LuaGuard>,
+    ) -> Result<Option<AppDataRefMut<T>>, BorrowMutError> {
         let data = unsafe { &*self.container.get() }
-            .get(&TypeId::of::<T>())?
-            .borrow_mut();
-        self.borrow.set(self.borrow.get() + 1);
-        Some(AppDataRefMut {
-            data: RefMut::filter_map(data, |data| data.downcast_mut()).ok()?,
-            borrow: &self.borrow,
-            _guard: guard,
-        })
+            .get(&TypeId::of::<T>())
+            .map(|c| c.try_borrow_mut())
+            .transpose()?
+            .and_then(|data| RefMut::filter_map(data, |data| data.downcast_mut()).ok());
+        match data {
+            Some(data) => {
+                self.borrow.set(self.borrow.get() + 1);
+                Ok(Some(AppDataRefMut {
+                    data,
+                    borrow: &self.borrow,
+                    _guard: guard,
+                }))
+            }
+            None => Ok(None),
+        }
     }
 
     #[track_caller]

diff --git a/tests/tests.rs b/tests/tests.rs
@@ -844,10 +844,12 @@ fn test_application_data() -> Result<()> {
     assert_eq!(format!("{s:?}"), "\"test1\"");
 
     // Borrowing immutably and mutably of the same type is not allowed
+    assert!(lua.try_app_data_mut::<&str>().is_err());
     match catch_unwind(AssertUnwindSafe(|| lua.app_data_mut::<&str>().unwrap())) {
         Ok(_) => panic!("expected panic"),
         Err(_) => {}
     }
+    assert!(lua.try_app_data_ref::<Vec<&str>>().is_err());
     drop((s, v));
 
     // Test that application data is accessible from anywhere