bones_scripting/
lua.rs

1use crate::prelude::*;
2
3use bevy_tasks::{ComputeTaskPool, TaskPool, ThreadExecutor};
4use bones_asset::dashmap::mapref::one::{MappedRef, MappedRefMut};
5use bones_lib::ecs::utils::*;
6
7use parking_lot::Mutex;
8pub use piccolo;
9use piccolo::{
10    compiler::{LineNumber, ParseError},
11    registry::{Fetchable, Stashable},
12    Closure, Context, Executor, FromValue, Lua, PrototypeError, StashedClosure, Table, UserData,
13    Value,
14};
15use send_wrapper::SendWrapper;
16use std::{any::Any, rc::Rc, sync::Arc};
17
18#[macro_use]
19mod freeze;
20use freeze::*;
21
22mod asset;
23pub use asset::*;
24
25mod ext;
26pub use ext::*;
27
28pub mod bindings;
29
30/// Install the lua scripting plugin.
31pub fn lua_game_plugin(game: &mut Game) {
32    // Register asset type.
33    LuaScript::register_schema();
34
35    // Add `SchemaLuaMetatable` type data for common types.
36    bindings::register_lua_typedata();
37
38    // Initialize the lua engine resource.
39    game.init_shared_resource::<LuaEngine>();
40}
41
42/// A [`SessionPlugin] that will run the provided lua plugins
43pub struct LuaPluginLoaderSessionPlugin(pub Arc<Vec<Handle<LuaPlugin>>>);
44
45/// Resource containing the lua plugins that have been installed in this session.
46#[derive(HasSchema, Deref, DerefMut, Default, Clone)]
47pub struct LuaPlugins(pub Arc<Vec<Handle<LuaPlugin>>>);
48
49impl SessionPlugin for LuaPluginLoaderSessionPlugin {
50    fn install(self, session: &mut SessionBuilder) {
51        session.insert_resource(LuaPlugins(self.0));
52
53        for lua_stage in [
54            CoreStage::First,
55            CoreStage::PreUpdate,
56            CoreStage::Update,
57            CoreStage::PostUpdate,
58            CoreStage::Last,
59        ] {
60            session.stages.add_system_to_stage(
61                lua_stage,
62                move |engine: Res<LuaEngine>,
63                      asset_server: Res<AssetServer>,
64                      lua_plugins: Res<LuaPlugins>,
65                      world: &World| {
66                    engine.exec(|lua| {
67                        Frozen::<Freeze![&'freeze World]>::in_scope(world, |world| {
68                            lua.enter(|ctx| {
69                                let env = ctx.singletons().get(ctx, bindings::env);
70                                let worldref = WorldRef(world);
71                                worldref.add_to_env(ctx, env);
72                            });
73
74                            for plugin_handle in lua_plugins.iter() {
75                                let Some(plugin) = asset_server.try_get(*plugin_handle) else {
76                                    return;
77                                };
78                                let plugin = plugin.unwrap();
79
80                                // Load the plugin if necessary
81                                if !plugin.has_loaded() {
82                                    if let Err(e) = plugin.load(engine.executor.clone(), lua) {
83                                        tracing::error!("Error loading lua plugin: {e}");
84                                    }
85                                }
86
87                                let mut systems = plugin.systems.borrow_mut();
88                                let systems = systems.as_loaded_mut();
89
90                                for (has_run, closure) in &mut systems.startup {
91                                    if !*has_run {
92                                        let executor = lua.enter(|ctx| {
93                                            let closure = ctx.registry().fetch(closure);
94                                            let ex = Executor::start(ctx, closure.into(), ());
95                                            ctx.registry().stash(&ctx, ex)
96                                        });
97                                        if let Err(e) = lua.execute::<()>(&executor) {
98                                            tracing::error!("Error running lua plugin system: {e}");
99                                        }
100
101                                        *has_run = true;
102                                    }
103                                }
104
105                                for (stage, closure) in &systems.core_stages {
106                                    if stage == &lua_stage {
107                                        let executor = lua.enter(|ctx| {
108                                            let closure = ctx.registry().fetch(closure);
109                                            let ex = Executor::start(ctx, closure.into(), ());
110                                            ctx.registry().stash(&ctx, ex)
111                                        });
112                                        if let Err(e) = lua.execute::<()>(&executor) {
113                                            tracing::error!("Error running lua plugin system: {e}");
114                                        }
115                                    }
116                                }
117                            }
118                        })
119                    });
120                },
121            );
122        }
123    }
124}
125
126/// A frozen reference to the ECS [`World`].
127///
128// This type can be converted into lua userdata for accessing the world from lua.
129#[derive(Deref, DerefMut, Clone)]
130pub struct WorldRef(Frozen<Freeze![&'freeze World]>);
131impl Default for WorldRef {
132    fn default() -> Self {
133        Self(Frozen::new())
134    }
135}
136impl<'gc> FromValue<'gc> for &'gc WorldRef {
137    fn from_value(_ctx: Context<'gc>, value: Value<'gc>) -> Result<Self, piccolo::TypeError> {
138        value.as_static_user_data::<WorldRef>()
139    }
140}
141
142impl WorldRef {
143    /// Convert this [`WorldRef`] into a Lua userdata.
144    pub fn into_userdata(self, ctx: Context<'_>) -> UserData<'_> {
145        let data = UserData::new_static(&ctx, self);
146        data.set_metatable(
147            &ctx,
148            Some(ctx.singletons().get(ctx, bindings::world::metatable)),
149        );
150        data
151    }
152
153    /// Add this world
154    fn add_to_env<'gc>(&self, ctx: Context<'gc>, env: Table<'gc>) {
155        ctx.globals()
156            .set(ctx, "world", self.clone().into_userdata(ctx))
157            .unwrap();
158        for (name, metatable) in [
159            ("world", bindings::world::metatable as fn(Context) -> Table),
160            ("components", bindings::components::metatable),
161            ("resources", bindings::resources::metatable),
162            ("assets", bindings::assets::metatable),
163        ] {
164            let data = UserData::new_static(&ctx, self.clone());
165            data.set_metatable(&ctx, Some(ctx.singletons().get(ctx, metatable)));
166            env.set(ctx, name, data).unwrap();
167        }
168    }
169}
170
171/// Resource used to access the lua scripting engine.
172#[derive(HasSchema, Clone)]
173#[schema(no_default)]
174pub struct LuaEngine {
175    /// The thread-local task executor that is used to spawn any tasks that need access to the
176    /// lua engine which can only be accessed on it's own thread.
177    executor: Arc<ThreadExecutor<'static>>,
178    /// The lua engine state container.
179    state: Arc<SendWrapper<EngineState>>,
180}
181
182/// Internal state for [`LuaEngine`].
183struct EngineState {
184    /// The Lua engine.
185    lua: Mutex<Lua>,
186    /// Persisted lua data we need stored in Rust, such as the environment table, world
187    /// metatable, etc.
188    data: LuaSingletons,
189    /// Cache of the content IDs of loaded scripts, and their compiled lua closures.
190    compiled_scripts: Mutex<HashMap<Cid, StashedClosure>>,
191}
192
193// TODO: Don't Use Function Pointers to Index Lua Singletons.
194// Unfortunately, function pointers with different signatures may be unified by
195// LLVM when activating, for intance, LTO, and that can cause unexpected behavior
196// when using them as indexes into a HashMap for singletons.
197/// Struct for accessing and initializing lua singletons.
198///
199/// This is stored in a lua global and accessed conveniently through our [`CtxExt`] trait
200/// so that we can easily initialize lua tables and callbacks throughout our lua bindings.
201pub struct LuaSingletons {
202    singletons: Rc<AtomicCell<HashMap<usize, Box<dyn Any>>>>,
203}
204impl Default for LuaSingletons {
205    fn default() -> Self {
206        Self {
207            singletons: Rc::new(AtomicCell::new(HashMap::default())),
208        }
209    }
210}
211impl LuaSingletons {
212    /// Fetch a lua singleton, initializing it if it has not yet been created.
213    ///
214    /// The singleton is defined by a function pointer that returns a stashable value.
215    fn get<
216        'gc,
217        S: Fetchable<'gc, Fetched = T> + 'static,
218        T: Stashable<'gc, Stashed = S> + Clone + Copy + 'gc,
219    >(
220        &self,
221        ctx: Context<'gc>,
222        singleton: fn(Context<'gc>) -> T,
223    ) -> T {
224        let map = self.singletons.borrow_mut();
225        let id = singleton as usize;
226        if let Some(entry) = map.get(&id) {
227            let stashed = entry.downcast_ref::<S>().expect(
228                "Encountered two functions with different return types and \
229                the same function pointer.",
230            );
231            ctx.registry().fetch(stashed)
232        } else {
233            drop(map); // Make sure we don't deadlock
234            let v = singleton(ctx);
235            let stashed = ctx.registry().stash(&ctx, v);
236            self.singletons.borrow_mut().insert(id, Box::new(stashed));
237            v
238        }
239    }
240}
241
242impl Default for EngineState {
243    fn default() -> Self {
244        // Initialize an empty lua engine and our lua data.
245        let mut lua = Lua::core();
246        lua.try_enter(|ctx| {
247            // Insert our lua singletons.
248            ctx.globals().set(
249                ctx,
250                "luasingletons",
251                UserData::new_static(&ctx, LuaSingletons::default()),
252            )?;
253            Ok(())
254        })
255        .unwrap();
256        Self {
257            lua: Mutex::new(lua),
258            data: default(),
259            compiled_scripts: default(),
260        }
261    }
262}
263
264impl Default for LuaEngine {
265    /// Initialize the Lua engine.
266    fn default() -> Self {
267        // Make sure the compute task pool is initialized
268        ComputeTaskPool::init(TaskPool::new);
269
270        #[cfg(not(target_arch = "wasm32"))]
271        let executor = {
272            let (send, recv) = async_channel::bounded(1);
273
274            // Spawn the executor task that will be used for the lua engine.
275            let pool = ComputeTaskPool::get();
276            pool.spawn_local(async move {
277                let executor = Arc::new(ThreadExecutor::new());
278                send.try_send(executor.clone()).unwrap();
279
280                let ticker = (*executor).ticker().unwrap();
281                loop {
282                    ticker.tick().await;
283                }
284            })
285            .detach();
286            pool.with_local_executor(|local| while local.try_tick() {});
287
288            recv.try_recv().unwrap()
289        };
290
291        #[cfg(target_arch = "wasm32")]
292        let executor = Arc::new(ThreadExecutor::new());
293
294        LuaEngine {
295            executor,
296            state: Arc::new(SendWrapper::new(default())),
297        }
298    }
299}
300
301impl LuaEngine {
302    /// Access the lua engine to run code on it.
303    pub fn exec<'a, F: FnOnce(&mut Lua) + Send + 'a>(&self, f: F) {
304        let pool = ComputeTaskPool::get();
305
306        // Create a new scope spawned on the lua engine thread.
307        pool.scope_with_executor(false, Some(&self.executor), |scope| {
308            scope.spawn_on_external(async {
309                f(&mut self.state.lua.lock());
310            });
311        });
312    }
313
314    /// Run a lua script as a system on the given world.
315    pub fn run_script_system(&self, world: &World, script: Handle<LuaScript>) {
316        self.exec(|lua| {
317            Frozen::<Freeze![&'freeze World]>::in_scope(world, |world| {
318                // Wrap world reference so that it can be converted to lua userdata.
319                let worldref = WorldRef(world);
320
321                let executor = lua.try_enter(|ctx| {
322                    // Fetch the env table
323                    let env = self.state.data.get(ctx, bindings::env);
324
325                    // Compile the script
326                    let closure = worldref.with(|world| {
327                        let asset_server = world.resource::<AssetServer>();
328                        let cid = *asset_server
329                            .store
330                            .asset_ids
331                            .get(&script.untyped())
332                            .ok_or_else(|| {
333                                tracing::warn!("Script asset not loaded.");
334                                PrototypeError::Parser(ParseError {
335                                    kind: piccolo::compiler::ParseErrorKind::EndOfStream {
336                                        expected: None,
337                                    },
338                                    line_number: LineNumber(0),
339                                })
340                            })?;
341
342                        let mut compiled_scripts = self.state.compiled_scripts.lock();
343                        let closure = compiled_scripts.get(&cid);
344
345                        Ok::<_, PrototypeError>(match closure {
346                            Some(closure) => ctx.registry().fetch(closure),
347                            None => {
348                                let asset = asset_server.store.assets.get(&cid).unwrap();
349                                let source = &asset.data.cast_ref::<LuaScript>().source;
350                                // TODO: Provide a meaningfull name to loaded scripts.
351                                let closure =
352                                    Closure::load_with_env(ctx, None, source.as_bytes(), env)?;
353                                compiled_scripts.insert(cid, ctx.registry().stash(&ctx, closure));
354
355                                closure
356                            }
357                        })
358                    })?;
359
360                    // Insert the world ref into the global scope
361                    worldref.add_to_env(ctx, env);
362
363                    let ex = Executor::start(ctx, closure.into(), ());
364                    let ex = ctx.registry().stash(&ctx, ex);
365                    Ok(ex)
366                });
367
368                if let Err(e) = executor.and_then(|ex| lua.execute::<()>(&ex)) {
369                    tracing::error!("{e}");
370                }
371            });
372        });
373    }
374}