feat: Implement global namespace registration (#202)

makspll · web-flow · commit c1b6375bd3a2 · 2025-01-13T23:34:08.000Z
* feat: add more functionality to the script registry

* add global namespace alias

* register global functions in lua
diff --git a/crates/bevy_mod_scripting_core/src/bindings/function/namespace.rs b/crates/bevy_mod_scripting_core/src/bindings/function/namespace.rs
@@ -63,13 +63,20 @@ pub enum Namespace {
     OnType(TypeId),
 }
 
+/// A type which implements [`IntoNamespace`] by always converting to the global namespace
+pub struct GlobalNamespace;
+
 pub trait IntoNamespace {
     fn into_namespace() -> Namespace;
 }
 
 impl<T: ?Sized + 'static> IntoNamespace for T {
     fn into_namespace() -> Namespace {
-        Namespace::OnType(TypeId::of::<T>())
+        if TypeId::of::<T>() == TypeId::of::<GlobalNamespace>() {
+            Namespace::Global
+        } else {
+            Namespace::OnType(TypeId::of::<T>())
+        }
     }
 }
 
diff --git a/crates/bevy_mod_scripting_core/src/bindings/function/script_function.rs b/crates/bevy_mod_scripting_core/src/bindings/function/script_function.rs
@@ -360,8 +360,8 @@ impl ScriptFunctionRegistryArc {
 
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct FunctionKey {
-    name: Cow<'static, str>,
-    namespace: Namespace,
+    pub name: Cow<'static, str>,
+    pub namespace: Namespace,
 }
 
 #[derive(Debug, Default)]
@@ -372,6 +372,8 @@ pub struct ScriptFunctionRegistry {
 impl ScriptFunctionRegistry {
     /// Register a script function with the given name. If the name already exists,
     /// the new function will be registered as an overload of the function.
+    ///
+    /// If you want to overwrite an existing function, use [`ScriptFunctionRegistry::overwrite`]
     pub fn register<F, M>(
         &mut self,
         namespace: Namespace,
@@ -380,21 +382,44 @@ impl ScriptFunctionRegistry {
     ) where
         F: ScriptFunction<'static, M>,
     {
-        self.register_overload(namespace, name, func);
+        self.register_overload(namespace, name, func, false);
+    }
+
+    /// Overwrite a function with the given name. If the function does not exist, it will be registered as a new function.
+    pub fn overwrite<F, M>(
+        &mut self,
+        namespace: Namespace,
+        name: impl Into<Cow<'static, str>>,
+        func: F,
+    ) where
+        F: ScriptFunction<'static, M>,
+    {
+        self.register_overload(namespace, name, func, true);
+    }
+
+    /// Remove a function from the registry if it exists. Returns the removed function if it was found.
+    pub fn remove(
+        &mut self,
+        namespace: Namespace,
+        name: impl Into<Cow<'static, str>>,
+    ) -> Option<DynamicScriptFunction> {
+        let name = name.into();
+        self.functions.remove(&FunctionKey { name, namespace })
     }
 
     fn register_overload<F, M>(
         &mut self,
         namespace: Namespace,
         name: impl Into<Cow<'static, str>>,
         func: F,
+        overwrite: bool,
     ) where
         F: ScriptFunction<'static, M>,
     {
         // always start with non-suffixed registration
         // TODO: we do alot of string work, can we make this all more efficient?
         let name: Cow<'static, str> = name.into();
-        if !self.contains(namespace, name.clone()) {
+        if overwrite || !self.contains(namespace, name.clone()) {
             let func = func
                 .into_dynamic_script_function()
                 .with_name(name.clone())
@@ -636,4 +661,34 @@ mod test {
         assert_eq!(all_functions[0].info.name(), "test");
         assert_eq!(all_functions[1].info.name(), "test-1");
     }
+
+    #[test]
+    fn test_overwrite_script_function() {
+        let mut registry = ScriptFunctionRegistry::default();
+        let fn_ = |a: usize, b: usize| a + b;
+        let namespace = Namespace::Global;
+        registry.register(namespace, "test", fn_);
+        let fn_2 = |a: usize, b: i32| a + (b as usize);
+        registry.overwrite(namespace, "test", fn_2);
+
+        let all_functions = registry
+            .iter_overloads(namespace, "test")
+            .expect("Failed to get overloads")
+            .collect::<Vec<_>>();
+
+        assert_eq!(all_functions.len(), 1);
+        assert_eq!(all_functions[0].info.name(), "test");
+    }
+
+    #[test]
+    fn test_remove_script_function() {
+        let mut registry = ScriptFunctionRegistry::default();
+        let fn_ = |a: usize, b: usize| a + b;
+        let namespace = Namespace::Global;
+        registry.register(namespace, "test", fn_);
+        let removed = registry.remove(namespace, "test");
+        assert!(removed.is_some());
+        let removed = registry.remove(namespace, "test");
+        assert!(removed.is_none());
+    }
 }
diff --git a/crates/bevy_mod_scripting_functions/src/test_functions.rs b/crates/bevy_mod_scripting_functions/src/test_functions.rs
@@ -8,7 +8,7 @@ use bevy::{
 use bevy_mod_scripting_core::{
     bindings::{
         function::{
-            namespace::NamespaceBuilder,
+            namespace::{GlobalNamespace, NamespaceBuilder},
             script_function::{CallerContext, DynamicScriptFunctionMut},
         },
         pretty_print::DisplayWithWorld,
@@ -79,4 +79,7 @@ pub fn register_test_functions(world: &mut App) {
                 }
             },
         );
+
+    NamespaceBuilder::<GlobalNamespace>::new_unregistered(world)
+        .register("global_hello_world", || Ok("hi!"));
 }
diff --git a/crates/languages/bevy_mod_scripting_lua/src/lib.rs b/crates/languages/bevy_mod_scripting_lua/src/lib.rs
@@ -1,18 +1,19 @@
 use bevy::{
-    app::{App, Plugin},
+    app::Plugin,
     ecs::{entity::Entity, world::World},
 };
 use bevy_mod_scripting_core::{
     asset::{AssetPathToLanguageMapper, Language},
     bindings::{
-        script_value::ScriptValue, ThreadWorldContainer, WorldCallbackAccess, WorldContainer,
+        function::namespace::Namespace, script_value::ScriptValue, ThreadWorldContainer,
+        WorldCallbackAccess, WorldContainer,
     },
     context::{ContextBuilder, ContextInitializer, ContextPreHandlingInitializer},
     error::ScriptError,
     event::CallbackLabel,
     reflection_extensions::PartialReflectExt,
     script::ScriptId,
-    AddContextInitializer, IntoScriptPluginParams, ScriptingPlugin,
+    IntoScriptPluginParams, ScriptingPlugin,
 };
 use bindings::{
     reference::{LuaReflectReference, LuaStaticReflectReference},
@@ -48,16 +49,62 @@ impl Default for LuaScriptingPlugin {
                 language_mapper: Some(AssetPathToLanguageMapper {
                     map: lua_language_mapper,
                 }),
-                context_initializers: vec![|_script_id, context| {
-                    context
-                        .globals()
-                        .set(
-                            "world",
-                            LuaStaticReflectReference(std::any::TypeId::of::<World>()),
-                        )
-                        .map_err(ScriptError::from_mlua_error)?;
-                    Ok(())
-                }],
+                context_initializers: vec![
+                    |_script_id, context| {
+                        // set the world global
+                        context
+                            .globals()
+                            .set(
+                                "world",
+                                LuaStaticReflectReference(std::any::TypeId::of::<World>()),
+                            )
+                            .map_err(ScriptError::from_mlua_error)?;
+                        Ok(())
+                    },
+                    |_script_id, context: &mut Lua| {
+                        // set static globals
+                        let world = ThreadWorldContainer.get_world();
+                        let type_registry = world.type_registry();
+                        let type_registry = type_registry.read();
+
+                        for registration in type_registry.iter() {
+                            // only do this for non generic types
+                            // we don't want to see `Vec<Entity>:function()` in lua
+                            if !registration.type_info().generics().is_empty() {
+                                continue;
+                            }
+
+                            if let Some(global_name) =
+                                registration.type_info().type_path_table().ident()
+                            {
+                                let ref_ = LuaStaticReflectReference(registration.type_id());
+                                context
+                                    .globals()
+                                    .set(global_name, ref_)
+                                    .map_err(ScriptError::from_mlua_error)?;
+                            }
+                        }
+
+                        // go through functions in the global namespace and add them to the lua context
+                        let script_function_registry = world.script_function_registry();
+                        let script_function_registry = script_function_registry.read();
+
+                        for (key, function) in script_function_registry
+                            .iter_all()
+                            .filter(|(k, _)| k.namespace == Namespace::Global)
+                        {
+                            context
+                                .globals()
+                                .set(
+                                    key.name.to_string(),
+                                    LuaScriptValue::from(ScriptValue::Function(function.clone())),
+                                )
+                                .map_err(ScriptError::from_mlua_error)?;
+                        }
+
+                        Ok(())
+                    },
+                ],
                 context_pre_handling_initializers: vec![|script_id, entity, context| {
                     let world = ThreadWorldContainer.get_world();
                     context
@@ -89,33 +136,6 @@ impl Plugin for LuaScriptingPlugin {
     fn build(&self, app: &mut bevy::prelude::App) {
         self.scripting_plugin.build(app);
     }
-
-    fn cleanup(&self, app: &mut App) {
-        // find all registered types, and insert dummy for calls
-
-        app.add_context_initializer::<LuaScriptingPlugin>(|_script_id, context: &mut Lua| {
-            let world = ThreadWorldContainer.get_world();
-            let type_registry = world.type_registry();
-            let type_registry = type_registry.read();
-
-            for registration in type_registry.iter() {
-                // only do this for non generic types
-                // we don't want to see `Vec<Entity>:function()` in lua
-                if !registration.type_info().generics().is_empty() {
-                    continue;
-                }
-
-                if let Some(global_name) = registration.type_info().type_path_table().ident() {
-                    let ref_ = LuaStaticReflectReference(registration.type_id());
-                    context
-                        .globals()
-                        .set(global_name, ref_)
-                        .map_err(ScriptError::from_mlua_error)?;
-                }
-            }
-            Ok(())
-        });
-    }
 }
 
 pub fn lua_context_load(
diff --git a/crates/languages/bevy_mod_scripting_lua/tests/data/globals/dynamic_globals_are_in_scope.lua b/crates/languages/bevy_mod_scripting_lua/tests/data/globals/dynamic_globals_are_in_scope.lua
@@ -0,0 +1 @@
+assert(global_hello_world() == "hi!", "global_hello_world() == 'hi!'")

Original file line number	Diff line number	Diff line change
`@@ -63,13 +63,20 @@ pub enum Namespace {`
`63`	`63`	`OnType(TypeId),`
`64`	`64`	`}`
`65`	`65`
	`66`	+/// A type which implements [`IntoNamespace`] by always converting to the global namespace
	`67`	`+pub struct GlobalNamespace;`
	`68`	`+`
`66`	`69`	`pub trait IntoNamespace {`
`67`	`70`	`fn into_namespace() -> Namespace;`
`68`	`71`	`}`
`69`	`72`
`70`	`73`	`impl<T: ?Sized + 'static> IntoNamespace for T {`
`71`	`74`	`fn into_namespace() -> Namespace {`
`72`		`- Namespace::OnType(TypeId::of::<T>())`
	`75`	`+ if TypeId::of::<T>() == TypeId::of::<GlobalNamespace>() {`
	`76`	`+ Namespace::Global`
	`77`	`+ } else {`
	`78`	`+ Namespace::OnType(TypeId::of::<T>())`
	`79`	`+ }`
`73`	`80`	`}`
`74`	`81`	`}`
`75`	`82`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+assert(global_hello_world() == "hi!", "global_hello_world() == 'hi!'")`