From 4d3ac6d8c5bc2fbc916959fd18fc39e9e8f5c3b4 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Tue, 21 Sep 2021 00:26:48 +0100
Subject: [PATCH] Add new "application data" api

---
 src/lua.rs     | 74 +++++++++++++++++++++++++++++++++++++++++++++++++-
 tests/tests.rs | 31 +++++++++++++++++++++
 2 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/src/lua.rs b/src/lua.rs
index 6fbb6df..716afee 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -1,5 +1,6 @@
-use std::any::TypeId;
+use std::any::{Any, TypeId};
 use std::cell::{Ref, RefCell, RefMut, UnsafeCell};
+use std::collections::HashMap;
 use std::ffi::CString;
 use std::fmt;
 use std::marker::PhantomData;
@@ -68,6 +69,8 @@ struct ExtraData {
     registered_userdata_mt: FxHashMap<*const c_void, Option<TypeId>>,
     registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
 
+    app_data: RefCell<HashMap<TypeId, Box<dyn Any>>>,
+
     libs: StdLib,
     mem_info: Option<Box<MemoryInfo>>,
     safe: bool, // Same as in the Lua struct
@@ -460,6 +463,7 @@ impl Lua {
             registered_userdata: FxHashMap::default(),
             registered_userdata_mt: FxHashMap::default(),
             registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))),
+            app_data: RefCell::new(HashMap::new()),
             ref_thread,
             libs: StdLib::NONE,
             mem_info: None,
@@ -1578,6 +1582,74 @@ impl Lua {
         }
     }
 
+    /// Sets or replaces an application data object of type `T`.
+    ///
+    /// Application data could be accessed at any time by using [`Lua::app_data_ref()`] or [`Lua::app_data_mut()`]
+    /// methods where `T` is the data type.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use mlua::{Lua, Result};
+    ///
+    /// fn hello(lua: &Lua, _: ()) -> Result<()> {
+    ///     let mut s = lua.app_data_mut::<&str>().unwrap();
+    ///     assert_eq!(*s, "hello");
+    ///     *s = "world";
+    ///     Ok(())
+    /// }
+    ///
+    /// fn main() -> Result<()> {
+    ///     let lua = Lua::new();
+    ///     lua.set_app_data("hello");
+    ///     lua.create_function(hello)?.call(())?;
+    ///     let s = lua.app_data_ref::<&str>().unwrap();
+    ///     assert_eq!(*s, "world");
+    ///     Ok(())
+    /// }
+    /// ```
+    pub fn set_app_data<T: 'static + MaybeSend>(&self, data: T) {
+        let extra = unsafe { &mut (*self.extra.get()) };
+        extra
+            .app_data
+            .try_borrow_mut()
+            .expect("cannot borrow mutably app data container")
+            .insert(TypeId::of::<T>(), Box::new(data));
+    }
+
+    /// Gets a reference to an application data object stored by [`Lua::set_app_data()`] of type `T`.
+    pub fn app_data_ref<T: 'static>(&self) -> Option<Ref<T>> {
+        let extra = unsafe { &(*self.extra.get()) };
+        let app_data = extra
+            .app_data
+            .try_borrow()
+            .expect("cannot borrow app data container");
+        let value = app_data.get(&TypeId::of::<T>())?.downcast_ref::<T>()? as *const _;
+        Some(Ref::map(app_data, |_| unsafe { &*value }))
+    }
+
+    /// Gets a mutable reference to an application data object stored by [`Lua::set_app_data()`] of type `T`.
+    pub fn app_data_mut<T: 'static>(&self) -> Option<RefMut<T>> {
+        let extra = unsafe { &(*self.extra.get()) };
+        let mut app_data = extra
+            .app_data
+            .try_borrow_mut()
+            .expect("cannot mutably borrow app data container");
+        let value = app_data.get_mut(&TypeId::of::<T>())?.downcast_mut::<T>()? as *mut _;
+        Some(RefMut::map(app_data, |_| unsafe { &mut *value }))
+    }
+
+    /// Removes an application data of type `T`.
+    pub fn remove_app_data<T: 'static>(&self) -> Option<T> {
+        let extra = unsafe { &mut (*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().ok().map(|data| *data))
+    }
+
     // Uses 2 stack spaces, does not call checkstack
     pub(crate) unsafe fn push_value(&self, value: Value) -> Result<()> {
         match value {
diff --git a/tests/tests.rs b/tests/tests.rs
index 9b450ae..7973cb0 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -847,6 +847,37 @@ fn test_mismatched_registry_key() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_application_data() -> Result<()> {
+    let lua = Lua::new();
+
+    lua.set_app_data("test1");
+    lua.set_app_data(vec!["test2"]);
+
+    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");
+        Ok(())
+    })?;
+    f.call(())?;
+
+    assert_eq!(*lua.app_data_ref::<&str>().unwrap(), "test1");
+    assert_eq!(
+        *lua.app_data_ref::<Vec<&str>>().unwrap(),
+        vec!["test2", "test3"]
+    );
+
+    lua.remove_app_data::<Vec<&str>>();
+    assert!(matches!(lua.app_data_ref::<Vec<&str>>(), None));
+
+    Ok(())
+}
+
 #[test]
 fn test_recursion() -> Result<()> {
     let lua = Lua::new();