From 51020737093b8dd3f44018bd02a0a54a36d834da Mon Sep 17 00:00:00 2001
From: Bilal Itani <beelol@users.noreply.github.com>
Date: Wed, 25 Jun 2025 19:46:54 -0500
Subject: [PATCH] Add Luau behavior system

---
 Cargo.toml                    |   5 +
 core/script/Cargo.toml        |  23 +++
 core/script/macros/Cargo.toml |  16 ++
 core/script/macros/src/lib.rs |  44 +++++
 core/script/src/lib.rs        | 291 ++++++++++++++++++++++++++++++++++
 tests/luau_behavior.rs        |  51 ++++++
 6 files changed, 430 insertions(+)
 create mode 100644 core/script/Cargo.toml
 create mode 100644 core/script/macros/Cargo.toml
 create mode 100644 core/script/macros/src/lib.rs
 create mode 100644 core/script/src/lib.rs
 create mode 100644 tests/luau_behavior.rs

diff --git a/Cargo.toml b/Cargo.toml
index ed2b369f5fe77..d01ae7332e308 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,6 +28,9 @@ members = [
   "examples/mobile",
   # Examples of using Bevy on no_std platforms.
   "examples/no_std/*",
+  # Luau behavior system
+  "core/script",
+  "core/script/macros",
   # Benchmarks
   "benches",
   # Internal tools that are not published.
@@ -568,6 +571,8 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1.0.140"
 bytemuck = "1.7"
 bevy_render = { path = "crates/bevy_render", version = "0.16.0-dev", default-features = false }
+# for temp file usage in tests
+tempfile = "3"
 # The following explicit dependencies are needed for proc macros to work inside of examples as they are part of the bevy crate itself.
 bevy_ecs = { path = "crates/bevy_ecs", version = "0.16.0-dev", default-features = false }
 bevy_state = { path = "crates/bevy_state", version = "0.16.0-dev", default-features = false }
diff --git a/core/script/Cargo.toml b/core/script/Cargo.toml
new file mode 100644
index 0000000000000..758162cdd46f4
--- /dev/null
+++ b/core/script/Cargo.toml
@@ -0,0 +1,23 @@
+[package]
+name = "luau_behavior"
+version = "0.1.0"
+edition = "2024"
+publish = false
+
+[dependencies]
+mlua = { version = "0.9", features = ["luau"] }
+notify = "6"
+bevy_app = { path = "../../crates/bevy_app", version = "0.16.0-dev" }
+bevy_ecs = { path = "../../crates/bevy_ecs", version = "0.16.0-dev" }
+bevy_input = { path = "../../crates/bevy_input", version = "0.16.0-dev" }
+bevy_time = { path = "../../crates/bevy_time", version = "0.16.0-dev" }
+bevy_utils = { path = "../../crates/bevy_utils", version = "0.16.0-dev" }
+luau_behavior_macros = { path = "macros" }
+crossbeam-channel = "0.5"
+
+[features]
+default = ["std"]
+std = []
+
+[lints]
+workspace = true
diff --git a/core/script/macros/Cargo.toml b/core/script/macros/Cargo.toml
new file mode 100644
index 0000000000000..46ac3afae1383
--- /dev/null
+++ b/core/script/macros/Cargo.toml
@@ -0,0 +1,16 @@
+[package]
+name = "luau_behavior_macros"
+version = "0.1.0"
+edition = "2024"
+publish = false
+
+[lib]
+proc-macro = true
+
+[dependencies]
+quote = "1"
+syn = { version = "2", features = ["full"] }
+proc-macro2 = "1"
+
+[lints]
+workspace = true
diff --git a/core/script/macros/src/lib.rs b/core/script/macros/src/lib.rs
new file mode 100644
index 0000000000000..72a5508f47f90
--- /dev/null
+++ b/core/script/macros/src/lib.rs
@@ -0,0 +1,44 @@
+use proc_macro::TokenStream;
+use quote::{quote, format_ident};
+use syn::{parse_macro_input, ItemFn, FnArg, Pat, Type};
+
+#[proc_macro_attribute]
+pub fn lua_api(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(item as ItemFn);
+    let ident = &input.sig.ident;
+    // collect arg patterns and types
+    let mut arg_names = Vec::new();
+    let mut arg_types = Vec::new();
+    for arg in &input.sig.inputs {
+        if let FnArg::Typed(pat) = arg {
+            if let Pat::Ident(ident_pat) = &*pat.pat {
+                arg_names.push(ident_pat.ident.clone());
+                arg_types.push(*pat.ty.clone());
+            }
+        }
+    }
+    let register_name = format_ident!("register_{}", ident);
+    let output = quote! {
+        #input
+        #[doc(hidden)]
+        pub fn #register_name(lua: &mlua::Lua) -> mlua::Result<()> {
+            let func = lua.create_function(|_, args: (#(#arg_types),*)| {
+                let (#(#arg_names),*) = args;
+                #ident(#(#arg_names),*);
+                Ok(())
+            })?;
+            let globals = lua.globals();
+            let game = match globals.get::<_, mlua::Table>("game") {
+                Ok(t) => t,
+                Err(_) => {
+                    let t = lua.create_table()?;
+                    globals.set("game", t.clone())?;
+                    t
+                }
+            };
+            game.set(stringify!(#ident), func)?;
+            Ok(())
+        }
+    };
+    output.into()
+}
diff --git a/core/script/src/lib.rs b/core/script/src/lib.rs
new file mode 100644
index 0000000000000..5ef649de62e25
--- /dev/null
+++ b/core/script/src/lib.rs
@@ -0,0 +1,291 @@
+#![cfg_attr(docsrs, feature(doc_auto_cfg))]
+#![forbid(unsafe_code)]
+#![no_std]
+
+#[cfg(feature = "std")]
+extern crate std;
+extern crate alloc;
+
+use alloc::{string::String, sync::Arc, vec::Vec};
+use std::path::{Path, PathBuf};
+use std::sync::Mutex;
+
+use bevy_app::{App, Plugin, PostUpdate, Update};
+use bevy_ecs::{component::Component, event::Event, prelude::*, system::Resource};
+use bevy_input::keyboard::KeyCode;
+use bevy_input::Input;
+use bevy_time::Time;
+use bevy_utils::tracing::{error, info};
+use crossbeam_channel::{unbounded, Receiver};
+use mlua::{Lua, Function, Table};
+use notify::{RecommendedWatcher, RecursiveMode, Watcher, EventKind};
+use luau_behavior_macros::lua_api;
+
+/// Events emitted by sample api functions
+#[derive(Event)]
+pub struct MoveToEvent {
+    pub entity: Entity,
+    pub x: f32,
+    pub y: f32,
+}
+
+#[derive(Event)]
+pub struct AttackEvent {
+    pub entity: Entity,
+    pub target: Entity,
+}
+
+#[derive(Event)]
+pub struct GetStateEvent {
+    pub entity: Entity,
+}
+
+#[derive(Default)]
+struct LuaEventQueue {
+    move_to: Vec<MoveToEvent>,
+    attack: Vec<AttackEvent>,
+    get_state: Vec<GetStateEvent>,
+}
+
+/// Global Lua resource
+#[derive(Resource)]
+pub struct LuaResource {
+    pub lua: Lua,
+    events: Arc<Mutex<LuaEventQueue>>, // shared with Lua functions
+}
+
+impl LuaResource {
+    pub fn lua_ref(&self) -> &Lua { &self.lua }
+}
+
+impl Default for LuaResource {
+    fn default() -> Self {
+        let lua = Lua::new();
+        let events = Arc::new(Mutex::new(LuaEventQueue::default()));
+        // register api functions
+        lua.context(|ctx| {
+            EVENTS.with(|c| *c.borrow_mut() = Some(events.clone()));
+            register_move_to(ctx).unwrap();
+            register_attack(ctx).unwrap();
+            register_get_state(ctx).unwrap();
+        });
+        Self { lua, events }
+    }
+}
+
+/// Behavior handle identifying a loaded script
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+pub struct BehaviorHandle(usize);
+
+struct BehaviorData {
+    path: PathBuf,
+    bytecode: Arc<Vec<u8>>,
+    watcher: RecommendedWatcher,
+}
+
+#[derive(Resource)]
+pub struct BehaviorManager {
+    handles: Vec<BehaviorData>,
+    rx: Receiver<(usize, Vec<u8>)>,
+    tx: crossbeam_channel::Sender<(usize, Vec<u8>)>,
+}
+
+impl BehaviorManager {
+    pub fn new() -> Self {
+        let (tx, rx) = unbounded();
+        Self { handles: Vec::new(), rx, tx }
+    }
+
+    pub fn load(&mut self, lua: &Lua, path: &str) -> BehaviorHandle {
+        let idx = self.handles.len();
+        let src = std::fs::read_to_string(path).expect("read script");
+        let func: Function = lua.load(&src).into_function().unwrap();
+        let bytecode = Arc::new(func.dump(false));
+        let tx = self.tx.clone();
+        let path_buf = PathBuf::from(path);
+        let mut watcher: RecommendedWatcher = RecommendedWatcher::new(move |res| {
+            if let Ok(event) = res {
+                if matches!(event.kind, EventKind::Modify(_)) {
+                    if let Ok(src) = std::fs::read_to_string(&path_buf) {
+                        if let Ok(func) = lua.load(&src).into_function() {
+                            let bytes = func.dump(false);
+                            let _ = tx.send((idx, bytes));
+                        }
+                    }
+                }
+            }
+        }).unwrap();
+        watcher.watch(Path::new(path), RecursiveMode::NonRecursive).unwrap();
+        self.handles.push(BehaviorData { path: path.into(), bytecode, watcher });
+        BehaviorHandle(idx)
+    }
+
+    pub fn reload(&mut self, lua: &Lua, handle: BehaviorHandle) {
+        if let Some(data) = self.handles.get_mut(handle.0) {
+            if let Ok(src) = std::fs::read_to_string(&data.path) {
+                if let Ok(func) = lua.load(&src).into_function() {
+                    let new = Arc::new(func.dump(false));
+                    data.bytecode = new.clone();
+                    let bytes = (*new).clone();
+                    let _ = self.tx.send((handle.0, bytes));
+                }
+            }
+        }
+    }
+
+    pub fn attach_instance(&self, lua: &Lua, world: &mut World, entity: Entity, handle: BehaviorHandle) {
+        let data = &self.handles[handle.0];
+        let func: Function = lua.load(&*data.bytecode).into_function().unwrap();
+        let table: Table = func.call(()).unwrap();
+        world.entity_mut(entity).insert(LuauBehavior {
+            handle,
+            bytecode: data.bytecode.clone(),
+            state: table,
+        });
+    }
+
+    pub fn detach_and_clone(&self, lua: &Lua, world: &mut World, entity: Entity) {
+        if let Some(mut b) = world.get_mut::<LuauBehavior>(entity) {
+            let bytes = (*b.bytecode).clone();
+            b.bytecode = Arc::new(bytes.clone());
+            let func: Function = lua.load(&bytes).into_function().unwrap();
+            b.state = func.call(()).unwrap();
+        }
+    }
+}
+
+#[derive(Component)]
+pub struct LuauBehavior {
+    handle: BehaviorHandle,
+    bytecode: Arc<Vec<u8>>,
+    state: Table<'static>,
+}
+
+#[derive(Event)]
+struct BehaviorReloadEvent {
+    handle: BehaviorHandle,
+    old: Arc<Vec<u8>>,
+    new: Arc<Vec<u8>>,
+}
+
+fn process_fs_events(mut manager: ResMut<BehaviorManager>, mut ev: EventWriter<BehaviorReloadEvent>) {
+    for (id, bytes) in manager.rx.try_iter() {
+        if let Some(data) = manager.handles.get_mut(id) {
+            let old = data.bytecode.clone();
+            let new = Arc::new(bytes);
+            data.bytecode = new.clone();
+            ev.send(BehaviorReloadEvent { handle: BehaviorHandle(id), old, new });
+        }
+    }
+}
+
+fn apply_reload(
+    mut q: Query<&mut LuauBehavior>,
+    lua_res: Res<LuaResource>,
+    mut events: EventReader<BehaviorReloadEvent>,
+) {
+    let lua = &lua_res.lua;
+    for ev in events.read() {
+        for mut beh in q.iter_mut() {
+            if beh.handle == ev.handle && Arc::ptr_eq(&beh.bytecode, &ev.old) {
+                beh.bytecode = ev.new.clone();
+                let func: Function = lua.load(&*beh.bytecode).into_function().unwrap();
+                beh.state = func.call(()).unwrap();
+            }
+        }
+    }
+}
+
+fn run_lua_behaviors(
+    lua_res: Res<LuaResource>,
+    time: Res<Time>,
+    mut move_writer: EventWriter<MoveToEvent>,
+    mut attack_writer: EventWriter<AttackEvent>,
+    mut get_state_writer: EventWriter<GetStateEvent>,
+    mut query: Query<&mut LuauBehavior>,
+) {
+    let lua = &lua_res.lua;
+    for mut beh in query.iter_mut() {
+        let dt = time.delta_seconds_f64();
+        lua.context(|ctx| {
+            if let Ok(update) = beh.state.get::<_, Function>("update") {
+                if let Err(e) = update.call::<_, ()>((beh.state.clone(), dt)) {
+                    error!("lua error: {e}");
+                }
+            }
+        });
+    }
+    let mut queue = lua_res.events.lock().unwrap();
+    for ev in queue.move_to.drain(..) { move_writer.send(ev); }
+    for ev in queue.attack.drain(..) { attack_writer.send(ev); }
+    for ev in queue.get_state.drain(..) { get_state_writer.send(ev); }
+}
+
+fn developer_detach_helper(
+    input: Res<Input<KeyCode>>, 
+    mut manager: ResMut<BehaviorManager>,
+    lua_res: Res<LuaResource>,
+    mut commands: Commands,
+    query: Query<Entity, With<LuauBehavior>>,
+) {
+    if input.just_pressed(KeyCode::R) {
+        if let Some(entity) = query.iter().next() {
+            let world = commands.as_mut().world_mut();
+            let lua = &lua_res.lua;
+            manager.detach_and_clone(lua, world, entity);
+        }
+    }
+}
+
+pub struct LuauBehaviorPlugin;
+
+impl Plugin for LuauBehaviorPlugin {
+    fn build(&self, app: &mut App) {
+        app.insert_resource(LuaResource::default());
+        app.insert_resource(BehaviorManager::new())
+            .add_event::<MoveToEvent>()
+            .add_event::<AttackEvent>()
+            .add_event::<GetStateEvent>()
+            .add_event::<BehaviorReloadEvent>()
+            .add_systems(Update, (process_fs_events, developer_detach_helper))
+            .add_systems(PostUpdate, (apply_reload, run_lua_behaviors));
+    }
+}
+
+pub use {
+    LuauBehaviorPlugin, BehaviorManager, BehaviorHandle, LuauBehavior, LuaResource,
+    MoveToEvent, AttackEvent, GetStateEvent,
+};
+
+// Sample API functions
+
+use core::cell::RefCell;
+thread_local! {
+    static EVENTS: RefCell<Option<Arc<Mutex<LuaEventQueue>>>> = const { RefCell::new(None) };
+}
+
+fn with_events(f: impl FnOnce(&mut LuaEventQueue)) {
+    EVENTS.with(|cell| {
+        if let Some(ref rc) = *cell.borrow() {
+            if let Ok(mut q) = rc.lock() {
+                f(&mut q);
+            }
+        }
+    });
+}
+
+#[lua_api]
+pub fn move_to(entity: u64, x: f32, y: f32) {
+    with_events(|q| q.move_to.push(MoveToEvent { entity: Entity::from_bits(entity), x, y }));
+}
+
+#[lua_api]
+pub fn attack(entity: u64, target: u64) {
+    with_events(|q| q.attack.push(AttackEvent { entity: Entity::from_bits(entity), target: Entity::from_bits(target) }));
+}
+
+#[lua_api]
+pub fn get_state(entity: u64) {
+    with_events(|q| q.get_state.push(GetStateEvent { entity: Entity::from_bits(entity) }));
+}
+
diff --git a/tests/luau_behavior.rs b/tests/luau_behavior.rs
new file mode 100644
index 0000000000000..f915309091b73
--- /dev/null
+++ b/tests/luau_behavior.rs
@@ -0,0 +1,51 @@
+use bevy::prelude::*;
+use luau_behavior::*;
+use std::fs::write;
+
+#[test]
+fn behavior_reload_and_override() {
+    let dir = tempfile::tempdir().unwrap();
+    let script_path = dir.path().join("b.lua");
+    write(&script_path, "return { update=function(self,dt) self.x=(self.x or 0)+1 end }").unwrap();
+
+    let mut app = App::new();
+    app.add_plugins(MinimalPlugins).add_plugin(LuauBehaviorPlugin);
+    let handle;
+    let entity;
+    {
+        let mut manager = app.world.resource_mut::<BehaviorManager>();
+        let lua = app.world.resource::<LuaResource>().lua_ref();
+        handle = manager.load(lua, script_path.to_str().unwrap());
+        entity = app.world.spawn_empty().id();
+        manager.attach_instance(lua, app.world_mut(), entity, handle);
+    }
+    app.update();
+    let beh = app.world.entity(entity).get::<LuauBehavior>().unwrap();
+    assert_eq!(beh.state.get::<_, i64>("x").unwrap(), 1);
+
+    write(&script_path, "return { update=function(self,dt) self.x=(self.x or 0)+2 end }").unwrap();
+    {
+        let lua = app.world.resource::<LuaResource>().lua_ref();
+        app.world.resource_mut::<BehaviorManager>().reload(lua, handle);
+    }
+    app.update();
+    let beh = app.world.entity(entity).get::<LuauBehavior>().unwrap();
+    assert_eq!(beh.state.get::<_, i64>("x").unwrap(), 3);
+
+    {
+        let lua = app.world.resource::<LuaResource>().lua_ref();
+        app.world.resource_mut::<BehaviorManager>().detach_and_clone(lua, app.world_mut(), entity);
+    }
+    write(&script_path, "return { update=function(self,dt) self.x=(self.x or 0)+10 end }").unwrap();
+    {
+        let lua = app.world.resource::<LuaResource>().lua_ref();
+        app.world.resource_mut::<BehaviorManager>().reload(lua, handle);
+        let other = app.world.spawn_empty().id();
+        app.world.resource_mut::<BehaviorManager>().attach_instance(lua, app.world_mut(), other, handle);
+        app.update();
+        let beh1 = app.world.entity(entity).get::<LuauBehavior>().unwrap();
+        let beh2 = app.world.entity(other).get::<LuauBehavior>().unwrap();
+        assert_eq!(beh1.state.get::<_, i64>("x").unwrap(), 13);
+        assert_eq!(beh2.state.get::<_, i64>("x").unwrap(), 3);
+    }
+}