diff --git a/src/table.rs b/src/table.rs index bb45ea5..a02a9f1 100644 --- a/src/table.rs +++ b/src/table.rs @@ -1,7 +1,7 @@ use std::marker::PhantomData; use std::os::raw::c_int; -use crate::error::Result; +use crate::error::{Error, Result}; use crate::ffi; use crate::function::Function; use crate::types::{Integer, LuaRef}; @@ -222,13 +222,6 @@ impl<'lua> Table<'lua> { Ok(false) } - /// Removes a key from the table, returning the value at the key - /// if the key was previously in the table. - pub fn raw_remove>(&self, key: K) -> Result<()> { - self.raw_set(key, Nil)?; - Ok(()) - } - /// Sets a key-value pair without invoking metamethods. pub fn raw_set, V: ToLua<'lua>>(&self, key: K, value: V) -> Result<()> { let lua = self.0.lua; @@ -269,6 +262,70 @@ impl<'lua> Table<'lua> { V::from_lua(value, lua) } + /// Inserts element value at position idx to the table, shifting up the elements from table[idx]. + /// The worst case complexity is O(n), where n is the table length. + pub fn raw_insert>(&self, idx: Integer, value: V) -> Result<()> { + let lua = self.0.lua; + let size = self.raw_len(); + if idx < 1 || idx > size + 1 { + return Err(Error::RuntimeError("index out of bounds".to_string())); + } + + let value = value.to_lua(lua)?; + unsafe { + let _sg = StackGuard::new(lua.state); + assert_stack(lua.state, 6); + + lua.push_ref(&self.0); + lua.push_value(value)?; + + protect_lua_closure(lua.state, 2, 0, |state| { + for i in (idx..size + 1).rev() { + // table[i+1] = table[i] + ffi::lua_rawgeti(state, -2, i); + ffi::lua_rawseti(state, -3, i + 1); + } + ffi::lua_rawseti(state, -2, idx); + }) + } + } + + /// Removes a key from the table. + /// + /// If `key` is an integer, mlua shifts down the elements from table[key+1], + /// and erases element table[key]. The complexity is O(n) in worst case, + /// where n is the table length. + /// + /// For othey key types this is equivalent to setting table[key] = nil. + pub fn raw_remove>(&self, key: K) -> Result<()> { + let lua = self.0.lua; + let key = key.to_lua(lua)?; + match key { + Value::Integer(idx) => { + let size = self.raw_len(); + if idx < 1 || idx > size { + return Err(Error::RuntimeError("index out of bounds".to_string())); + } + unsafe { + let _sg = StackGuard::new(lua.state); + assert_stack(lua.state, 6); + + lua.push_ref(&self.0); + + protect_lua_closure(lua.state, 1, 0, |state| { + for i in idx..size { + ffi::lua_rawgeti(state, -1, i + 1); + ffi::lua_rawseti(state, -2, i); + } + ffi::lua_pushnil(state); + ffi::lua_rawseti(state, -2, size); + }) + } + } + _ => self.raw_set(key, Nil), + } + } + /// Returns the result of the Lua `#` operator. /// /// This might invoke the `__len` metamethod. Use the [`raw_len`] method if that is not desired. diff --git a/tests/memory.rs b/tests/memory.rs index 0bcdeb6..53c08bd 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -16,6 +16,7 @@ use mlua::{Lua, Result, UserData}; #[test] fn test_gc_control() -> Result<()> { let lua = Lua::new(); + let globals = lua.globals(); #[cfg(any(feature = "lua53", feature = "lua52"))] { @@ -30,9 +31,8 @@ fn test_gc_control() -> Result<()> { impl UserData for MyUserdata {} let rc = Arc::new(()); - lua.globals() - .set("userdata", lua.create_userdata(MyUserdata(rc.clone()))?)?; - lua.globals().raw_remove("userdata")?; + globals.set("userdata", lua.create_userdata(MyUserdata(rc.clone()))?)?; + globals.raw_remove("userdata")?; assert_eq!(Arc::strong_count(&rc), 2); lua.gc_collect()?; diff --git a/tests/table.rs b/tests/table.rs index 1d788e2..26d9516 100644 --- a/tests/table.rs +++ b/tests/table.rs @@ -91,10 +91,32 @@ fn test_table() -> Result<()> { globals.set("table4", lua.create_sequence_from(vec![1, 2, 3, 4, 5])?)?; let table4 = globals.get::<_, Table>("table4")?; assert_eq!( - table4.pairs().collect::>>()?, + table4 + .clone() + .pairs() + .collect::>>()?, vec![(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)] ); + table4.raw_insert(4, 35)?; + table4.raw_insert(7, 7)?; + assert_eq!( + table4 + .clone() + .pairs() + .collect::>>()?, + vec![(1, 1), (2, 2), (3, 3), (4, 35), (5, 4), (6, 5), (7, 7)] + ); + + table4.raw_remove(1)?; + assert_eq!( + table4 + .clone() + .pairs() + .collect::>>()?, + vec![(1, 2), (2, 3), (3, 35), (4, 4), (5, 5), (6, 7)] + ); + Ok(()) }