Defer metatable return on userdata creation until the end

Relates to #477

Author: khvzak

src/state/raw.rs

@@ -832,7 +832,7 @@ impl RawLua {
 
     pub(crate) unsafe fn push_userdata_metatable<T>(&self, mut registry: UserDataRegistry<T>) -> Result<()> {
         let state = self.state();
-        let _sg = StackGuard::with_top(state, ffi::lua_gettop(state) + 1);
+        let mut stack_guard = StackGuard::new(state);
         check_stack(state, 13)?;
 
         // Prepare metatable, add meta methods first and then meta fields
@@ -863,8 +863,6 @@ impl RawLua {
         }
         let metatable_index = ffi::lua_absindex(state, -1);
 
-        let mut extra_tables_count = 0;
-
         let fields_nrec = registry.fields.len();
         if fields_nrec > 0 {
             // If `__index` is a table then update it in-place
@@ -909,7 +907,6 @@ impl RawLua {
                 rawset_field(state, -2, &k)?;
             }
             field_getters_index = Some(ffi::lua_absindex(state, -1));
-            extra_tables_count += 1;
         }
 
         let mut field_setters_index = None;
@@ -921,7 +918,6 @@ impl RawLua {
                 rawset_field(state, -2, &k)?;
             }
             field_setters_index = Some(ffi::lua_absindex(state, -1));
-            extra_tables_count += 1;
         }
 
         let mut methods_index = None;
@@ -958,7 +954,6 @@ impl RawLua {
                 }
                 _ => {
                     methods_index = Some(ffi::lua_absindex(state, -1));
-                    extra_tables_count += 1;
                 }
             }
         }
@@ -980,8 +975,8 @@ impl RawLua {
             extra_init,
         )?;
 
-        // Pop extra tables to get metatable on top of the stack
-        ffi::lua_pop(state, extra_tables_count);
+        // Update stack guard to keep metatable after return
+        stack_guard.keep(1);
 
         Ok(())
     }

src/util/mod.rs

@@ -66,6 +66,11 @@ impl StackGuard {
     pub(crate) fn with_top(state: *mut ffi::lua_State, top: c_int) -> StackGuard {
         StackGuard { state, top }
     }
+
+    #[inline]
+    pub(crate) fn keep(&mut self, n: c_int) {
+        self.top += n;
+    }
 }
 
 impl Drop for StackGuard {
@@ -129,6 +134,15 @@ pub(crate) unsafe fn push_table(
     }
 }
 
+// Uses 4 stack spaces, does not call checkstack.
+pub(crate) unsafe fn rawget_field(state: *mut ffi::lua_State, table: c_int, field: &str) -> Result<c_int> {
+    ffi::lua_pushvalue(state, table);
+    protect_lua!(state, 1, 1, |state| {
+        ffi::lua_pushlstring(state, field.as_ptr() as *const c_char, field.len());
+        ffi::lua_rawget(state, -2)
+    })
+}
+
 // Uses 4 stack spaces, does not call checkstack.
 pub(crate) unsafe fn rawset_field(state: *mut ffi::lua_State, table: c_int, field: &str) -> Result<()> {
     ffi::lua_pushvalue(state, table);

src/util/userdata.rs

@@ -3,7 +3,7 @@ use std::os::raw::{c_int, c_void};
 use std::{ptr, str};
 
 use crate::error::Result;
-use crate::util::{check_stack, get_metatable_ptr, push_string, push_table, rawset_field, TypeKey};
+use crate::util::{check_stack, get_metatable_ptr, push_table, rawget_field, rawset_field, TypeKey};
 
 // Pushes the userdata and attaches a metatable with __gc method.
 // Internally uses 3 stack spaces, does not call checkstack.
@@ -154,14 +154,11 @@ pub(crate) unsafe fn init_userdata_metatable(
     methods: Option<c_int>,
     extra_init: Option<fn(*mut ffi::lua_State) -> Result<()>>,
 ) -> Result<()> {
-    ffi::lua_pushvalue(state, metatable);
-
     if field_getters.is_some() || methods.is_some() {
         // Push `__index` generator function
         init_userdata_metatable_index(state)?;
 
-        push_string(state, b"__index", true)?;
-        let index_type = ffi::lua_rawget(state, -3);
+        let index_type = rawget_field(state, metatable, "__index")?;
         match index_type {
             ffi::LUA_TNIL | ffi::LUA_TTABLE | ffi::LUA_TFUNCTION => {
                 for &idx in &[field_getters, methods] {
@@ -175,28 +172,27 @@ pub(crate) unsafe fn init_userdata_metatable(
                 // Generate `__index`
                 protect_lua!(state, 4, 1, fn(state) ffi::lua_call(state, 3, 1))?;
             }
-            _ => mlua_panic!("improper __index type {}", index_type),
+            _ => mlua_panic!("improper `__index` type: {}", index_type),
         }
 
-        rawset_field(state, -2, "__index")?;
+        rawset_field(state, metatable, "__index")?;
     }
 
     if let Some(field_setters) = field_setters {
         // Push `__newindex` generator function
         init_userdata_metatable_newindex(state)?;
 
-        push_string(state, b"__newindex", true)?;
-        let newindex_type = ffi::lua_rawget(state, -3);
+        let newindex_type = rawget_field(state, metatable, "__newindex")?;
         match newindex_type {
             ffi::LUA_TNIL | ffi::LUA_TTABLE | ffi::LUA_TFUNCTION => {
                 ffi::lua_pushvalue(state, field_setters);
                 // Generate `__newindex`
                 protect_lua!(state, 3, 1, fn(state) ffi::lua_call(state, 2, 1))?;
             }
-            _ => mlua_panic!("improper __newindex type {}", newindex_type),
+            _ => mlua_panic!("improper `__newindex` type: {}", newindex_type),
         }
 
-        rawset_field(state, -2, "__newindex")?;
+        rawset_field(state, metatable, "__newindex")?;
     }
 
     // Additional initialization
@@ -205,9 +201,7 @@ pub(crate) unsafe fn init_userdata_metatable(
     }
 
     ffi::lua_pushboolean(state, 0);
-    rawset_field(state, -2, "__metatable")?;
-
-    ffi::lua_pop(state, 1);
+    rawset_field(state, metatable, "__metatable")?;
 
     Ok(())
 }