Refactor application data container.

Now it's allowed at the same time mutably and immutably borrow different types.
Each value in the application data container is stored in it's own `RefCell` wrapper.
Also added new function `Lua::try_set_app_data()`.

Author: khvzak

src/chunk.rs

@@ -454,7 +454,7 @@ impl<'lua, 'a> Chunk<'lua, 'a> {
                     } else {
                         let mut cache = ChunksCache(HashMap::new());
                         cache.0.insert(text_source, binary_source.as_ref().to_vec());
-                        self.lua.set_app_data(cache);
+                        let _ = self.lua.try_set_app_data(cache);
                     }
                 }
             }

src/lib.rs

@@ -116,7 +116,7 @@ pub use crate::stdlib::StdLib;
 pub use crate::string::String;
 pub use crate::table::{Table, TableExt, TablePairs, TableSequence};
 pub use crate::thread::{Thread, ThreadStatus};
-pub use crate::types::{Integer, LightUserData, Number, RegistryKey};
+pub use crate::types::{AppDataRef, AppDataRefMut, Integer, LightUserData, Number, RegistryKey};
 pub use crate::userdata::{
     AnyUserData, MetaMethod, UserData, UserDataFields, UserDataMetatable, UserDataMethods,
     UserDataRef, UserDataRefMut,

src/lua.rs

@@ -1,5 +1,5 @@
-use std::any::{Any, TypeId};
-use std::cell::{Ref, RefCell, RefMut, UnsafeCell};
+use std::any::TypeId;
+use std::cell::{RefCell, UnsafeCell};
 use std::ffi::{CStr, CString};
 use std::fmt;
 use std::marker::PhantomData;
@@ -8,6 +8,7 @@ use std::ops::Deref;
 use std::os::raw::{c_char, c_int, c_void};
 use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe, Location};
 use std::ptr::NonNull;
+use std::result::Result as StdResult;
 use std::sync::atomic::{AtomicPtr, Ordering};
 use std::sync::{Arc, Mutex};
 use std::{mem, ptr, str};
@@ -25,8 +26,8 @@ use crate::string::String;
 use crate::table::Table;
 use crate::thread::Thread;
 use crate::types::{
-    Callback, CallbackUpvalue, DestructedUserdata, Integer, LightUserData, LuaRef, MaybeSend,
-    Number, RegistryKey,
+    AppData, AppDataRef, AppDataRefMut, Callback, CallbackUpvalue, DestructedUserdata, Integer,
+    LightUserData, LuaRef, MaybeSend, Number, RegistryKey,
 };
 use crate::userdata::{AnyUserData, MetaMethod, UserData, UserDataCell};
 use crate::userdata_impl::{UserDataProxy, UserDataRegistrar};
@@ -85,10 +86,8 @@ pub(crate) struct ExtraData {
     // When Lua instance dropped, setting `None` would prevent collecting `RegistryKey`s
     registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
 
-    #[cfg(not(feature = "send"))]
-    app_data: RefCell<FxHashMap<TypeId, Box<dyn Any>>>,
-    #[cfg(feature = "send")]
-    app_data: RefCell<FxHashMap<TypeId, Box<dyn Any + Send>>>,
+    // Container to store arbitrary data (extensions)
+    app_data: AppData,
 
     safe: bool,
     libs: StdLib,
@@ -508,7 +507,7 @@ impl Lua {
             registered_userdata_mt: FxHashMap::default(),
             last_checked_userdata_mt: (ptr::null(), None),
             registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))),
-            app_data: RefCell::new(FxHashMap::default()),
+            app_data: AppData::default(),
             safe: false,
             libs: StdLib::NONE,
             mem_state: None,
@@ -2222,51 +2221,45 @@ impl Lua {
     /// }
     /// ```
     #[track_caller]
-    pub fn set_app_data<T: 'static + MaybeSend>(&self, data: T) -> Option<T> {
+    pub fn set_app_data<T: MaybeSend + 'static>(&self, data: T) -> Option<T> {
+        let extra = unsafe { &*self.extra.get() };
+        extra.app_data.insert(data)
+    }
+
+    /// Tries to set or replace an application data object of type `T`.
+    ///
+    /// Returns:
+    /// - `Ok(Some(old_data))` if the data object of type `T` was successfully replaced.
+    /// - `Ok(None)` if the data object of type `T` was successfully inserted.
+    /// - `Err(data)` if the data object of type `T` was not inserted because the container is currently borrowed.
+    ///
+    /// See [`Lua::set_app_data()`] for examples.
+    pub fn try_set_app_data<T: MaybeSend + 'static>(&self, data: T) -> StdResult<Option<T>, T> {
         let extra = unsafe { &*self.extra.get() };
-        extra
-            .app_data
-            .try_borrow_mut()
-            .expect("cannot borrow mutably app data container")
-            .insert(TypeId::of::<T>(), Box::new(data))
-            .and_then(|data| data.downcast::<T>().ok().map(|data| *data))
+        extra.app_data.try_insert(data)
     }
 
     /// Gets a reference to an application data object stored by [`Lua::set_app_data()`] of type `T`.
     ///
     /// # Panics
     ///
-    /// Panics if the app data container is currently mutably borrowed. Multiple immutable reads can be
-    /// taken out at the same time.
+    /// Panics if the data object of type `T` is currently mutably borrowed. Multiple immutable reads
+    /// can be taken out at the same time.
     #[track_caller]
-    pub fn app_data_ref<T: 'static>(&self) -> Option<Ref<T>> {
+    pub fn app_data_ref<T: 'static>(&self) -> Option<AppDataRef<T>> {
         let extra = unsafe { &*self.extra.get() };
-        let app_data = extra
-            .app_data
-            .try_borrow()
-            .expect("cannot borrow app data container");
-        Ref::filter_map(app_data, |data| {
-            data.get(&TypeId::of::<T>())?.downcast_ref::<T>()
-        })
-        .ok()
+        extra.app_data.borrow()
     }
 
     /// Gets a mutable reference to an application data object stored by [`Lua::set_app_data()`] of type `T`.
     ///
     /// # Panics
     ///
-    /// Panics if the app data container is currently borrowed.
+    /// Panics if the data object of type `T` is currently borrowed.
     #[track_caller]
-    pub fn app_data_mut<T: 'static>(&self) -> Option<RefMut<T>> {
+    pub fn app_data_mut<T: 'static>(&self) -> Option<AppDataRefMut<T>> {
         let extra = unsafe { &*self.extra.get() };
-        let app_data = extra
-            .app_data
-            .try_borrow_mut()
-            .expect("cannot mutably borrow app data container");
-        RefMut::filter_map(app_data, |data| {
-            data.get_mut(&TypeId::of::<T>())?.downcast_mut::<T>()
-        })
-        .ok()
+        extra.app_data.borrow_mut()
     }
 
     /// Removes an application data of type `T`.
@@ -2277,12 +2270,7 @@ impl Lua {
     #[track_caller]
     pub fn remove_app_data<T: 'static>(&self) -> Option<T> {
         let extra = unsafe { &*self.extra.get() };
-        extra
-            .app_data
-            .try_borrow_mut()
-            .expect("cannot mutably borrow app data container")
-            .remove(&TypeId::of::<T>())
-            .and_then(|data| data.downcast::<T>().ok().map(|data| *data))
+        extra.app_data.remove()
     }
 
     // Uses 2 stack spaces, does not call checkstack

src/types.rs

@@ -1,13 +1,18 @@
-use std::cell::UnsafeCell;
+use std::any::{Any, TypeId};
+use std::cell::{Cell, Ref, RefCell, RefMut, UnsafeCell};
 use std::hash::{Hash, Hasher};
+use std::ops::{Deref, DerefMut};
 use std::os::raw::{c_int, c_void};
+use std::result::Result as StdResult;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 use std::{fmt, mem, ptr};
 
 #[cfg(feature = "lua54")]
 use std::ffi::CStr;
 
+use rustc_hash::FxHashMap;
+
 #[cfg(feature = "async")]
 use futures_util::future::LocalBoxFuture;
 
@@ -286,6 +291,150 @@ impl LuaOwnedRef {
     }
 }
 
+#[derive(Debug, Default)]
+pub(crate) struct AppData {
+    #[cfg(not(feature = "send"))]
+    container: UnsafeCell<FxHashMap<TypeId, RefCell<Box<dyn Any>>>>,
+    #[cfg(feature = "send")]
+    container: UnsafeCell<FxHashMap<TypeId, RefCell<Box<dyn Any + Send>>>>,
+    borrow: Cell<usize>,
+}
+
+impl AppData {
+    #[track_caller]
+    pub(crate) fn insert<T: MaybeSend + 'static>(&self, data: T) -> Option<T> {
+        match self.try_insert(data) {
+            Ok(data) => data,
+            Err(_) => panic!("cannot mutably borrow app data container"),
+        }
+    }
+
+    pub(crate) fn try_insert<T: MaybeSend + 'static>(&self, data: T) -> StdResult<Option<T>, T> {
+        if self.borrow.get() != 0 {
+            return Err(data);
+        }
+        // SAFETY: we checked that there are no other references to the container
+        Ok(unsafe { &mut *self.container.get() }
+            .insert(TypeId::of::<T>(), RefCell::new(Box::new(data)))
+            .and_then(|data| data.into_inner().downcast::<T>().ok().map(|data| *data)))
+    }
+
+    #[track_caller]
+    pub(crate) fn borrow<T: 'static>(&self) -> Option<AppDataRef<T>> {
+        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,
+        })
+    }
+
+    #[track_caller]
+    pub(crate) fn borrow_mut<T: 'static>(&self) -> Option<AppDataRefMut<T>> {
+        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,
+        })
+    }
+
+    #[track_caller]
+    pub(crate) fn remove<T: 'static>(&self) -> Option<T> {
+        if self.borrow.get() != 0 {
+            panic!("cannot mutably borrow app data container");
+        }
+        // SAFETY: we checked that there are no other references to the container
+        unsafe { &mut *self.container.get() }
+            .remove(&TypeId::of::<T>())?
+            .into_inner()
+            .downcast::<T>()
+            .ok()
+            .map(|data| *data)
+    }
+}
+
+/// A wrapper type for an immutably borrowed value from an app data container.
+///
+/// This type is similar to [`Ref`].
+pub struct AppDataRef<'a, T: ?Sized + 'a> {
+    data: Ref<'a, T>,
+    borrow: &'a Cell<usize>,
+}
+
+impl<T: ?Sized> Drop for AppDataRef<'_, T> {
+    fn drop(&mut self) {
+        self.borrow.set(self.borrow.get() - 1);
+    }
+}
+
+impl<T: ?Sized> Deref for AppDataRef<'_, T> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.data
+    }
+}
+
+impl<T: ?Sized + fmt::Display> fmt::Display for AppDataRef<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
+impl<T: ?Sized + fmt::Debug> fmt::Debug for AppDataRef<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
+/// A wrapper type for a mutably borrowed value from an app data container.
+///
+/// This type is similar to [`RefMut`].
+pub struct AppDataRefMut<'a, T: ?Sized + 'a> {
+    data: RefMut<'a, T>,
+    borrow: &'a Cell<usize>,
+}
+
+impl<T: ?Sized> Drop for AppDataRefMut<'_, T> {
+    fn drop(&mut self) {
+        self.borrow.set(self.borrow.get() - 1);
+    }
+}
+
+impl<T: ?Sized> Deref for AppDataRefMut<'_, T> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.data
+    }
+}
+
+impl<T: ?Sized> DerefMut for AppDataRefMut<'_, T> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.data
+    }
+}
+
+impl<T: ?Sized + fmt::Display> fmt::Display for AppDataRefMut<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
+impl<T: ?Sized + fmt::Debug> fmt::Debug for AppDataRefMut<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
 #[cfg(test)]
 mod assertions {
     use super::*;

tests/tests.rs

@@ -887,19 +887,47 @@ fn test_application_data() -> Result<()> {
     lua.set_app_data("test1");
     lua.set_app_data(vec!["test2"]);
 
+    // Borrow &str immutably and Vec<&str> mutably
+    let s = lua.app_data_ref::<&str>().unwrap();
+    let mut v = lua.app_data_mut::<Vec<&str>>().unwrap();
+    v.push("test3");
+
+    // Insert of new data or removal should fail now
+    assert!(lua.try_set_app_data::<i32>(123).is_err());
+    match catch_unwind(AssertUnwindSafe(|| lua.set_app_data::<i32>(123))) {
+        Ok(_) => panic!("expected panic"),
+        Err(_) => {}
+    }
+    match catch_unwind(AssertUnwindSafe(|| lua.remove_app_data::<i32>())) {
+        Ok(_) => panic!("expected panic"),
+        Err(_) => {}
+    }
+
+    // Check display and debug impls
+    assert_eq!(format!("{s}"), "test1");
+    assert_eq!(format!("{s:?}"), "\"test1\"");
+
+    // Borrowing immutably and mutably of the same type is not allowed
+    match catch_unwind(AssertUnwindSafe(|| lua.app_data_mut::<&str>().unwrap())) {
+        Ok(_) => panic!("expected panic"),
+        Err(_) => {}
+    }
+    drop((s, v));
+
+    // Test that application data is accessible from anywhere
     let f = lua.create_function(|lua, ()| {
-        {
-            let data1 = lua.app_data_ref::<&str>().unwrap();
-            assert_eq!(*data1, "test1");
-        }
-        let mut data2 = lua.app_data_mut::<Vec<&str>>().unwrap();
-        assert_eq!(*data2, vec!["test2"]);
-        data2.push("test3");
+        let mut data1 = lua.app_data_mut::<&str>().unwrap();
+        assert_eq!(*data1, "test1");
+        *data1 = "test4";
+
+        let data2 = lua.app_data_ref::<Vec<&str>>().unwrap();
+        assert_eq!(*data2, vec!["test2", "test3"]);
+
         Ok(())
     })?;
     f.call(())?;
 
-    assert_eq!(*lua.app_data_ref::<&str>().unwrap(), "test1");
+    assert_eq!(*lua.app_data_ref::<&str>().unwrap(), "test4");
     assert_eq!(
         *lua.app_data_ref::<Vec<&str>>().unwrap(),
         vec!["test2", "test3"]