From 526e7418d820f538b3e0e0c09240511393e7837b Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sat, 9 May 2020 11:24:31 +0100
Subject: [PATCH] Add MetaMethod::Close to support Lua 5.4 to-be-closed
 variables

---
 src/userdata.rs   | 16 ++++++++++++++++
 src/util.rs       | 11 +++++++++++
 tests/userdata.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/src/userdata.rs b/src/userdata.rs
index 469ad64..5fcf3eb 100644
--- a/src/userdata.rs
+++ b/src/userdata.rs
@@ -80,6 +80,16 @@ pub enum MetaMethod {
     ///
     /// This is not an operator, but it will be called by the built-in `pairs` function.
     Pairs,
+    /// The `__close` metamethod.
+    ///
+    /// Executed when a variable, that marked as to-be-closed, goes out of scope.
+    ///
+    /// More information about to-be-closed variabled can be found in the Lua 5.4
+    /// [documentation][lua_doc].
+    ///
+    /// [lua_doc]: https://www.lua.org/manual/5.4/manual.html#3.3.8
+    #[cfg(feature = "lua54")]
+    Close,
 }
 
 impl MetaMethod {
@@ -92,6 +102,7 @@ impl MetaMethod {
             MetaMethod::Mod => b"__mod",
             MetaMethod::Pow => b"__pow",
             MetaMethod::Unm => b"__unm",
+
             #[cfg(any(feature = "lua54", feature = "lua53"))]
             MetaMethod::IDiv => b"__idiv",
             #[cfg(any(feature = "lua54", feature = "lua53"))]
@@ -106,6 +117,7 @@ impl MetaMethod {
             MetaMethod::Shl => b"__shl",
             #[cfg(any(feature = "lua54", feature = "lua53"))]
             MetaMethod::Shr => b"__shr",
+
             MetaMethod::Concat => b"__concat",
             MetaMethod::Len => b"__len",
             MetaMethod::Eq => b"__eq",
@@ -115,8 +127,12 @@ impl MetaMethod {
             MetaMethod::NewIndex => b"__newindex",
             MetaMethod::Call => b"__call",
             MetaMethod::ToString => b"__tostring",
+
             #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
             MetaMethod::Pairs => b"__pairs",
+
+            #[cfg(feature = "lua54")]
+            MetaMethod::Close => b"__close",
         }
     }
 }
diff --git a/src/util.rs b/src/util.rs
index 1d21cdd..1d98fb7 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -661,12 +661,19 @@ pub unsafe fn init_error_registry(state: *mut ffi::lua_State) {
         cstr!("__mod"),
         cstr!("__pow"),
         cstr!("__unm"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__idiv"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__band"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__bor"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__bxor"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__bnot"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__shl"),
+        #[cfg(any(feature = "lua54", feature = "lua53"))]
         cstr!("__shr"),
         cstr!("__concat"),
         cstr!("__len"),
@@ -677,8 +684,12 @@ pub unsafe fn init_error_registry(state: *mut ffi::lua_State) {
         cstr!("__newindex"),
         cstr!("__call"),
         cstr!("__tostring"),
+        #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
         cstr!("__pairs"),
+        #[cfg(any(feature = "lua53", feature = "lua52"))]
         cstr!("__ipairs"),
+        #[cfg(feature = "lua54")]
+        cstr!("__close"),
     ] {
         ffi::lua_pushstring(state, method);
         ffi::lua_pushcfunction(state, destructed_error);
diff --git a/tests/userdata.rs b/tests/userdata.rs
index fa40a00..8afc795 100644
--- a/tests/userdata.rs
+++ b/tests/userdata.rs
@@ -1,5 +1,8 @@
 use std::sync::Arc;
 
+#[cfg(feature = "lua54")]
+use std::sync::atomic::{AtomicI64, Ordering};
+
 use mlua::{
     AnyUserData, ExternalError, Function, Lua, MetaMethod, Result, String, UserData,
     UserDataMethods, Value,
@@ -154,6 +157,52 @@ fn test_metamethods() -> Result<()> {
     Ok(())
 }
 
+#[test]
+#[cfg(feature = "lua54")]
+fn test_metamethod_close() -> Result<()> {
+    #[derive(Clone)]
+    struct MyUserData(Arc<AtomicI64>);
+
+    impl UserData for MyUserData {
+        fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
+            methods.add_method("get", |_, data, ()| Ok(data.0.load(Ordering::Relaxed)));
+            methods.add_meta_method(MetaMethod::Close, |_, data, _err: Value| {
+                data.0.store(0, Ordering::Relaxed);
+                Ok(())
+            });
+        }
+    }
+
+    let lua = Lua::new();
+    let globals = lua.globals();
+
+    let ud = MyUserData(Arc::new(AtomicI64::new(-1)));
+    let ud2 = ud.clone();
+
+    globals.set(
+        "new_userdata",
+        lua.create_function(move |_lua, val: i64| {
+            let ud = ud2.clone();
+            ud.0.store(val, Ordering::Relaxed);
+            Ok(ud)
+        })?,
+    )?;
+
+    lua.load(
+        r#"
+        do
+            local ud <close> = new_userdata(7)
+            assert(ud:get() == 7)
+        end
+    "#,
+    )
+    .exec()?;
+
+    assert_eq!(ud.0.load(Ordering::Relaxed), 0);
+
+    Ok(())
+}
+
 #[test]
 fn test_gc_userdata() -> Result<()> {
     struct MyUserdata {