Reduce error_guard code to as little as possible

Also ensure that on error in error_guard the stack is in a predictable place.

3 files changed, 67 insertions, 74 deletions

diff --git a/src/lua.rs b/src/lua.rs
index 0bc3794..034163d 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -223,12 +223,13 @@ impl<'lua> LuaTable<'lua> {
         let key = key.to_lua(lua)?;
         let value = value.to_lua(lua)?;
         unsafe {
-            check_stack(lua.state, 3);
+            check_stack(lua.state, 5);
             lua.push_ref(lua.state, &self.0);
             lua.push_value(lua.state, key);
             lua.push_value(lua.state, value);
-            error_guard(lua.state, 3, 0, |state| {
+            error_guard(lua.state, 3, |state| {
                 ffi::lua_settable(state, -3);
+                ffi::lua_pop(state, 1);
                 Ok(())
             })
         }
@@ -243,15 +244,15 @@ impl<'lua> LuaTable<'lua> {
         let lua = self.0.lua;
         let key = key.to_lua(lua)?;
         unsafe {
-            check_stack(lua.state, 2);
+            check_stack(lua.state, 4);
             lua.push_ref(lua.state, &self.0);
             lua.push_value(lua.state, key.to_lua(lua)?);
-            let res = error_guard(lua.state, 2, 0, |state| {
+            error_guard(lua.state, 2, |state| {
                 ffi::lua_gettable(state, -2);
-                let res = lua.pop_value(state);
-                ffi::lua_pop(state, 1);
-                Ok(res)
+                Ok(())
             })?;
+            let res = lua.pop_value(lua.state);
+            ffi::lua_pop(lua.state, 1);
             V::from_lua(res, lua)
         }
     }
@@ -261,15 +262,16 @@ impl<'lua> LuaTable<'lua> {
         let lua = self.0.lua;
         let key = key.to_lua(lua)?;
         unsafe {
-            check_stack(lua.state, 2);
+            check_stack(lua.state, 4);
             lua.push_ref(lua.state, &self.0);
             lua.push_value(lua.state, key);
-            error_guard(lua.state, 2, 0, |state| {
+            error_guard(lua.state, 2, |state| {
                 ffi::lua_gettable(state, -2);
-                let has = ffi::lua_isnil(state, -1) == 0;
-                ffi::lua_pop(state, 2);
-                Ok(has)
-            })
+                Ok(())
+            })?;
+            let has = ffi::lua_isnil(lua.state, -1) == 0;
+            ffi::lua_pop(lua.state, 2);
+            Ok(has)
         }
     }
 
@@ -311,10 +313,12 @@ impl<'lua> LuaTable<'lua> {
     pub fn len(&self) -> LuaResult<LuaInteger> {
         let lua = self.0.lua;
         unsafe {
-            error_guard(lua.state, 0, 0, |state| {
-                check_stack(state, 1);
-                lua.push_ref(state, &self.0);
-                Ok(ffi::luaL_len(state, -1))
+            check_stack(lua.state, 3);
+            lua.push_ref(lua.state, &self.0);
+            error_guard(lua.state, 1, |state| {
+                let len = ffi::luaL_len(state, -1);
+                ffi::lua_pop(state, 1);
+                Ok(len)
             })
         }
     }
@@ -381,31 +385,29 @@ where
             let lua = self.table.lua;
 
             unsafe {
-                check_stack(lua.state, 4);
+                check_stack(lua.state, 6);
 
                 lua.push_ref(lua.state, &self.table);
                 lua.push_ref(lua.state, &next_key);
 
-                match error_guard(lua.state, 2, 0, |state| if ffi::lua_next(state, -2) != 0 {
-                    ffi::lua_pushvalue(state, -2);
-                    let key = lua.pop_value(state);
-                    let value = lua.pop_value(state);
-                    let next_key = lua.pop_ref(lua.state);
-                    ffi::lua_pop(lua.state, 1);
-                    Ok(Some((key, value, next_key)))
-                } else {
-                    ffi::lua_pop(lua.state, 1);
-                    Ok(None)
-                }) {
-                    Ok(Some((key, value, next_key))) => {
-                        self.next_key = Some(next_key);
+                match error_guard(lua.state, 2, |state| Ok(ffi::lua_next(state, -2) != 0)) {
+                    Ok(true) => {
+                        ffi::lua_pushvalue(lua.state, -2);
+                        let key = lua.pop_value(lua.state);
+                        let value = lua.pop_value(lua.state);
+                        self.next_key = Some(lua.pop_ref(lua.state));
+                        ffi::lua_pop(lua.state, 1);
+
                         Some((|| {
-                             let key = K::from_lua(key, lua)?;
-                             let value = V::from_lua(value, lua)?;
-                             Ok((key, value))
-                         })())
+                            let key = K::from_lua(key, lua)?;
+                            let value = V::from_lua(value, lua)?;
+                            Ok((key, value))
+                        })())
+                    }
+                    Ok(false) => {
+                        ffi::lua_pop(lua.state, 1);
+                        None
                     }
-                    Ok(None) => None,
                     Err(e) => Some(Err(e)),
                 }
             }
@@ -436,27 +438,20 @@ where
             let lua = self.table.lua;
 
             unsafe {
-                check_stack(lua.state, 2);
+                check_stack(lua.state, 4);
 
                 lua.push_ref(lua.state, &self.table);
-                match error_guard(
-                    lua.state,
-                    1,
-                    0,
-                    |state| if ffi::lua_geti(state, -1, index) != ffi::LUA_TNIL {
-                        let value = lua.pop_value(state);
-                        ffi::lua_pop(state, 1);
-                        Ok(Some(value))
-                    } else {
-                        ffi::lua_pop(state, 2);
-                        Ok(None)
-                    },
-                ) {
-                    Ok(Some(r)) => {
+                match error_guard(lua.state, 1, |state| Ok(ffi::lua_geti(state, -1, index) != ffi::LUA_TNIL)) {
+                    Ok(true) => {
+                        let value = lua.pop_value(lua.state);
+                        ffi::lua_pop(lua.state, 1);
                         self.index = Some(index + 1);
-                        Some(V::from_lua(r, lua))
+                        Some(V::from_lua(value, lua))
                     }
-                    Ok(None) => None,
+                    Ok(false) => {
+                        ffi::lua_pop(lua.state, 2);
+                        None
+                    },
                     Err(e) => Some(Err(e)),
                 }
             }
diff --git a/src/tests.rs b/src/tests.rs
index 92e21ba..a02ac37 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -810,16 +810,10 @@ fn test_expired_userdata() {
         id: u8,
     }
 
-    impl Drop for Userdata {
-        fn drop(&mut self) {
-            println!("dropping {}", self.id);
-        }
-    }
-
     impl LuaUserDataType for Userdata {
         fn add_methods(methods: &mut LuaUserDataMethods<Self>) {
             methods.add_method("access", |lua, this, _| {
-                println!("accessing userdata {}", this.id);
+                assert!(this.id == 123);
                 lua.pack(())
             });
         }
diff --git a/src/util.rs b/src/util.rs
index a3f188b..8fdba9a 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -20,7 +20,7 @@ macro_rules! cstr {
 macro_rules! lua_panic {
     ($state:expr) => {
         {
-            $crate::ffi::lua_settop($state, 0);
+            $crate::ffi::lua_settor($state, 0);
             panic!("rlua internal error");
         }
     };
@@ -151,16 +151,15 @@ where
     res
 }
 
-// Call the given rust function in a protected lua context, similar to pcall.
-// The stack given to the protected function is a separate protected stack. This
-// catches all calls to lua_error, but ffi functions that can call lua_error are
-// still longjmps, and have all the same dangers as longjmps, so extreme care
-// must still be taken in code that uses this function.  Does not call
-// lua_checkstack, and uses 2 extra stack spaces.
+// Call the given rust function in a protected lua context, similar to pcall.  The stack given to
+// the protected function is a separate protected stack. This catches all calls to lua_error, but
+// ffi functions that can call lua_error are still longjmps, and have all the same dangers as
+// longjmps, so extreme care must still be taken in code that uses this function.  Does not call
+// lua_checkstack, and uses 2 extra stack spaces.  On error, the stack position is set to just below
+// the given arguments.
 pub unsafe fn error_guard<F, R>(
     state: *mut ffi::lua_State,
     nargs: c_int,
-    nresults: c_int,
     func: F,
 ) -> LuaResult<R>
 where
@@ -179,7 +178,6 @@ where
     unsafe fn cpcall<F>(
         state: *mut ffi::lua_State,
         nargs: c_int,
-        nresults: c_int,
         mut func: F,
     ) -> LuaResult<()>
     where
@@ -189,15 +187,21 @@ where
         ffi::lua_insert(state, -(nargs + 1));
         ffi::lua_pushlightuserdata(state, &mut func as *mut F as *mut c_void);
         mem::forget(func);
-        handle_error(state, pcall_with_traceback(state, nargs + 1, nresults))
+        handle_error(state, pcall_with_traceback(state, nargs + 1, ffi::LUA_MULTRET))
     }
 
     let mut res = None;
-    cpcall(state, nargs, nresults, |state| {
+    let top = ffi::lua_gettop(state);
+    if let Err(err) = cpcall(state, nargs, |state| {
         res = Some(callback_error(state, || func(state)));
-        ffi::lua_gettop(state)
-    })?;
-    Ok(res.unwrap())
+        let c = ffi::lua_gettop(state);
+        c
+    }) {
+        ffi::lua_settop(state, top - nargs);
+        Err(err)
+    } else {
+        Ok(res.unwrap())
+    }
 }
 
 // If the return code indicates an error, pops the error off of the stack and