Add Scope::add_destructor to attach custom destructors

khvzak · khvzak · commit 5b8681dcf2d9 · 2024-10-31T14:35:21.000Z
diff --git a/src/scope.rs b/src/scope.rs
@@ -19,7 +19,9 @@ use crate::util::{
 /// See [`Lua::scope`] for more details.
 pub struct Scope<'scope, 'env: 'scope> {
     lua: LuaGuard,
+    // Internal destructors run first, then user destructors (based on the declaration order)
     destructors: Destructors<'env>,
+    user_destructors: UserDestructors<'env>,
     _scope_invariant: PhantomData<&'scope mut &'scope ()>,
     _env_invariant: PhantomData<&'env mut &'env ()>,
 }
@@ -29,11 +31,14 @@ type DestructorCallback<'a> = Box<dyn FnOnce(&RawLua, ValueRef) -> Vec<Box<dyn F
 // Implement Drop on Destructors instead of Scope to avoid compilation error
 struct Destructors<'a>(RefCell<Vec<(ValueRef, DestructorCallback<'a>)>>);
 
+struct UserDestructors<'a>(RefCell<Vec<Box<dyn FnOnce() + 'a>>>);
+
 impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
     pub(crate) fn new(lua: LuaGuard) -> Self {
         Scope {
             lua,
             destructors: Destructors(RefCell::new(Vec::new())),
+            user_destructors: UserDestructors(RefCell::new(Vec::new())),
             _scope_invariant: PhantomData,
             _env_invariant: PhantomData,
         }
@@ -215,6 +220,31 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
         }
     }
 
+    /// Adds a destructor function to be run when the scope ends.
+    ///
+    /// This functionality is useful for cleaning up any resources after the scope ends.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use mlua::{Error, Lua, Result};
+    /// # fn main() -> Result<()> {
+    /// let lua = Lua::new();
+    /// let ud = lua.create_any_userdata(String::from("hello"))?;
+    /// lua.scope(|scope| {
+    ///     scope.add_destructor(|| {
+    ///         _ = ud.take::<String>();
+    ///     });
+    ///     // Run the code that uses `ud` here
+    ///    Ok(())
+    /// })?;
+    /// assert!(matches!(ud.borrow::<String>(), Err(Error::UserDataDestructed)));
+    /// # Ok(())
+    /// # }
+    pub fn add_destructor(&'scope self, destructor: impl FnOnce() + 'env) {
+        self.user_destructors.0.borrow_mut().push(Box::new(destructor));
+    }
+
     unsafe fn create_callback(&'scope self, f: ScopedCallback<'scope>) -> Result<Function> {
         let f = mem::transmute::<ScopedCallback, Callback>(f);
         let f = self.lua.create_callback(f)?;
@@ -271,3 +301,12 @@ impl Drop for Destructors<'_> {
         }
     }
 }
+
+impl Drop for UserDestructors<'_> {
+    fn drop(&mut self) {
+        let destructors = mem::take(&mut *self.0.borrow_mut());
+        for destructor in destructors {
+            destructor();
+        }
+    }
+}
diff --git a/src/state.rs b/src/state.rs
@@ -1470,9 +1470,9 @@ impl Lua {
     /// lifetimes only outlive the scope lifetime.
     pub fn scope<'env, R>(
         &self,
-        f: impl for<'scope> FnOnce(&'scope mut Scope<'scope, 'env>) -> Result<R>,
+        f: impl for<'scope> FnOnce(&'scope Scope<'scope, 'env>) -> Result<R>,
     ) -> Result<R> {
-        f(&mut Scope::new(self.lock_arc()))
+        f(&Scope::new(self.lock_arc()))
     }
 
     /// Attempts to coerce a Lua value into a String in a manner consistent with Lua's internal
diff --git a/tests/scope.rs b/tests/scope.rs
@@ -1,6 +1,7 @@
 use std::cell::Cell;
 use std::rc::Rc;
 use std::string::String as StdString;
+use std::sync::Arc;
 
 use mlua::{
     AnyUserData, Error, Function, Lua, MetaMethod, ObjectLike, Result, String, UserData, UserDataFields,
@@ -66,6 +67,27 @@ fn test_scope_outer_lua_access() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_scope_capture_scope() -> Result<()> {
+    let lua = Lua::new();
+
+    let i = Cell::new(0);
+    lua.scope(|scope| {
+        let f = scope.create_function(|_, ()| {
+            scope.create_function(|_, n: u32| {
+                i.set(i.get() + n);
+                Ok(())
+            })
+        })?;
+        f.call::<Function>(())?.call::<()>(10)?;
+        Ok(())
+    })?;
+
+    assert_eq!(i.get(), 10);
+
+    Ok(())
+}
+
 #[test]
 fn test_scope_userdata_fields() -> Result<()> {
     struct MyUserData<'a>(&'a Cell<i64>);
@@ -463,6 +485,42 @@ fn test_scope_any_userdata_ref_mut() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_scope_destructors() -> Result<()> {
+    let lua = Lua::new();
+
+    lua.register_userdata_type::<Arc<StdString>>(|reg| {
+        reg.add_meta_method("__tostring", |_, data, ()| Ok(data.to_string()));
+    })?;
+
+    let arc_str = Arc::new(StdString::from("foo"));
+
+    let ud = lua.create_any_userdata(arc_str.clone())?;
+    lua.scope(|scope| {
+        scope.add_destructor(|| {
+            assert!(ud.take::<Arc<StdString>>().is_ok());
+        });
+        Ok(())
+    })?;
+    assert_eq!(Arc::strong_count(&arc_str), 1);
+
+    // Try destructing the userdata while it's borrowed
+    let ud = lua.create_any_userdata(arc_str.clone())?;
+    ud.borrow_scoped::<Arc<StdString>, _>(|arc_str| {
+        assert_eq!(arc_str.as_str(), "foo");
+        lua.scope(|scope| {
+            scope.add_destructor(|| {
+                assert!(ud.take::<Arc<StdString>>().is_err());
+            });
+            Ok(())
+        })
+        .unwrap();
+        assert_eq!(arc_str.as_str(), "foo");
+    })?;
+
+    Ok(())
+}
+
 fn modify_userdata(lua: &Lua, ud: &AnyUserData) -> Result<()> {
     lua.load(
         r#"
