Add raw_insert() and raw_remove() for tables (represented as lists)

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<K: ToLua<'lua>>(&self, key: K) -> Result<()> {
-        self.raw_set(key, Nil)?;
-        Ok(())
-    }
-
     /// Sets a key-value pair without invoking metamethods.
     pub fn raw_set<K: ToLua<'lua>, 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<V: ToLua<'lua>>(&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<K: ToLua<'lua>>(&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.

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()?;

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::<Result<Vec<(i64, i64)>>>()?,
+        table4
+            .clone()
+            .pairs()
+            .collect::<Result<Vec<(i64, i64)>>>()?,
         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::<Result<Vec<(i64, i64)>>>()?,
+        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::<Result<Vec<(i64, i64)>>>()?,
+        vec![(1, 2), (2, 3), (3, 35), (4, 4), (5, 5), (6, 7)]
+    );
+
     Ok(())
 }