From b0d401421e449da599d9076efe62c55f92994223 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Wed, 22 Mar 2023 11:20:47 +0100
Subject: [PATCH] Implement mechanism for implementing roblox methods only on
 certain classes

---
 .../lib-roblox/src/instance/data_model.rs     | 20 +++++++
 packages/lib-roblox/src/instance/mod.rs       |  7 ++-
 packages/lib-roblox/src/shared/classes.rs     | 56 +++++++++++++++++++
 packages/lib-roblox/src/shared/mod.rs         |  1 +
 4 files changed, 82 insertions(+), 2 deletions(-)
 create mode 100644 packages/lib-roblox/src/instance/data_model.rs
 create mode 100644 packages/lib-roblox/src/shared/classes.rs

diff --git a/packages/lib-roblox/src/instance/data_model.rs b/packages/lib-roblox/src/instance/data_model.rs
new file mode 100644
index 0000000..8d00422
--- /dev/null
+++ b/packages/lib-roblox/src/instance/data_model.rs
@@ -0,0 +1,20 @@
+use mlua::prelude::*;
+
+use crate::shared::classes::add_class_restricted_method;
+
+use super::Instance;
+
+pub fn add_methods<'lua, M: LuaUserDataMethods<'lua, Instance>>(methods: &mut M) {
+    add_class_restricted_method(
+        methods,
+        "DataModel",
+        "GetService",
+        |_, _, _service_name: String| Ok(()),
+    );
+    add_class_restricted_method(
+        methods,
+        "DataModel",
+        "FindService",
+        |_, _, _service_name: String| Ok(()),
+    );
+}
diff --git a/packages/lib-roblox/src/instance/mod.rs b/packages/lib-roblox/src/instance/mod.rs
index 0a1b830..6c48318 100644
--- a/packages/lib-roblox/src/instance/mod.rs
+++ b/packages/lib-roblox/src/instance/mod.rs
@@ -15,6 +15,8 @@ use crate::{
     shared::instance::{class_exists, class_is_a, find_property_info},
 };
 
+mod data_model;
+
 lazy_static::lazy_static! {
     static ref INTERNAL_DOM: RwLock<WeakDom> =
         RwLock::new(WeakDom::new(DomInstanceBuilder::new("ROOT")));
@@ -865,8 +867,9 @@ impl LuaUserData for Instance {
                 .find_ancestor(|ancestor| ancestor.referent() == instance.dom_ref)
                 .is_some())
         });
-        // FUTURE: We could pass the "methods" struct to some other functions
-        // here to add inheritance-like behavior and class-specific methods
+        // Here we add inheritance-like behavior for instances by creating
+        // methods that are restricted to specific classnames / base classes
+        data_model::add_methods(methods);
     }
 }
 
diff --git a/packages/lib-roblox/src/shared/classes.rs b/packages/lib-roblox/src/shared/classes.rs
new file mode 100644
index 0000000..bb0fb98
--- /dev/null
+++ b/packages/lib-roblox/src/shared/classes.rs
@@ -0,0 +1,56 @@
+use mlua::prelude::*;
+
+use crate::instance::Instance;
+
+use super::instance::class_is_a;
+
+pub(crate) fn add_class_restricted_method<'lua, M: LuaUserDataMethods<'lua, Instance>, A, R, F>(
+    methods: &mut M,
+    class_name: &'static str,
+    method_name: &'static str,
+    method: F,
+) where
+    A: FromLuaMulti<'lua>,
+    R: ToLuaMulti<'lua>,
+    F: 'static + Fn(&'lua Lua, &Instance, A) -> LuaResult<R>,
+{
+    methods.add_method(method_name, move |lua, this, args| {
+        if class_is_a(this.get_class_name(), class_name).unwrap_or(false) {
+            method(lua, this, args)
+        } else {
+            Err(LuaError::RuntimeError(format!(
+                "{} is not a valid member of {}",
+                method_name, class_name
+            )))
+        }
+    });
+}
+
+#[allow(dead_code)]
+pub(crate) fn add_class_restricted_method_mut<
+    'lua,
+    M: LuaUserDataMethods<'lua, Instance>,
+    A,
+    R,
+    F,
+>(
+    methods: &mut M,
+    class_name: &'static str,
+    method_name: &'static str,
+    method: F,
+) where
+    A: FromLuaMulti<'lua>,
+    R: ToLuaMulti<'lua>,
+    F: 'static + Fn(&'lua Lua, &mut Instance, A) -> LuaResult<R>,
+{
+    methods.add_method_mut(method_name, move |lua, this, args| {
+        if class_is_a(this.get_class_name(), class_name).unwrap_or(false) {
+            method(lua, this, args)
+        } else {
+            Err(LuaError::RuntimeError(format!(
+                "{} is not a valid member of {}",
+                method_name, class_name
+            )))
+        }
+    });
+}
diff --git a/packages/lib-roblox/src/shared/mod.rs b/packages/lib-roblox/src/shared/mod.rs
index 10d13e0..b57e1c4 100644
--- a/packages/lib-roblox/src/shared/mod.rs
+++ b/packages/lib-roblox/src/shared/mod.rs
@@ -1,2 +1,3 @@
+pub(crate) mod classes;
 pub(crate) mod instance;
 pub(crate) mod userdata;