Start work on custom instance properties and methods in roblox builtin

src/lune/builtins/roblox/mod.rs

@@ -6,7 +6,7 @@ use crate::{
     roblox::{
         self,
         document::{Document, DocumentError, DocumentFormat, DocumentKind},
-        instance::Instance,
+        instance::{registry::InstanceRegistry, Instance},
         reflection::Database as ReflectionDatabase,
     },
 };
@@ -31,6 +31,8 @@ pub fn create(lua: &'static Lua) -> LuaResult<LuaTable> {
         .with_async_function("serializeModel", serialize_model)?
         .with_function("getAuthCookie", get_auth_cookie)?
         .with_function("getReflectionDatabase", get_reflection_database)?
+        .with_function("implementProperty", implement_property)?
+        .with_function("implementMethod", implement_method)?
         .build_readonly()
 }
 
@@ -105,3 +107,43 @@ fn get_auth_cookie(_: &Lua, raw: Option<bool>) -> LuaResult<Option<String>> {
 fn get_reflection_database(_: &Lua, _: ()) -> LuaResult<ReflectionDatabase> {
     Ok(*REFLECTION_DATABASE.get_or_init(ReflectionDatabase::new))
 }
+
+fn implement_property(
+    lua: &Lua,
+    (class_name, property_name, property_getter, property_setter): (
+        String,
+        String,
+        LuaFunction,
+        Option<LuaFunction>,
+    ),
+) -> LuaResult<()> {
+    let property_setter = match property_setter {
+        Some(setter) => setter,
+        None => {
+            let property_name = property_name.clone();
+            lua.create_function(move |_, _: LuaMultiValue| {
+                Err::<(), _>(LuaError::runtime(format!(
+                    "Property '{property_name}' is read-only"
+                )))
+            })?
+        }
+    };
+    // TODO: Wrap getter and setter functions in async compat layers,
+    // the roblox library does not know about the Lune runtime or the
+    // scheduler and users may want to call async functions, some of
+    // which are not obvious that they are async such as print, warn, ...
+    InstanceRegistry::insert_property_getter(lua, &class_name, &property_name, property_getter)
+        .into_lua_err()?;
+    InstanceRegistry::insert_property_setter(lua, &class_name, &property_name, property_setter)
+        .into_lua_err()?;
+    Ok(())
+}
+
+fn implement_method(
+    lua: &Lua,
+    (class_name, method_name, method): (String, String, LuaFunction),
+) -> LuaResult<()> {
+    // TODO: Same as above, wrap the provided method in an async compat layer
+    InstanceRegistry::insert_method(lua, &class_name, &method_name, method).into_lua_err()?;
+    Ok(())
+}

src/roblox/instance/base.rs

@@ -15,7 +15,7 @@ use crate::roblox::{
     shared::instance::{class_is_a, find_property_info},
 };
 
-use super::{data_model, Instance};
+use super::{data_model, registry::InstanceRegistry, Instance};
 
 pub fn add_methods<'lua, M: LuaUserDataMethods<'lua, Instance>>(m: &mut M) {
     m.add_meta_method(LuaMetaMethod::ToString, |lua, this, ()| {
@@ -267,6 +267,10 @@ fn instance_property_get<'lua>(
         }
     } else if let Some(inst) = this.find_child(|inst| inst.name == prop_name) {
         Ok(LuaValue::UserData(lua.create_userdata(inst)?))
+    } else if let Some(getter) = InstanceRegistry::find_property_getter(lua, this, &prop_name) {
+        getter.call(this.clone())
+    } else if let Some(method) = InstanceRegistry::find_method(lua, this, &prop_name) {
+        Ok(LuaValue::Function(method))
     } else {
         Err(LuaError::RuntimeError(format!(
             "{} is not a valid member of {}",
@@ -317,40 +321,39 @@ fn instance_property_set<'lua>(
         _ => {}
     }
 
-    let info = match find_property_info(&this.class_name, &prop_name) {
-        Some(b) => b,
-        None => {
-            return Err(LuaError::RuntimeError(format!(
-                "{} is not a valid member of {}",
-                prop_name, this
+    if let Some(info) = find_property_info(&this.class_name, &prop_name) {
+        if let Some(enum_name) = info.enum_name {
+            match LuaUserDataRef::<EnumItem>::from_lua(prop_value, lua) {
+                Ok(given_enum) if given_enum.parent.desc.name == enum_name => {
+                    this.set_property(prop_name, DomValue::Enum((*given_enum).clone().into()));
+                    Ok(())
+                }
+                Ok(given_enum) => Err(LuaError::RuntimeError(format!(
+                    "Failed to set property '{}' - expected Enum.{}, got Enum.{}",
+                    prop_name, enum_name, given_enum.parent.desc.name
+                ))),
+                Err(e) => Err(e),
+            }
+        } else if let Some(dom_type) = info.value_type {
+            match prop_value.lua_to_dom_value(lua, Some(dom_type)) {
+                Ok(dom_value) => {
+                    this.set_property(prop_name, dom_value);
+                    Ok(())
+                }
+                Err(e) => Err(e.into()),
+            }
+        } else {
+            Err(LuaError::RuntimeError(format!(
+                "Failed to set property '{}' - malformed property info",
+                prop_name
             )))
         }
-    };
-
-    if let Some(enum_name) = info.enum_name {
-        match LuaUserDataRef::<EnumItem>::from_lua(prop_value, lua) {
-            Ok(given_enum) if given_enum.parent.desc.name == enum_name => {
-                this.set_property(prop_name, DomValue::Enum((*given_enum).clone().into()));
-                Ok(())
-            }
-            Ok(given_enum) => Err(LuaError::RuntimeError(format!(
-                "Failed to set property '{}' - expected Enum.{}, got Enum.{}",
-                prop_name, enum_name, given_enum.parent.desc.name
-            ))),
-            Err(e) => Err(e),
-        }
-    } else if let Some(dom_type) = info.value_type {
-        match prop_value.lua_to_dom_value(lua, Some(dom_type)) {
-            Ok(dom_value) => {
-                this.set_property(prop_name, dom_value);
-                Ok(())
-            }
-            Err(e) => Err(e.into()),
-        }
+    } else if let Some(setter) = InstanceRegistry::find_property_setter(lua, this, &prop_name) {
+        setter.call((this.clone(), prop_value))
     } else {
         Err(LuaError::RuntimeError(format!(
-            "Failed to set property '{}' - malformed property info",
-            prop_name
+            "{} is not a valid member of {}",
+            prop_name, this
         )))
     }
 }

src/roblox/instance/data_model.rs

@@ -12,8 +12,8 @@ use super::Instance;
 
 pub const CLASS_NAME: &str = "DataModel";
 
-pub fn add_fields<'lua, M: LuaUserDataFields<'lua, Instance>>(m: &mut M) {
-    add_class_restricted_getter(m, CLASS_NAME, "Workspace", data_model_get_workspace);
+pub fn add_fields<'lua, F: LuaUserDataFields<'lua, Instance>>(f: &mut F) {
+    add_class_restricted_getter(f, CLASS_NAME, "Workspace", data_model_get_workspace);
 }
 
 pub fn add_methods<'lua, M: LuaUserDataMethods<'lua, Instance>>(m: &mut M) {

src/roblox/instance/mod.rs

@@ -23,6 +23,8 @@ pub(crate) mod data_model;
 pub(crate) mod terrain;
 pub(crate) mod workspace;
 
+pub mod registry;
+
 const PROPERTY_NAME_ATTRIBUTES: &str = "Attributes";
 const PROPERTY_NAME_TAGS: &str = "Tags";
 
@@ -735,6 +737,9 @@ impl LuaExportsTable<'_> for Instance {
     and methods we support here - we should only implement methods that
     are necessary for modifying the dom and / or having ergonomic access
     to the dom, not try to replicate Roblox engine behavior of instances
+
+    If a user wants to replicate Roblox engine behavior, they can use the
+    instance registry, and register properties + methods from the lua side
 */
 impl LuaUserData for Instance {
     fn add_fields<'lua, F: LuaUserDataFields<'lua, Self>>(fields: &mut F) {

src/roblox/instance/registry.rs

@@ -0,0 +1,225 @@
+use std::{
+    borrow::Borrow,
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use mlua::{prelude::*, AppDataRef};
+use thiserror::Error;
+
+use super::Instance;
+
+type InstanceRegistryMap = HashMap<String, HashMap<String, LuaRegistryKey>>;
+
+#[derive(Debug, Clone, Error)]
+pub enum InstanceRegistryError {
+    #[error("class name '{0}' is not valid")]
+    InvalidClassName(String),
+    #[error("class '{class_name}' already registered method '{method_name}'")]
+    MethodAlreadyExists {
+        class_name: String,
+        method_name: String,
+    },
+    #[error("class '{class_name}' already registered property '{property_name}'")]
+    PropertyAlreadyExists {
+        class_name: String,
+        property_name: String,
+    },
+}
+
+#[derive(Debug, Clone)]
+pub struct InstanceRegistry {
+    getters: Arc<Mutex<InstanceRegistryMap>>,
+    setters: Arc<Mutex<InstanceRegistryMap>>,
+    methods: Arc<Mutex<InstanceRegistryMap>>,
+}
+
+impl InstanceRegistry {
+    // NOTE: We lazily create the instance registry instead
+    // of always creating it together with the roblox builtin
+    // since it is less commonly used and it simplifies some app
+    // data borrowing relationship problems we'd otherwise have
+    fn get_or_create(lua: &Lua) -> AppDataRef<'_, Self> {
+        if lua.app_data_ref::<Self>().is_none() {
+            lua.set_app_data(Self {
+                getters: Arc::new(Mutex::new(HashMap::new())),
+                setters: Arc::new(Mutex::new(HashMap::new())),
+                methods: Arc::new(Mutex::new(HashMap::new())),
+            });
+        }
+        lua.app_data_ref::<Self>()
+            .expect("Missing InstanceRegistry in app data")
+    }
+
+    pub fn insert_method<'lua>(
+        lua: &'lua Lua,
+        class_name: &str,
+        method_name: &str,
+        method: LuaFunction<'lua>,
+    ) -> Result<(), InstanceRegistryError> {
+        let registry = Self::get_or_create(lua);
+
+        let mut methods = registry
+            .methods
+            .lock()
+            .expect("Failed to lock instance registry methods");
+
+        let class_methods = methods.entry(class_name.to_string()).or_default();
+        if class_methods.contains_key(method_name) {
+            return Err(InstanceRegistryError::MethodAlreadyExists {
+                class_name: class_name.to_string(),
+                method_name: method_name.to_string(),
+            });
+        }
+
+        let key = lua
+            .create_registry_value(method)
+            .expect("Failed to store method in lua registry");
+        class_methods.insert(method_name.to_string(), key);
+
+        Ok(())
+    }
+
+    pub fn insert_property_getter<'lua>(
+        lua: &'lua Lua,
+        class_name: &str,
+        property_name: &str,
+        property_getter: LuaFunction<'lua>,
+    ) -> Result<(), InstanceRegistryError> {
+        let registry = Self::get_or_create(lua);
+
+        let mut getters = registry
+            .getters
+            .lock()
+            .expect("Failed to lock instance registry getters");
+
+        let class_getters = getters.entry(class_name.to_string()).or_default();
+        if class_getters.contains_key(property_name) {
+            return Err(InstanceRegistryError::PropertyAlreadyExists {
+                class_name: class_name.to_string(),
+                property_name: property_name.to_string(),
+            });
+        }
+
+        let key = lua
+            .create_registry_value(property_getter)
+            .expect("Failed to store getter in lua registry");
+        class_getters.insert(property_name.to_string(), key);
+
+        Ok(())
+    }
+
+    pub fn insert_property_setter<'lua>(
+        lua: &'lua Lua,
+        class_name: &str,
+        property_name: &str,
+        property_setter: LuaFunction<'lua>,
+    ) -> Result<(), InstanceRegistryError> {
+        let registry = Self::get_or_create(lua);
+
+        let mut setters = registry
+            .setters
+            .lock()
+            .expect("Failed to lock instance registry getters");
+
+        let class_setters = setters.entry(class_name.to_string()).or_default();
+        if class_setters.contains_key(property_name) {
+            return Err(InstanceRegistryError::PropertyAlreadyExists {
+                class_name: class_name.to_string(),
+                property_name: property_name.to_string(),
+            });
+        }
+
+        let key = lua
+            .create_registry_value(property_setter)
+            .expect("Failed to store getter in lua registry");
+        class_setters.insert(property_name.to_string(), key);
+
+        Ok(())
+    }
+
+    pub fn find_method<'lua>(
+        lua: &'lua Lua,
+        instance: &Instance,
+        method_name: &str,
+    ) -> Option<LuaFunction<'lua>> {
+        let registry = Self::get_or_create(lua);
+        let methods = registry
+            .methods
+            .lock()
+            .expect("Failed to lock instance registry methods");
+
+        class_name_chain(&instance.class_name)
+            .iter()
+            .find_map(|&class_name| {
+                methods
+                    .get(class_name)
+                    .and_then(|class_methods| class_methods.get(method_name))
+                    .map(|key| lua.registry_value::<LuaFunction>(key).unwrap())
+            })
+    }
+
+    pub fn find_property_getter<'lua>(
+        lua: &'lua Lua,
+        instance: &Instance,
+        property_name: &str,
+    ) -> Option<LuaFunction<'lua>> {
+        let registry = Self::get_or_create(lua);
+        let getters = registry
+            .getters
+            .lock()
+            .expect("Failed to lock instance registry getters");
+
+        class_name_chain(&instance.class_name)
+            .iter()
+            .find_map(|&class_name| {
+                getters
+                    .get(class_name)
+                    .and_then(|class_getters| class_getters.get(property_name))
+                    .map(|key| lua.registry_value::<LuaFunction>(key).unwrap())
+            })
+    }
+
+    pub fn find_property_setter<'lua>(
+        lua: &'lua Lua,
+        instance: &Instance,
+        property_name: &str,
+    ) -> Option<LuaFunction<'lua>> {
+        let registry = Self::get_or_create(lua);
+        let setters = registry
+            .setters
+            .lock()
+            .expect("Failed to lock instance registry setters");
+
+        class_name_chain(&instance.class_name)
+            .iter()
+            .find_map(|&class_name| {
+                setters
+                    .get(class_name)
+                    .and_then(|class_setters| class_setters.get(property_name))
+                    .map(|key| lua.registry_value::<LuaFunction>(key).unwrap())
+            })
+    }
+}
+
+pub fn class_name_chain(class_name: &str) -> Vec<&str> {
+    let db = rbx_reflection_database::get();
+
+    let mut list = vec![class_name];
+    let mut current_name = class_name;
+
+    loop {
+        let class_descriptor = db
+            .classes
+            .get(current_name)
+            .expect("Got invalid class name");
+        if let Some(sup) = &class_descriptor.superclass {
+            current_name = sup.borrow();
+            list.push(current_name);
+        } else {
+            break;
+        }
+    }
+
+    list
+}

src/roblox/instance/workspace.rs

@@ -8,9 +8,9 @@ use super::Instance;
 
 pub const CLASS_NAME: &str = "Workspace";
 
-pub fn add_fields<'lua, M: LuaUserDataFields<'lua, Instance>>(m: &mut M) {
-    add_class_restricted_getter(m, CLASS_NAME, "Terrain", workspace_get_terrain);
-    add_class_restricted_getter(m, CLASS_NAME, "CurrentCamera", workspace_get_camera);
+pub fn add_fields<'lua, F: LuaUserDataFields<'lua, Instance>>(f: &mut F) {
+    add_class_restricted_getter(f, CLASS_NAME, "Terrain", workspace_get_terrain);
+    add_class_restricted_getter(f, CLASS_NAME, "CurrentCamera", workspace_get_camera);
 }
 
 /**