elua/src/table.rs

use mlua_sys as lua;

use crate::{ensure_type, AsLua, FromLua, FromLuaStack, LuaRef, PushToLuaStack, Result, StackGuard, State, Value};

#[derive(Clone)]
pub struct Table<'a> {
    state: &'a State,
    pub(crate) lref: LuaRef<'a>,
    /// a boolean indicating if this Table is a Metatable
    mt_marker: bool,
}

impl<'a> Table<'a> {
    /// Create a new table
    pub fn new(state: &'a State) -> Result<Self> {
        let lref = unsafe {
            state.ensure_stack(1)?;

            let s = state.state_ptr();
            lua::lua_newtable(s);
            LuaRef::from_stack(state)?
        };
        
        Ok(Self {
            state,
            lref,
            mt_marker: false,
        })
    }

    pub fn new_meta_table(state: &'a State, name: &str) -> Result<Self> {
        let name_term = format!("{}\0", name);
        let name_term_c = name_term.as_str().as_ptr() as *const i8;
        
        let meta_ref = unsafe {
            let _g = StackGuard::new(state);
            state.ensure_stack(2)?;

            let s = state.state_ptr();
            if lua::luaL_newmetatable(s, name_term_c) == 0 {
                // lua::luaL_getmetatable does not return the type that was
                // retrieved from the registry
                lua::lua_pushstring(s, name_term_c);
                let ty = lua::lua_rawget(s, lua::LUA_REGISTRYINDEX);
                
                if ty != lua::LUA_TTABLE {
                    return Err(crate::Error::RegistryConflict(name.to_string()));
                }
            }
            LuaRef::from_stack(state)?
        };
        
        Ok(Self {
            state,
            lref: meta_ref,
            mt_marker: true
        })
    }

    /// Construct a table with a lua reference to one.
    pub fn with_ref(state: &'a State, lua_ref: LuaRef<'a>, is_metatable: bool) -> Result<Self> {
        let s = state.state_ptr();
        
        unsafe {
            let _g = StackGuard::new(state);
            
            lua_ref.push_to_lua_stack(state)?;
            if lua::lua_istable(s, -1) == 0 {
                panic!("Provided reference is not a table")
            }
        }

        Ok(Self {
            state,
            lref: lua_ref,
            mt_marker: is_metatable,
        })
    }

    /// Set a key to a value in the table.
    /// 
    /// This may trigger the `__newindex` metamethod, see [`Table::raw_set`]
    pub fn set<K, V>(&self, key: K, val: V) -> Result<()>
    where
        K: AsLua<'a>,
        V: AsLua<'a>
    {
        let s = self.state.state_ptr();
        unsafe {
            let _g = StackGuard::new(self.state);

            if self.mt_marker {
                self.state.ensure_stack(4)?;
            } else {
                self.state.ensure_stack(3)?;
            }

            self.lref.push_to_lua_stack(self.state)?;
            key.as_lua(self.state)?
                .push_to_lua_stack(self.state)?;
            val.as_lua(self.state)?
                .push_to_lua_stack(self.state)?;

            lua::lua_settable(s, -3);
        }

        Ok(())
    }

    /// Get a value from the table.
    /// 
    /// This may trigger the `__index` metamethod, see [`Table::raw_get`]
    pub fn get<K, V>(&self, key: K) -> Result<V>
    where
        K: AsLua<'a>,
        V: FromLua<'a>,
    {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(2)?;
            let _g = StackGuard::new(self.state);
            
            self.lref.push_to_lua_stack(self.state)?;
            let val = key.as_lua(self.state)?;
            val.push_to_lua_stack(self.state)?;

            lua::lua_gettable(s, -2); // table[key] is at top of stack
            
            let val = Value::from_lua_stack(self.state)?;
            let val = V::from_lua(self.state, val)?;

            Ok(val)
        }
    }

    /// Returns the length of the table.
    /// 
    /// This may trigger the `__len` metamethod, see [`Table::raw_len`]
    pub fn len(&self) -> Result<u64> {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(1)?;
            let _g = StackGuard::new(self.state);

            self.lref.push_to_lua_stack(self.state)?;
            lua::lua_len(s, -1);

            let len = lua::lua_tonumber(s, -1);

            Ok(len as u64)
        }
    }

    /// Returns a boolean indicating if this table has a key
    pub fn has_key<K>(&self, key: K) -> Result<bool>
    where
        K: AsLua<'a>,
    {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(2)?;
            let _g = StackGuard::new(self.state);
            
            self.lref.push_to_lua_stack(self.state)?;
            let key_val = key.as_lua(self.state)?;
            key_val.push_to_lua_stack(self.state)?;

            // table[key] is at top of stack
            if lua::lua_gettable(s, -2) == lua::LUA_TNIL {
                Ok(true)
            } else {
                Ok(false)
            }
        }
    }

    /// Set a key to a value in the table without calling any meta methods.
    pub fn raw_set<K, V>(&self, key: K, val: V) -> Result<()>
    where
        K: AsLua<'a>,
        V: AsLua<'a>
    {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(3)?;
            let _g = StackGuard::new(self.state);
            
            self.lref.push_to_lua_stack(self.state)?;
            key.as_lua(self.state)?
                .push_to_lua_stack(self.state)?;
            val.as_lua(self.state)?
                .push_to_lua_stack(self.state)?;

            lua::lua_rawset(s, -3);
        }

        Ok(())
    }

    /// Get a value from the table without calling any meta methods.
    pub fn raw_get<K, V>(&self, key: K) -> Result<V>
    where
        K: AsLua<'a>,
        V: FromLua<'a>,
    {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(2)?;
            let _g = StackGuard::new(self.state);

            self.lref.push_to_lua_stack(self.state)?;
            key.as_lua(self.state)?
                .push_to_lua_stack(self.state)?;

            lua::lua_rawget(s, -2); // table[key] is at top of stack
            
            let val = Value::from_lua_stack(self.state)?;
            V::from_lua(self.state, val)
        }
    }

    /// Get the length of the table without calling any meta methods.
    pub fn raw_len(&self) -> Result<u64> {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(1)?;
            let _g = StackGuard::new(self.state);

            self.lref.push_to_lua_stack(self.state)?;
            let len = lua::lua_rawlen(s, -1);
            lua::lua_pop(s, 1); // pop table
            
            Ok(len as u64)
        }
    }

    /// Returns a boolean indicating if this table has a key without calling any meta methods.
    pub fn raw_has_key<K>(&self, key: K) -> Result<bool>
    where
        K: AsLua<'a>,
    {
        let s = self.state.state_ptr();
        unsafe {
            self.state.ensure_stack(2)?;
            let _g = StackGuard::new(self.state);

            self.lref.push_to_lua_stack(self.state)?;
            key.as_lua(self.state)?
                .push_to_lua_stack(self.state)?;
            
            // table[key] is at top of stack
            Ok(lua::lua_rawget(s, -2) != lua::LUA_TNIL)
        }
    }

    /// Set something in the metatable
    /// 
    /// Does nothing if this table is not a metatable
    pub fn set_meta<K, V>(&self, key: K, val: V) -> Result<()>
    where
        K: AsLua<'a>,
        V: AsLua<'a>
    {
        if self.mt_marker {
            unsafe {
                let s = self.state.state_ptr();
                self.state.ensure_stack(3)?;
                let _g = StackGuard::new(self.state);

                self.lref.push_to_lua_stack(self.state)?;
                key.as_lua(self.state)?
                    .push_to_lua_stack(self.state)?;
                val.as_lua(self.state)?
                    .push_to_lua_stack(self.state)?;
                lua::lua_settable(s, -3);
            }
        }

        Ok(())
    }
}

impl<'a> PushToLuaStack<'a> for Table<'a> {
    unsafe fn push_to_lua_stack(&self, state: &State) -> Result<()> {
        // no need to ensure stack, the LuaRef does it for us
        self.lref.push_to_lua_stack(state)?;

        Ok(())
    }
}

impl<'a> FromLuaStack<'a> for Table<'a> {
    unsafe fn from_lua_stack(state: &'a State) -> Result<Self> {
        ensure_type(state, lua::LUA_TTABLE, -1)?;
        Table::with_ref(state, LuaRef::from_stack(state)?, false)
    }
}

impl<'a> FromLua<'a> for Table<'a> {
    fn from_lua(_lua: &'a State, val: Value<'a>) -> crate::Result<Self> {
        val.into_table()
    }
}