[Collection] Implement Userdata serialization

moriokii · Jul 7, 2024 · 9a50784 · 9a50784
1 parent bcd1393
commit 9a50784
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/src/functions/mongodbcollection.rs b/src/functions/mongodbcollection.rs
@@ -3,16 +3,18 @@ use std::ffi::{CStr, CString};
 use futures::TryStreamExt;
 use mongodb::{Collection, Database};
 use mongodb::bson::{Bson, Document};
-use rglua::lua::{lua_pushboolean, lua_setmetatable, luaL_checkstring, luaL_getmetatable, LuaState};
+use rglua::lua::{luaL_checkstring, luaL_getmetatable, lua_pushboolean, lua_setmetatable, LuaState};
 use rglua::prelude::{lua_gettop, lua_istable, lua_newtable, lua_next, lua_pop, lua_pushnil, lua_pushnumber, lua_pushstring, lua_rawseti, lua_settable, lua_tonumber, lua_tostring, lua_type};
+use rglua::userdata::{Angle, Vector};
 
 use crate::logger::{log, LogLevel};
 use crate::mongo::MONGO_WORKER;
-use crate::utils::luautils::{read_userdata, write_userdata};
+use crate::utils::luautils::{check_userdata, read_userdata, write_userdata};
 
 const LUA_TNUMBER: i32 = 3;
 const LUA_TSTRING: i32 = 4;
 const LUA_TTABLE: i32 = 5;
+const LUA_TUDATA: i32 = 7;
 const LUA_TNIL: i32 = 0;
 
 unsafe fn get_table_key(l: LuaState, key_type: i32) -> Result<String, String> {
@@ -72,6 +74,26 @@ fn lua_table_to_bson(l: LuaState, index: i32) -> Result<Document, String> {
                     let nested_doc = lua_table_to_bson(l, lua_gettop(l))?;
                     doc.insert(key, Bson::Document(nested_doc));
                 }
+                LUA_TUDATA => {
+                    if let Ok(angle) = check_userdata::<Angle>(l, -1, cstr!("Angle")) {
+                        let mut value = Document::new();
+                        value.insert("type", "Angle");
+                        value.insert("p", angle.p);
+                        value.insert("y", angle.y);
+                        value.insert("r", angle.r);
+
+                        doc.insert(key, Bson::Document(value));
+                    } else if let Ok(vector) = check_userdata::<Vector>(l, -1, cstr!("Vector")) {
+                        let mut value = Document::new();
+                        value.insert("type", "Vector");
+                        value.insert("x", vector.x);
+                        value.insert("y", vector.y);
+                        value.insert("z", vector.z);
+
+                        doc.insert(key, Bson::Document(value));
+                    }
+                }
+
                 _ => {
                     lua_pop(l, 1);
                     return Err("Unsupported value type".to_string());

diff --git a/src/utils/luautils.rs b/src/utils/luautils.rs
@@ -1,4 +1,4 @@
-use rglua::prelude::{lua_newuserdata, lua_touserdata, LuaState};
+use rglua::{lua::{lua_getfield, lua_getmetatable, lua_pop, lua_rawequal, LuaString, Userdata, REGISTRYINDEX}, prelude::{lua_newuserdata, lua_touserdata, LuaState}};
 
 pub fn write_userdata<T>(l: LuaState, data: T) {
     let data_ptr = lua_newuserdata(l, std::mem::size_of::<T>()) as *mut T;
@@ -14,4 +14,25 @@ pub fn read_userdata<T: Clone>(l: LuaState) -> Result<T, String> {
     } else {
         Ok(unsafe { (*data_ptr).clone() })
     }
+}
+
+pub fn check_userdata<T: Clone>(l: LuaState, idx: i32, key: LuaString) -> Result<T, String> {
+    let udata = lua_touserdata(l, idx);
+    if udata.is_null() {
+        return Err("Invalid userdata".to_string());
+    }
+
+    if lua_getmetatable(l, idx) == 0 {
+        return Err("Userdata without metatable".to_string());
+    }
+
+    lua_getfield(l, REGISTRYINDEX, key);
+
+    if lua_rawequal(l, -1, -2) == 1 {
+        lua_pop(l, 2);
+        Ok(unsafe { (*( (*(udata as *mut Userdata)).data as *mut T)).clone() })
+    } else {
+        lua_pop(l, 2);
+        Err("Not rawequal".to_string())
+    }
 }