Add automatic Lua "user accessible registry" keys

Also, during the implementation of this, I noticed a problem with the 0.10
memory safety, which is that luaL_ref is also memory unsafe.  I attempted to
change the API to support luaL_ref potentially returning Result, but this change
will cause an enormous amount of API chaos, (just as an example, it becomes
impossible to implement Clone for LuaRef as is).  Instead, luaL_ref now is
guarded by gc_guard.

5 files changed, 148 insertions, 31 deletions

diff --git a/src/ffi.rs b/src/ffi.rs
index 707cfc9..fdbc924 100644
--- a/src/ffi.rs
+++ b/src/ffi.rs
@@ -119,6 +119,7 @@ extern "C" {
     pub fn lua_geti(state: *mut lua_State, index: c_int, i: lua_Integer) -> c_int;
     pub fn lua_rawget(state: *mut lua_State, index: c_int) -> c_int;
     pub fn lua_rawgeti(state: *mut lua_State, index: c_int, n: lua_Integer) -> c_int;
+    pub fn lua_rawseti(state: *mut lua_State, index: c_int, n: lua_Integer);
     pub fn lua_getmetatable(state: *mut lua_State, index: c_int) -> c_int;
 
     pub fn lua_createtable(state: *mut lua_State, narr: c_int, nrec: c_int);
diff --git a/src/lua.rs b/src/lua.rs
index c34c0a9..03485f2 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -14,7 +14,7 @@ use ffi;
 use error::*;
 use util::*;
 use value::{FromLua, FromLuaMulti, MultiValue, Nil, ToLua, ToLuaMulti, Value};
-use types::{Callback, Integer, LightUserData, LuaRef, Number};
+use types::{Callback, Integer, LightUserData, LuaRef, Number, RegistryKey};
 use string::String;
 use table::Table;
 use function::Function;
@@ -415,11 +415,19 @@ impl Lua {
     ) -> Result<()> {
         unsafe {
             stack_err_guard(self.state, 0, || {
-                check_stack(self.state, 5);
+                check_stack(self.state, 6);
+
+                ffi::lua_pushlightuserdata(
+                    self.state,
+                    &USER_REGISTRY_KEY as *const u8 as *mut c_void,
+                );
+                ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
+
                 push_string(self.state, name)?;
                 self.push_value(self.state, t.to_lua(self)?);
-                protect_lua_call(self.state, 2, 0, |state| {
-                    ffi::lua_settable(state, ffi::LUA_REGISTRYINDEX);
+
+                protect_lua_call(self.state, 3, 0, |state| {
+                    ffi::lua_settable(state, -3);
                 })
             })
         }
@@ -427,28 +435,109 @@ impl Lua {
 
     /// Get a value from the Lua registry based on a string name.
     ///
-    /// Any Lua instance which shares the underlying main state may call `get_registry` to get a
-    /// value previously set by `set_registry`.
+    /// Any Lua instance which shares the underlying main state may call `named_registry_value` to
+    /// get a value previously set by `set_named_registry_value`.
     pub fn named_registry_value<'lua, T: FromLua<'lua>>(&'lua self, name: &str) -> Result<T> {
         unsafe {
             stack_err_guard(self.state, 0, || {
-                check_stack(self.state, 4);
+                check_stack(self.state, 5);
+
+                ffi::lua_pushlightuserdata(
+                    self.state,
+                    &USER_REGISTRY_KEY as *const u8 as *mut c_void,
+                );
+                ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
+
                 push_string(self.state, name)?;
-                protect_lua_call(self.state, 1, 1, |state| {
-                    ffi::lua_gettable(state, ffi::LUA_REGISTRYINDEX)
-                })?;
+                protect_lua_call(self.state, 2, 1, |state| ffi::lua_gettable(state, -2))?;
+
                 T::from_lua(self.pop_value(self.state), self)
             })
         }
     }
 
-    /// Clears a named value in the Lua registry.
+    /// Removes a named value in the Lua registry.
     ///
     /// Equivalent to calling `Lua::set_named_registry_value` with a value of Nil.
     pub fn unset_named_registry_value<'lua>(&'lua self, name: &str) -> Result<()> {
         self.set_named_registry_value(name, Nil)
     }
 
+    /// Place a value in the Lua registry with an auto-generated key.
+    ///
+    /// This value will be available to rust from all `Lua` instances which share the same main
+    /// state.  The value will NOT be automatically removed when the `RegistryKey` is dropped, you
+    /// MUST call `remove_registry_value` to remove it.
+    pub fn create_registry_value<'lua, T: ToLua<'lua>>(&'lua self, t: T) -> Result<RegistryKey> {
+        unsafe {
+            stack_guard(self.state, 0, || {
+                check_stack(self.state, 3);
+
+                ffi::lua_pushlightuserdata(
+                    self.state,
+                    &USER_REGISTRY_KEY as *const u8 as *mut c_void,
+                );
+                ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
+
+                self.push_value(self.state, t.to_lua(self)?);
+                let id = gc_guard(self.state, || ffi::luaL_ref(self.state, -2));
+
+                ffi::lua_pop(self.state, 1);
+
+                Ok(RegistryKey(id))
+            })
+        }
+    }
+
+    /// Get a value from the Lua registry by its `RegistryKey`
+    ///
+    /// Any Lua instance which shares the underlying main state may call `registry_value` to get a
+    /// value previously placed by `create_registry_value`.
+    pub fn registry_value<'lua, T: FromLua<'lua>>(&'lua self, key: &RegistryKey) -> Result<T> {
+        unsafe {
+            stack_err_guard(self.state, 0, || {
+                check_stack(self.state, 2);
+
+                ffi::lua_pushlightuserdata(
+                    self.state,
+                    &USER_REGISTRY_KEY as *const u8 as *mut c_void,
+                );
+                ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
+
+                ffi::lua_rawgeti(self.state, -1, key.0 as ffi::lua_Integer);
+
+                let t = T::from_lua(self.pop_value(self.state), self)?;
+
+                ffi::lua_pop(self.state, 1);
+
+                Ok(t)
+            })
+        }
+    }
+
+    /// Removes a value from the Lua registry.
+    ///
+    /// You MUST call this function to remove a value placed in the registry with
+    /// `create_registry_value`
+    pub fn remove_registry_value<'lua>(&'lua self, key: RegistryKey) {
+        unsafe {
+            stack_guard(self.state, 0, || {
+                check_stack(self.state, 2);
+
+                ffi::lua_pushlightuserdata(
+                    self.state,
+                    &USER_REGISTRY_KEY as *const u8 as *mut c_void,
+                );
+                ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
+
+                ffi::lua_pushnil(self.state);
+                ffi::lua_rawseti(self.state, -2, key.0 as ffi::lua_Integer);
+
+                ffi::lua_pop(self.state, 1);
+            })
+        }
+    }
+
     // Uses 1 stack space, does not call checkstack
     pub(crate) unsafe fn push_value(&self, state: *mut ffi::lua_State, value: Value) {
         match value {
@@ -575,7 +664,7 @@ impl Lua {
     //
     // pop_ref uses 1 extra stack space and does not call checkstack
     pub(crate) unsafe fn pop_ref(&self, state: *mut ffi::lua_State) -> LuaRef {
-        let registry_id = ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX);
+        let registry_id = gc_guard(state, || ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX));
         LuaRef {
             lua: self,
             registry_id: registry_id,
@@ -810,8 +899,13 @@ impl Lua {
 
             ffi::lua_rawset(state, ffi::LUA_REGISTRYINDEX);
 
-            // Override pcall, xpcall, and setmetatable with versions that cannot be used to
-            // cause unsafety.
+            // Create "user registry" table
+
+            ffi::lua_pushlightuserdata(state, &USER_REGISTRY_KEY as *const u8 as *mut c_void);
+            ffi::lua_newtable(state);
+            ffi::lua_rawset(state, ffi::LUA_REGISTRYINDEX);
+
+            // Override pcall and xpcall with versions that cannot be used to catch rust panics.
 
             ffi::lua_rawgeti(state, ffi::LUA_REGISTRYINDEX, ffi::LUA_RIDX_GLOBALS);
 
@@ -892,3 +986,4 @@ impl Lua {
 
 static LUA_USERDATA_REGISTRY_KEY: u8 = 0;
 static FUNCTION_METATABLE_REGISTRY_KEY: u8 = 0;
+static USER_REGISTRY_KEY: u8 = 0;
diff --git a/src/tests.rs b/src/tests.rs
index 9faa524..e73bbd8 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -518,6 +518,24 @@ fn test_named_registry_value() {
     };
 }
 
+#[test]
+fn test_registry_value() {
+    let lua = Lua::new();
+
+    let mut r = Some(lua.create_registry_value::<i32>(42).unwrap());
+    let f = lua.create_function(move |lua, ()| {
+        if let Some(r) = r.take() {
+            assert_eq!(lua.registry_value::<i32>(&r)?, 42);
+            lua.remove_registry_value(r);
+        } else {
+            panic!();
+        }
+        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/types.rs b/src/types.rs
index 2a13265..21384a1 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -15,6 +15,9 @@ pub type Number = ffi::lua_Number;
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct LightUserData(pub *mut c_void);
 
+/// An auto generated key into the Lua registry.
+pub struct RegistryKey(pub(crate) c_int);
+
 pub(crate) type Callback<'lua> = Box<
     FnMut(&'lua Lua, MultiValue<'lua>) -> Result<MultiValue<'lua>> + 'lua,
 >;
diff --git a/src/util.rs b/src/util.rs
index aa203fc..0ce67fa 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -390,6 +390,23 @@ pub unsafe fn pop_wrapped_error(state: *mut ffi::lua_State) -> Option<Error> {
     }
 }
 
+// Runs the given function with the Lua garbage collector disabled.  `rlua` assumes that all
+// allocation failures are aborts, so when the garbage collector is disabled, 'm' functions that can
+// cause either an allocation error or a a `__gc` metamethod error are prevented from causing errors
+// at all.  The given function should never panic or longjmp, because this could inadverntently
+// disable the gc.  This is useful when error handling must allocate, and `__gc` errors at that time
+// would shadow more important errors, or be extremely difficult to handle safely.
+pub unsafe fn gc_guard<R, F: FnOnce() -> R>(state: *mut ffi::lua_State, f: F) -> R {
+    if ffi::lua_gc(state, ffi::LUA_GCISRUNNING, 0) != 0 {
+        ffi::lua_gc(state, ffi::LUA_GCSTOP, 0);
+        let r = f();
+        ffi::lua_gc(state, ffi::LUA_GCRESTART, 0);
+        r
+    } else {
+        f()
+    }
+}
+
 struct WrappedError(pub Error);
 struct WrappedPanic(pub Option<Box<Any + Send>>);
 
@@ -556,20 +573,3 @@ unsafe fn get_panic_metatable(state: *mut ffi::lua_State) -> c_int {
 
     ffi::LUA_TTABLE
 }
-
-// Runs the given function with the Lua garbage collector disabled.  `rlua` assumes that all
-// allocation failures are aborts, so when the garbage collector is disabled, 'm' functions that can
-// cause either an allocation error or a a `__gc` metamethod error are prevented from causing errors
-// at all.  The given function should never panic or longjmp, because this could inadverntently
-// disable the gc.  This is useful when error handling must allocate, and `__gc` errors at that time
-// would shadow more important errors, or be extremely difficult to handle safely.
-unsafe fn gc_guard<R, F: FnOnce() -> R>(state: *mut ffi::lua_State, f: F) -> R {
-    if ffi::lua_gc(state, ffi::LUA_GCISRUNNING, 0) != 0 {
-        ffi::lua_gc(state, ffi::LUA_GCSTOP, 0);
-        let r = f();
-        ffi::lua_gc(state, ffi::LUA_GCRESTART, 0);
-        r
-    } else {
-        f()
-    }
-}