From 727f99ee4d80825c03177fc99f0203615093ae6d Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sat, 20 Jan 2024 21:39:42 +0000
Subject: [PATCH] Implement IntoLua for `&RegistryKey` This would allow just
 passing registry keys to arguments with fasttrack to push directly into
 stack.

---
 src/conversion.rs | 22 +++++++++++++++++++++-
 tests/tests.rs    | 33 +++++++++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/src/conversion.rs b/src/conversion.rs
index 41fdf365..1321b421 100644
--- a/src/conversion.rs
+++ b/src/conversion.rs
@@ -16,7 +16,7 @@ use crate::lua::Lua;
 use crate::string::String;
 use crate::table::Table;
 use crate::thread::Thread;
-use crate::types::{LightUserData, MaybeSend};
+use crate::types::{LightUserData, MaybeSend, RegistryKey};
 use crate::userdata::{AnyUserData, UserData, UserDataRef, UserDataRefMut};
 use crate::value::{FromLua, IntoLua, Nil, Value};
 
@@ -240,6 +240,26 @@ impl<'lua> FromLua<'lua> for Error {
     }
 }
 
+impl<'lua> IntoLua<'lua> for &RegistryKey {
+    #[inline]
+    fn into_lua(self, lua: &'lua Lua) -> Result<Value<'lua>> {
+        lua.registry_value(self)
+    }
+
+    unsafe fn push_into_stack(self, lua: &'lua Lua) -> Result<()> {
+        if !lua.owns_registry_value(self) {
+            return Err(Error::MismatchedRegistryKey);
+        }
+
+        if self.is_nil() {
+            ffi::lua_pushnil(lua.state());
+        } else {
+            ffi::lua_rawgeti(lua.state(), ffi::LUA_REGISTRYINDEX, self.registry_id as _);
+        }
+        Ok(())
+    }
+}
+
 impl<'lua> IntoLua<'lua> for bool {
     #[inline]
     fn into_lua(self, _: &'lua Lua) -> Result<Value<'lua>> {
diff --git a/tests/tests.rs b/tests/tests.rs
index 2bfefe19..55f283df 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -8,8 +8,8 @@ use std::sync::Arc;
 use std::{error, f32, f64, fmt};
 
 use mlua::{
-    ChunkMode, Error, ExternalError, Function, Lua, LuaOptions, Nil, Result, StdLib, String, Table,
-    UserData, Value, Variadic,
+    ChunkMode, Error, ExternalError, Function, IntoLua, Lua, LuaOptions, Nil, Result, StdLib,
+    String, Table, UserData, Value, Variadic,
 };
 
 #[cfg(not(feature = "luau"))]
@@ -779,6 +779,35 @@ fn test_registry_value() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_registry_value_into_lua() -> Result<()> {
+    let lua = Lua::new();
+
+    let t = lua.create_table()?;
+    let r = lua.create_registry_value(t)?;
+    let f = lua.create_function(|_, t: Table| t.raw_set("hello", "world"))?;
+
+    f.call(&r)?;
+    let v = r.into_lua(&lua)?;
+    let t = v.as_table().unwrap();
+    assert_eq!(t.get::<_, String>("hello")?, "world");
+
+    // Try to set nil registry key
+    let r_nil = lua.create_registry_value(Value::Nil)?;
+    t.set("hello", &r_nil)?;
+    assert_eq!(t.get::<_, Value>("hello")?, Value::Nil);
+
+    // Check non-owned registry key
+    let lua2 = Lua::new();
+    let r2 = lua2.create_registry_value("abc")?;
+    assert!(matches!(
+        f.call::<_, ()>(&r2),
+        Err(Error::MismatchedRegistryKey)
+    ));
+
+    Ok(())
+}
+
 #[test]
 fn test_drop_registry_value() -> Result<()> {
     struct MyUserdata(Arc<()>);