From bfb6111e0a8ff7490fdc9a152565656dd3754ad9 Mon Sep 17 00:00:00 2001
From: kyren <kerriganw@gmail.com>
Date: Sat, 16 Dec 2017 17:44:13 -0500
Subject: [PATCH] API for registry access via string keys only (for now)

Also includes some fixes for stack usage and changes an assert_eq to lua_assert
---
 src/lua.rs   | 46 ++++++++++++++++++++++++++++++++++++++++------
 src/tests.rs | 13 +++++++++++++
 src/util.rs  |  2 ++
 3 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/src/lua.rs b/src/lua.rs
index 1a9d8ad..74bc76b 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -135,7 +135,7 @@ impl Lua {
     pub fn create_string(&self, s: &str) -> Result<String> {
         unsafe {
             stack_err_guard(self.state, 0, || {
-                check_stack(self.state, 2);
+                check_stack(self.state, 4);
                 push_string(self.state, s)?;
                 Ok(String(self.pop_ref(self.state)))
             })
@@ -405,7 +405,41 @@ impl Lua {
         T::from_lua_multi(value, self)
     }
 
-    // Used 1 stack space, does not call checkstack
+    /// Set a value in the Lua registry based on a string key.
+    ///
+    /// This value will be available to rust from all `Lua` instances which share the same main
+    /// state.
+    pub fn set_registry<'lua, T: ToLua<'lua>>(&'lua self, registry_key: &str, t: T) -> Result<()> {
+        unsafe {
+            stack_err_guard(self.state, 0, || {
+                check_stack(self.state, 5);
+                push_string(self.state, registry_key)?;
+                self.push_value(self.state, t.to_lua(self)?);
+                protect_lua_call(self.state, 2, 0, |state| {
+                    ffi::lua_settable(state, ffi::LUA_REGISTRYINDEX);
+                })
+            })
+        }
+    }
+
+    /// Get a value from the Lua registry based on a string key.
+    ///
+    /// Any Lua instance which shares the underlying main state may call `get_registry` to get a
+    /// value previously set by `set_registry`.
+    pub fn get_registry<'lua, T: FromLua<'lua>>(&'lua self, registry_key: &str) -> Result<T> {
+        unsafe {
+            stack_err_guard(self.state, 0, || {
+                check_stack(self.state, 4);
+                push_string(self.state, registry_key)?;
+                protect_lua_call(self.state, 1, 1, |state| {
+                    ffi::lua_gettable(state, ffi::LUA_REGISTRYINDEX)
+                })?;
+                T::from_lua(self.pop_value(self.state), self)
+            })
+        }
+    }
+
+    // Uses 1 stack space, does not call checkstack
     pub(crate) unsafe fn push_value(&self, state: *mut ffi::lua_State, value: Value) {
         match value {
             Value::Nil => {
@@ -454,7 +488,7 @@ impl Lua {
         }
     }
 
-    // Used 1 stack space, does not call checkstack
+    // Uses 1 stack space, does not call checkstack
     pub(crate) unsafe fn pop_value(&self, state: *mut ffi::lua_State) -> Value {
         match ffi::lua_type(state, -1) {
             ffi::LUA_TNIL => {
@@ -510,9 +544,9 @@ impl Lua {
 
     // Used 1 stack space, does not call checkstack
     pub(crate) unsafe fn push_ref(&self, state: *mut ffi::lua_State, lref: &LuaRef) {
-        assert_eq!(
-            lref.lua.main_state,
-            self.main_state,
+        lua_assert!(
+            state,
+            lref.lua.main_state == self.main_state,
             "Lua instance passed Value created from a different Lua"
         );
 
diff --git a/src/tests.rs b/src/tests.rs
index 418c3d1..1bbb6e1 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -499,6 +499,19 @@ fn test_gc_error() {
     }
 }
 
+#[test]
+fn test_registry() {
+    let lua = Lua::new();
+
+    lua.set_registry::<i32>("test", 42).unwrap();
+    let f = lua.create_function(move |lua, ()| {
+        assert_eq!(lua.get_registry::<i32>("test")?, 42);
+        Ok(())
+    }).unwrap();
+
+    f.call::<_, ()>(()).unwrap();
+}
+
 // TODO: Need to use compiletest-rs or similar to make sure these don't compile.
 /*
 #[test]
diff --git a/src/util.rs b/src/util.rs
index db594da..aa203fc 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -211,12 +211,14 @@ pub unsafe fn pop_error(state: *mut ffi::lua_State, err_code: c_int) -> Error {
     }
 }
 
+// Internally uses 4 stack spaces, does not call checkstack
 pub unsafe fn push_string(state: *mut ffi::lua_State, s: &str) -> Result<()> {
     protect_lua_call(state, 0, 1, |state| {
         ffi::lua_pushlstring(state, s.as_ptr() as *const c_char, s.len());
     })
 }
 
+// Internally uses 4 stack spaces, does not call checkstack
 pub unsafe fn push_userdata<T>(state: *mut ffi::lua_State, t: T) -> Result<()> {
     let mut t = Some(t);
     protect_lua_call(state, 0, 1, |state| {