From 0e419483e8477858896b20456760fd03850a7cc4 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Wed, 21 Oct 2026 17:47:18 -0700 Subject: [PATCH 01/10] Look up a Grain module's functions in its program when called --- src/eval/expr.rs | 9 +- src/grain/mod.rs | 4 +- src/grain/program.rs | 57 --------- src/grain/vm/callback.rs | 71 ++--------- src/grain/vm/mod.rs | 18 ++- src/module/grain.rs | 258 ++++++++++++++++++++++++++++++++++++++ src/module/lazy.rs | 56 +++++++-- src/module/mod.rs | 133 ++++++++++++++++++-- tests/grain/allocation.rs | 27 ++++ tests/grain/module.rs | 96 ++++++++++++++ 10 files changed, 573 insertions(+), 156 deletions(-) create mode 100644 src/module/grain.rs diff --git a/src/eval/expr.rs b/src/eval/expr.rs index 2b0d33201..ce315ddc8 100644 --- a/src/eval/expr.rs +++ b/src/eval/expr.rs @@ -109,16 +109,11 @@ impl Engine { // Finally, see if it is a script function name and return a function pointer #[cfg(not(feature = "no_function"))] - if let Some(func) = global + if let Some(fn_def) = global .lib .iter() - .flat_map(|m| m.iter_fn()) - .filter(|(f, _)| f.is_script()) - .filter(|(_, meta)| meta.name == var_name.as_str()) - .map(|(f, _)| f) - .next() + .find_map(|m| m.get_script_fn_named(var_name, None)) { - let fn_def = func.get_script_fn_def().unwrap(); let val: Dynamic = crate::FnPtr { name: var_name.clone(), curry: <_>::default(), diff --git a/src/grain/mod.rs b/src/grain/mod.rs index 30abc4ff0..762a5d949 100644 --- a/src/grain/mod.rs +++ b/src/grain/mod.rs @@ -209,6 +209,6 @@ pub use vm::{Fault, Vm}; #[cfg(feature = "internals")] pub use program::Function; -#[cfg(not(feature = "no_module"))] +#[cfg(not(feature = "internals"))] #[cfg(not(feature = "no_function"))] -pub(crate) use {program::export_functions_to_module, vm::callback_wrappers}; +pub(crate) use program::Function; diff --git a/src/grain/program.rs b/src/grain/program.rs index 277b58f8b..76e426dab 100644 --- a/src/grain/program.rs +++ b/src/grain/program.rs @@ -816,63 +816,6 @@ impl<'a> Program<'a> { } } -/// Export non-private functions into a [`Module`]. -#[cfg(not(feature = "no_function"))] -#[cfg(not(feature = "no_module"))] -pub fn export_functions_to_module( - program: &SharedProgram, - module: &mut crate::Module, - env: &Shared, -) { - program - .functions() - .iter() - .filter(|&f| match f.access { - crate::func::FnAccess::Public => true, - crate::func::FnAccess::Private => false, - }) - .for_each(|f| { - let fn_name: ImmutableString = - program.name(f.name).expect("valid function name").into(); - let params: crate::FnArgsVec = f - .params - .iter() - .map(|&idx| program.name(idx).expect("valid param name").into()) - .collect(); - - let fn_def = crate::func::ScriptFuncDef { - name: fn_name, - access: f.access, - #[cfg(not(feature = "no_object"))] - this_type: f - .this_type - .and_then(|index| program.name(index)) - .map(Into::into), - params, - body: crate::func::ScriptFuncPayload::GrainVM { - program: program.clone(), - params: f.params.iter().copied().collect(), - chunk: f.chunk, - span: crate::types::Span::new( - program.position(f.chunk.entry() as usize), - program.position(f.chunk.end() as usize), - ), - }, - #[cfg(feature = "metadata")] - comments: <_>::default(), - }; - - let hash = module.set_script_fn(fn_def); - - if let Some(crate::func::RhaiFunc::Script { env: ref mut e, .. }) = - module.get_script_fn_by_hash_mut(hash) - { - // Encapsulated environment - *e = Some(env.clone()); - } - }); -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/grain/vm/callback.rs b/src/grain/vm/callback.rs index f83b6aa2b..bd2c6b613 100644 --- a/src/grain/vm/callback.rs +++ b/src/grain/vm/callback.rs @@ -6,70 +6,21 @@ //! [`Op::Call`](crate::bytecode::Op::Call) finds one by a name *index* that //! only the compiler and the call site share, and a `FnPtr` carries a string. //! -//! So a program that hands a pointer out registers one native wrapper per -//! compiled function for the length of the run. Direct dispatch is untouched — +//! So a program that hands a pointer out stacks a module of its compiled +//! functions for the length of the run. The module looks each function up in +//! the program only when Rhai asks for it, so a run pays for the functions it +//! calls back, not for every function it has. Direct dispatch is untouched — //! this is somewhere for Rhai to look, not somewhere we look. -#[cfg(feature = "no_std")] -use std::prelude::v1::*; - #[cfg(not(feature = "no_function"))] -use crate::{ - func::{ScriptFuncDef, ScriptFuncPayload}, - grain::program::SharedProgram, - types::Span, - Engine, Module, -}; +use crate::{grain::program::SharedProgram, module::GrainFunctions, Module}; -/// A wrapper per compiled function, for Rhai to resolve a pointer against. +/// A module of the compiled functions, for Rhai to resolve a pointer against. /// -/// Built once per run rather than cached on the program: the closures hold the -/// program, so anything the program held back would be a cycle. +/// Built once per run rather than cached on the program: the definitions it +/// makes hold the program, so anything the program held back would be a +/// cycle. #[cfg(not(feature = "no_function"))] -pub(crate) fn wrappers(engine: &Engine, program: &SharedProgram) -> Module { - let mut module = Module::new(); - - // What Rhai reports as the source of a function it found here. Its own - // script library is the AST's, so this is the same string by the same - // route. - if let Some(source) = program.source() { - module.set_id(source.clone()); - } - - for function in program.functions() { - let Some(name) = program.name(function.name) else { - continue; - }; - - let params = function - .params - .iter() - .map(|&index| program.name(index).unwrap_or("")) - .map(|name| engine.get_interned_string(name)) - .collect(); - - module.set_script_fn(ScriptFuncDef { - body: ScriptFuncPayload::GrainVM { - program: program.clone(), - params: function.params.iter().copied().collect(), - chunk: function.chunk, - span: Span::new( - program.position(function.chunk.entry() as usize), - program.position(function.chunk.end() as usize), - ), - }, - name: engine.get_interned_string(name), - access: function.access, - #[cfg(not(feature = "no_object"))] - this_type: function - .this_type - .and_then(|index| program.name(index)) - .map(|name| engine.get_interned_string(name)), - params, - #[cfg(feature = "metadata")] - comments: Default::default(), - }); - } - - module +pub(crate) fn wrappers(program: &SharedProgram) -> Module { + Module::from_grain_functions(GrainFunctions::new(program)) } diff --git a/src/grain/vm/mod.rs b/src/grain/vm/mod.rs index 2a51c2e50..1836963ed 100644 --- a/src/grain/vm/mod.rs +++ b/src/grain/vm/mod.rs @@ -32,10 +32,6 @@ use crate::{types::dynamic::Variant, CallFnOptions}; mod callback; -#[cfg(not(feature = "no_function"))] -#[cfg(not(feature = "no_module"))] -pub(crate) use callback::wrappers as callback_wrappers; - use crate::grain::bytecode::{ code, names, AssignOp, Chain, Chunk, Receiver, Root, Step, StepFlags, Tail, }; @@ -875,13 +871,25 @@ impl<'e> Vm<'e> { return if program.functions().is_empty() { self.run_with(program, scope, None) } else { - let wrappers = callback::wrappers(self.engine, program).into(); + let wrappers = callback::wrappers(program).into(); self.run_with(program, scope, Some(wrappers)) }; #[cfg(feature = "no_function")] return self.run_with(program, scope, None); } + /// Like [`eval_with_callbacks`](Self::eval_with_callbacks), with a module of the + /// [`Program`]'s functions the caller made, so the caller can keep it after the run. + #[cfg(not(feature = "no_module"))] + pub(crate) fn eval_with_wrappers( + &mut self, + scope: &mut Scope, + program: &SharedProgram, + wrappers: Option, + ) -> RhaiResult { + self.run_with(program, scope, wrappers) + } + fn run_with( &mut self, program: &Program, diff --git a/src/module/grain.rs b/src/module/grain.rs new file mode 100644 index 000000000..0c79857aa --- /dev/null +++ b/src/module/grain.rs @@ -0,0 +1,258 @@ +//! Functions in a Grain program, looked up only when called. +//! +//! A [`Module`] made from a Grain [`Program`][crate::grain::Program] registers +//! none of its functions up front. A lookup finds the function in the +//! program's own table, and the first lookup of each one makes its +//! [`ScriptFuncDef`], which every later lookup shares. A function nothing +//! calls costs one empty slot. + +use super::{FnNamespace, Module, ModuleFlags}; +use crate::func::{EncapsulatedEnviron, RhaiFunc, ScriptFuncDef, ScriptFuncPayload}; +use crate::grain::{Function, SharedProgram}; +use crate::types::Span; +use crate::{FnAccess, ImmutableString, Shared}; +use once_cell::race::OnceBox; +#[cfg(feature = "no_std")] +use std::prelude::v1::*; + +/// The functions of one Grain program, as a [`Module`] holds them. +#[derive(Clone)] +pub(crate) struct GrainFunctions { + program: SharedProgram, + /// The environment a function runs in, for a module's exports. `None` + /// for the library a run of the program itself stacks, which is already + /// that environment. + env: Option>, + /// Whether private functions are left out, as a module's exports leave + /// them out. + public_only: bool, + /// Each function's definition, made on its first lookup. Shared by + /// every [`GrainFunctions`] made from the same one, so a module's library + /// and its exports make each definition once between them. + defs: Shared<[OnceBox>]>, +} + +impl GrainFunctions { + /// Every function of `program`, private ones too, with no environment. + #[must_use] + pub(crate) fn new(program: &SharedProgram) -> Self { + let defs: Vec<_> = program.functions().iter().map(|_| OnceBox::new()).collect(); + Self { + program: program.clone(), + env: None, + public_only: false, + defs: defs.into(), + } + } + + /// The public functions only, run in `env`: what a module exports. + #[cfg(not(feature = "no_module"))] + #[must_use] + pub(crate) fn exports(&self, env: Shared) -> Self { + Self { + program: self.program.clone(), + env: Some(env), + public_only: true, + defs: self.defs.clone(), + } + } + + /// The functions this holds, with their index in the program. + fn iter(&self) -> impl Iterator { + let public_only = self.public_only; + self.program + .functions() + .iter() + .enumerate() + .filter(move |(_, f)| !public_only || f.access == FnAccess::Public) + } + + /// How many functions this holds. + #[must_use] + pub(super) fn len(&self) -> usize { + self.iter().count() + } + + fn name(&self, f: &Function) -> &str { + self.program.name(f.name).unwrap_or("") + } + + /// A method declared for a receiver type is in the global namespace, as + /// [`Module::set_script_fn`] puts it. + fn namespace(f: &Function) -> FnNamespace { + #[cfg(not(feature = "no_object"))] + if f.this_type.is_some() { + return FnNamespace::Global; + } + let _ = f; + FnNamespace::Internal + } + + /// The hash [`Module::set_script_fn`] gives `f` under no namespace, or + /// [`Module::build_index`] gives it under `path`. + fn hash<'a>(&self, f: &Function, path: impl IntoIterator) -> u64 { + let hash = crate::calc_fn_hash(path, self.name(f), f.params.len()); + #[cfg(not(feature = "no_object"))] + if let Some(this_type) = f.this_type.and_then(|i| self.program.name(i)) { + return crate::calc_typed_method_hash(hash, this_type); + } + hash + } + + /// Function `i`'s definition, made on the first call. + fn def(&self, i: usize) -> &Shared { + self.defs[i].get_or_init(|| { + let program = &self.program; + let f = &program.functions()[i]; + let name = |index: u32| -> ImmutableString { program.name(index).unwrap_or("").into() }; + Box::new( + ScriptFuncDef { + body: ScriptFuncPayload::GrainVM { + program: program.clone(), + params: f.params.iter().copied().collect(), + chunk: f.chunk, + span: Span::new( + program.position(f.chunk.entry() as usize), + program.position(f.chunk.end() as usize), + ), + }, + name: name(f.name), + access: f.access, + #[cfg(not(feature = "no_object"))] + this_type: f.this_type.map(name), + params: f.params.iter().map(|&p| name(p)).collect(), + #[cfg(feature = "metadata")] + comments: <_>::default(), + } + .into(), + ) + }) + } + + fn func(&self, i: usize) -> RhaiFunc { + RhaiFunc::Script { + fn_def: self.def(i).clone(), + env: self.env.clone(), + } + } + + /// The function [`Module::set_script_fn`] would have put under `hash`. + #[must_use] + pub(super) fn get(&self, hash: u64) -> Option { + self.iter() + .find(|(_, f)| self.hash(f, None) == hash) + .map(|(i, _)| self.func(i)) + } + + /// The function [`Module::build_index`] would have put under `hash`, for + /// a module reached by `path`. + #[must_use] + pub(super) fn get_qualified(&self, path: &[&str], hash: u64) -> Option { + self.iter() + .find(|(_, f)| { + self.hash(f, path.iter().copied()) == hash + || (Self::namespace(f) == FnNamespace::Global && self.hash(f, None) == hash) + }) + .map(|(i, _)| self.func(i)) + } + + /// Whether any function is in the global namespace. + #[must_use] + pub(super) fn has_global(&self) -> bool { + self.iter() + .any(|(_, f)| Self::namespace(f) == FnNamespace::Global) + } + + /// The first function named `name` that takes `num_params`, or any + /// number of parameters when that is `None`. + #[must_use] + pub(super) fn get_named( + &self, + name: &str, + num_params: Option, + ) -> Option<&Shared> { + self.iter() + .find(|(_, f)| num_params.map_or(true, |n| f.params.len() == n) && self.name(f) == name) + .map(|(i, _)| self.def(i)) + } + + /// Every function: its namespace, access, name, parameter count and + /// definition. Makes every definition. + #[allow(dead_code)] + pub(super) fn iter_defs( + &self, + ) -> impl Iterator)> { + self.iter().map(move |(i, f)| { + ( + Self::namespace(f), + f.access, + self.name(f), + f.params.len(), + self.def(i), + ) + }) + } + + /// Register every function into `module` as [`Module::set_script_fn`] + /// does, leaving any it already has. + fn register_into(&self, module: &mut Module) { + for (i, _) in self.iter() { + let hash = module.set_script_fn(self.def(i).clone()); + if let Some(RhaiFunc::Script { env, .. }) = module.get_script_fn_by_hash_mut(hash) { + env.clone_from(&self.env); + } + } + } +} + +impl Module { + /// A [`Module`] whose functions are `functions`, looked up only when + /// called. + #[must_use] + pub(crate) fn from_grain_functions(functions: GrainFunctions) -> Self { + let mut module = Self::new(); + if let Some(source) = functions.program.source() { + module.set_id(source.clone()); + } + module.set_grain_functions(functions); + module.build_index(); + module + } + + /// Hold `functions`, looked up only when called, in place of any Grain + /// functions held before. + pub(crate) fn set_grain_functions(&mut self, functions: GrainFunctions) { + self.grain_functions = Some(functions); + self.flags + .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); + } + + /// Register the Grain functions this holds as ordinary script functions. + /// Functions already registered take precedence. + pub(super) fn register_grain_functions(&mut self) { + if let Some(grain) = self.grain_functions.take() { + let functions = self.functions.take(); + grain.register_into(self); + if let Some(functions) = functions { + self.functions + .get_or_insert_with(|| super::new_hash_map(functions.len())) + .extend(functions); + } + self.flags + .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); + } + } + + /// This [`Module`] with its Grain functions registered, cloning it if it + /// holds any. + #[must_use] + pub(super) fn with_grain_registered(&self) -> std::borrow::Cow<'_, Self> { + if self.grain_functions.is_some() { + let mut module = self.clone(); + module.register_grain_functions(); + std::borrow::Cow::Owned(module) + } else { + std::borrow::Cow::Borrowed(self) + } + } +} diff --git a/src/module/lazy.rs b/src/module/lazy.rs index f5adf0f63..f8b26bc4e 100644 --- a/src/module/lazy.rs +++ b/src/module/lazy.rs @@ -39,12 +39,23 @@ impl LazyFunctions { impl Module { /// Look up a function in plugin module manifests by hash, calculating the hash of each /// function in turn. + /// + /// Also looks up the functions of a Grain program this [`Module`] holds. #[must_use] pub(super) fn get_lazy_fn(&self, hash_native: u64) -> Option { - self.lazy_functions - .as_ref()? - .find_fn(&mut |f| f.calc_hash(&[]) == hash_native) - .map(|f| RhaiFunc::StaticPlugin { func: f.func() }) + let found = self + .lazy_functions + .as_ref() + .and_then(|f| f.find_fn(&mut |f| f.calc_hash(&[]) == hash_native)) + .map(|f| RhaiFunc::StaticPlugin { func: f.func() }); + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let found = found.or_else(|| { + self.grain_functions + .as_ref() + .and_then(|g| g.get(hash_native)) + }); + found } /// Create a new [`Module`] from a plugin module's [manifest][crate::plugin::ModuleManifest]. /// @@ -118,10 +129,15 @@ impl Module { } /// Register all functions in plugin module manifests (including sub-modules), with full /// metadata, so that they are no longer looked up only when called. + /// + /// The functions of a Grain program this [`Module`] holds are registered too. pub fn register_lazy_functions(&mut self) -> &mut Self { if !self.has_lazy_functions() { return self; } + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + self.register_grain_functions(); if let Some(lazy_functions) = self.lazy_functions.take() { // Functions registered normally take precedence over functions in manifests let functions = self.functions.take(); @@ -143,12 +159,14 @@ impl Module { .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); self } - /// Does this [`Module`] (or any sub-module) contain functions in plugin module manifests, - /// which are looked up only when called? + /// Does this [`Module`] (or any sub-module) contain functions in plugin module manifests, or + /// the functions of a Grain program, which are looked up only when called? #[inline] #[must_use] pub fn has_lazy_functions(&self) -> bool { - self.lazy_functions.is_some() || self.modules.values().any(|m| m.has_lazy_functions()) + self.lazy_functions.is_some() + || self.has_grain_functions() + || self.modules.values().any(|m| m.has_lazy_functions()) } /// Get this [`Module`] with all functions in plugin module manifests registered, cloning it /// if necessary. @@ -177,13 +195,25 @@ impl Module { module: &'a Module, path: &mut crate::StaticVec<&'a str>, hash: u64, - ) -> Option<&'static FnManifestEntry> { + ) -> Option { // Functions in the module take precedence over those in sub-modules - let found = module.lazy_functions.as_ref().and_then(|f| { - f.find_fn(&mut |f| { - f.calc_hash(path) == hash - || (f.namespace() == FnNamespace::Global && f.calc_hash(&[]) == hash) + let found = module + .lazy_functions + .as_ref() + .and_then(|f| { + f.find_fn(&mut |f| { + f.calc_hash(path) == hash + || (f.namespace() == FnNamespace::Global && f.calc_hash(&[]) == hash) + }) }) + .map(|f| RhaiFunc::StaticPlugin { func: f.func() }); + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let found = found.or_else(|| { + module + .grain_functions + .as_ref() + .and_then(|g| g.get_qualified(path, hash)) }); found.or_else(|| { @@ -202,6 +232,6 @@ impl Module { let path = &mut crate::StaticVec::new(); path.push(""); - find_fn(self, path, hash_qualified_fn).map(|f| RhaiFunc::StaticPlugin { func: f.func() }) + find_fn(self, path, hash_qualified_fn) } } diff --git a/src/module/mod.rs b/src/module/mod.rs index 68e5e5054..f4a0f1802 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -1,5 +1,11 @@ //! Module defining external-loaded modules for Rhai. +#[cfg(feature = "grain")] +#[cfg(not(feature = "no_function"))] +mod grain; +#[cfg(feature = "grain")] +#[cfg(not(feature = "no_function"))] +pub(crate) use grain::GrainFunctions; mod lazy; mod namespace; /// Module containing all built-in [module resolvers][ModuleResolver]. @@ -667,6 +673,10 @@ pub struct Module { flags: ModuleFlags, /// Functions in plugin module manifests, looked up only when called. Never empty if [`Some`]. lazy_functions: Option, + /// Functions in a Grain program, looked up only when called. + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + grain_functions: Option, } impl Default for Module { @@ -799,6 +809,9 @@ impl Module { all_type_iterators: BTreeMap::new(), flags: ModuleFlags::INDEXED, lazy_functions: None, + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + grain_functions: None, } } @@ -934,6 +947,11 @@ impl Module { self.type_iterators.clear(); self.all_type_iterators.clear(); self.lazy_functions = None; + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + { + self.grain_functions = None; + } self.flags .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); } @@ -1121,6 +1139,18 @@ impl Module { .map_or(true, StraightHashMap::is_empty) && self.all_type_iterators.is_empty() && self.lazy_functions.is_none() + && !self.has_grain_functions() + } + + /// Does this [`Module`] hold the functions of a Grain program, looked up only when called? + #[inline(always)] + #[must_use] + pub(crate) const fn has_grain_functions(&self) -> bool { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + return self.grain_functions.is_some(); + #[cfg(any(not(feature = "grain"), feature = "no_function"))] + return false; } /// Is the [`Module`] indexed? @@ -1388,13 +1418,34 @@ impl Module { name: impl AsRef, num_params: usize, ) -> Option<&Shared> { - self.functions.as_ref().and_then(|lib| { - let name = name.as_ref(); + self.get_script_fn_named(name.as_ref(), Some(num_params)) + } + /// Get a shared reference to the first script-defined function in the [`Module`] named + /// `name` that takes `num_params` parameters, or any number when that is [`None`]. + #[cfg(not(feature = "no_function"))] + #[must_use] + pub(crate) fn get_script_fn_named( + &self, + name: &str, + num_params: Option, + ) -> Option<&Shared> { + let found = self.functions.as_ref().and_then(|lib| { lib.values() - .find(|(_, f)| f.num_params == num_params && f.name == name) + .find(|(f, m)| { + f.is_script() + && num_params.map_or(true, |n| m.num_params == n) + && m.name == name + }) .and_then(|(f, _)| f.get_script_fn_def()) - }) + }); + #[cfg(feature = "grain")] + let found = found.or_else(|| { + self.grain_functions + .as_ref() + .and_then(|g| g.get_named(name, num_params)) + }); + found } /// Get a mutable reference to the underlying [`BTreeMap`] of sub-modules, @@ -2012,6 +2063,13 @@ impl Module { /// The other [`Module`] is _consumed_ to merge into this [`Module`]. #[inline] pub fn combine(&mut self, other: Self) -> &mut Self { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let other = { + let mut other = other; + other.register_grain_functions(); + other + }; if let Some(lazy_functions) = other.lazy_functions { self.push_manifests(lazy_functions.manifests.into_vec()); } @@ -2046,6 +2104,13 @@ impl Module { /// Sub-modules are flattened onto the root [`Module`], with higher level overriding lower level. #[inline] pub fn combine_flatten(&mut self, other: Self) -> &mut Self { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let other = { + let mut other = other; + other.register_grain_functions(); + other + }; for m in other.modules.into_values() { self.combine_flatten(shared_take_or_clone(m)); } @@ -2081,6 +2146,9 @@ impl Module { /// Only items not existing in this [`Module`] are added. #[inline] pub fn fill_with(&mut self, other: &Self) -> &mut Self { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let other = &*other.with_grain_registered(); // Existing manifests take precedence, so these go before them if let Some(ref lazy_functions) = other.lazy_functions { let existing = self.lazy_functions.take().map(|f| f.manifests.into_vec()); @@ -2142,6 +2210,9 @@ impl Module { other: &Self, _filter: impl Fn(FnNamespace, FnAccess, bool, &str, usize) -> bool + Copy, ) -> &mut Self { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let other = &*other.with_grain_registered(); for (k, v) in &other.modules { let mut m = Self::new(); m.merge_filtered(v, _filter); @@ -2208,6 +2279,8 @@ impl Module { &mut self, filter: impl Fn(FnNamespace, FnAccess, &str, usize) -> bool, ) -> &mut Self { + #[cfg(feature = "grain")] + self.register_grain_functions(); self.functions = std::mem::take(&mut self.functions).map(|m| { m.into_iter() .filter(|(.., (f, m))| { @@ -2236,11 +2309,22 @@ impl Module { pub fn count(&self) -> (usize, usize, usize) { ( self.variables.len(), - self.functions.as_ref().map_or(0, StraightHashMap::len), + self.functions.as_ref().map_or(0, StraightHashMap::len) + self.num_grain_functions(), self.type_iterators.len(), ) } + /// The number of functions of a Grain program this [`Module`] holds. + #[inline(always)] + #[must_use] + fn num_grain_functions(&self) -> usize { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + return self.grain_functions.as_ref().map_or(0, |g| g.len()); + #[cfg(any(not(feature = "grain"), feature = "no_function"))] + return 0; + } + /// Get an iterator to the sub-modules in the [`Module`]. #[inline(always)] pub fn iter_sub_modules(&self) -> impl Iterator { @@ -2303,7 +2387,7 @@ impl Module { &Shared, ), > + '_ { - self.iter_fn().filter(|(f, _)| f.is_script()).map(|(f, m)| { + let registered = self.iter_fn().filter(|(f, _)| f.is_script()).map(|(f, m)| { ( m.namespace, m.access, @@ -2311,7 +2395,11 @@ impl Module { m.num_params, f.get_script_fn_def().expect("`ScriptFuncDef`"), ) - }) + }); + #[cfg(feature = "grain")] + return registered.chain(self.grain_functions.iter().flat_map(|g| g.iter_defs())); + #[cfg(not(feature = "grain"))] + return registered; } /// _(internals)_ Get an iterator over all script-defined functions in the [`Module`]. @@ -2568,12 +2656,22 @@ impl Module { #[cfg(not(feature = "no_function"))] let orig_constants = std::mem::take(&mut global.constants); + // The program's functions, looked up only when called: by its own top-level run, by + // its functions through the module's environment, and through the module's exports. + #[cfg(not(feature = "no_function"))] + let functions = grain::GrainFunctions::new(&program); + #[cfg(not(feature = "no_function"))] + let wrappers: Option = (!program.functions().is_empty()) + .then(|| Self::from_grain_functions(functions.clone()).into()); + #[cfg(feature = "no_function")] + let wrappers = None; + // Run the Grain program let mut vm = crate::grain::Vm::with_global_state( engine, std::mem::replace(global, engine.new_global_runtime_state()), ); - let result = vm.eval_with_callbacks(scope, &program); + let result = vm.eval_with_wrappers(scope, &program, wrappers.clone()); *global = vm.into_global_state(); // Create new module @@ -2611,9 +2709,8 @@ impl Module { #[cfg(not(feature = "no_function"))] let mut lib: crate::StaticVec<_> = program.lib().cloned().into_iter().collect(); #[cfg(not(feature = "no_function"))] - if !program.functions().is_empty() { - let wrappers = crate::grain::callback_wrappers(engine, &program); - lib.push(wrappers.into()); + if let Some(wrappers) = wrappers { + lib.push(wrappers); } #[cfg(not(feature = "no_function"))] @@ -2669,7 +2766,9 @@ impl Module { } #[cfg(not(feature = "no_function"))] - crate::grain::export_functions_to_module(&program, &mut module, &env); + if !program.functions().is_empty() { + module.set_grain_functions(functions.exports(env)); + } module.id = program.source().cloned(); @@ -2743,6 +2842,16 @@ impl Module { }); } + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + if module + .grain_functions + .as_ref() + .map_or(false, |g| g.has_global()) + { + contains_indexed_global_functions = true; + } + // Index all functions for (&hash, (f, m)) in module.functions.iter().flatten() { match m.namespace { diff --git a/tests/grain/allocation.rs b/tests/grain/allocation.rs index a29eaf275..3e12d1e14 100644 --- a/tests/grain/allocation.rs +++ b/tests/grain/allocation.rs @@ -310,6 +310,33 @@ fn allocation_footprint() { // scaling with length. assert!(ratios[2] > ratios[1] && ratios[1] > ratios[0], "the saving must grow with the program, got {ratios:?}",); + // What importing a module costs. Its functions are looked up in the program + // when they are called, so until then a module of many functions costs one + // empty slot per function more than a module of one. + #[cfg(not(feature = "no_module"))] + { + let module_of = |n: usize| { + let source: String = (0..n).map(|i| format!("fn f{i}(a, b) {{ a + b }}\n")).collect(); + let program = Compiler::new().compile(&engine.compile(&source).expect("must compile")).into_shared(); + let module = measure(|| rhai::Module::eval_grain_as_new(rhai::Scope::new(), program.clone(), &engine).expect("must eval")); + (module.bytes, module.value) + }; + let (one, _one) = module_of(1); + let (many, many_module) = module_of(64); + println!("\nmodule of 1 function {one:>8} bytes\nmodule of 64 {many:>8} bytes"); + assert!(many - one <= 63 * 16, "a module must not register its functions up front: {one} bytes for 1, {many} for 64"); + + // Calling one makes that one's definition, and only that one's. + let mut called = Engine::new(); + let mut resolver = rhai::module_resolvers::StaticModuleResolver::new(); + resolver.insert("m", many_module); + called.set_module_resolver(resolver); + let first = measure(|| called.eval::(r#"import "m" as m; m::f3(1, 2)"#).expect("must call")); + assert_eq!(first.value, 3); + let again = measure(|| called.eval::(r#"import "m" as m; m::f3(1, 2)"#).expect("must call")); + assert!(again.bytes <= first.bytes, "a second call must not make another definition: {} then {}", first.bytes, again.bytes); + } + // The instrument has to be working before anything can lean on it. assert!(ast.count > 0 && ast.bytes > 0, "the counters saw nothing; the global allocator is not installed",); assert!(ast.peak >= ast.bytes, "peak ({}) below retained ({}) means peak tracking is broken", ast.peak, ast.bytes,); diff --git a/tests/grain/module.rs b/tests/grain/module.rs index 8b75f3d03..a9da189ec 100644 --- a/tests/grain/module.rs +++ b/tests/grain/module.rs @@ -206,3 +206,99 @@ fn grain_fn_ptr_callbacks_and_closures() { let items: Vec = res.into_iter().map(|d| d.cast::()).collect(); assert_eq!(items, vec![6, 7, 8]); } + +/// A Grain module looks its functions up in the program when they are called, +/// so a lookup has to find what registering them would have found: the public +/// ones under their qualified names, and none of the private ones. +#[test] +#[cfg(not(feature = "no_function"))] +fn functions_are_found_without_being_registered() { + let mut engine = corpus::engine(); + let ast = engine + .compile( + r#" + fn one() { 1 } + fn add(a, b) { a + b } + private fn hidden() { 2 } + "#, + ) + .expect("must compile"); + + let program = Compiler::new().compile(&ast); + let module = Module::eval_grain_as_new(Scope::new(), program, &engine).expect("must eval grain module"); + assert_eq!(module.count().1, 2, "the private function is not exported"); + assert!(!module.is_empty()); + + let mut resolver = StaticModuleResolver::new(); + resolver.insert("m", module); + engine.set_module_resolver(resolver); + + assert_eq!(engine.eval::(r#"import "m" as m; m::one() + m::add(m::one(), 2)"#).unwrap(), 4); + // The same call again resolves to the definition the first one made. + assert_eq!(engine.eval::(r#"import "m" as m; let n = 0; for i in 0..3 { n += m::add(i, 1) } n"#).unwrap(), 6); + assert!(engine.run(r#"import "m" as m; m::add(1)"#).is_err(), "the parameter count is part of the name"); + assert!(engine.run(r#"import "m" as m; m::hidden()"#).is_err()); +} + +/// A method declared for a receiver type is in the global namespace, so it is +/// found without the module's name. +#[test] +#[cfg(not(feature = "no_function"))] +#[cfg(not(feature = "no_object"))] +fn a_typed_method_is_found_globally() { + let mut engine = corpus::engine(); + let ast = engine.compile(r#"fn int.doubled() { this * 2 }"#).expect("must compile"); + + let program = Compiler::new().compile(&ast); + let module = Module::eval_grain_as_new(Scope::new(), program, &engine).expect("must eval grain module"); + + let mut resolver = StaticModuleResolver::new(); + resolver.insert("m", module); + engine.set_module_resolver(resolver); + + assert_eq!(engine.eval::(r#"import "m" as m; let x = 21; x.doubled()"#).unwrap(), 42); +} + +/// Combining a Grain module into another registers its functions there, where +/// nothing looks them up in the program any more. +#[test] +#[cfg(not(feature = "no_function"))] +fn a_combined_module_keeps_its_functions() { + let mut engine = corpus::engine(); + let ast = engine.compile(r#"fn seven() { 7 }"#).expect("must compile"); + + let program = Compiler::new().compile(&ast); + let grain = Module::eval_grain_as_new(Scope::new(), program, &engine).expect("must eval grain module"); + + let mut combined = Module::new(); + combined.set_native_fn("eight", || Ok(8 as INT)); + combined.combine(grain); + combined.build_index(); + assert_eq!(combined.count().1, 2); + + let mut resolver = StaticModuleResolver::new(); + resolver.insert("m", combined); + engine.set_module_resolver(resolver); + + assert_eq!(engine.eval::(r#"import "m" as m; m::seven() + m::eight()"#).unwrap(), 15); +} + +/// Metadata lists a Grain module's functions, as it lists those of a plugin +/// module's manifest. +#[test] +#[cfg(feature = "metadata")] +#[cfg(not(feature = "no_function"))] +fn metadata_lists_the_functions() { + let engine = corpus::engine(); + let ast = engine.compile(r#"fn greet(name) { name }"#).expect("must compile"); + + let program = Compiler::new().compile(&ast); + let mut module = Module::eval_grain_as_new(Scope::new(), program, &engine).expect("must eval grain module"); + assert!(module.has_lazy_functions()); + + module.register_lazy_functions(); + assert!(!module.has_lazy_functions()); + let signatures: Vec = module.gen_fn_signatures_with_mapper(Into::into).collect(); + assert_eq!(signatures.len(), 1, "{signatures:?}"); + assert!(signatures[0].starts_with("greet(name"), "{signatures:?}"); +} From 4b276cffb9a21284b6c89bfec3e05b6356053e46 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Wed, 21 Oct 2026 17:47:18 -0700 Subject: [PATCH 02/10] Verify a chunk with a state per jump target, not per code byte --- src/grain/bytecode/verify.rs | 88 ++++++++++++++++++++++++++++++------ tests/grain/allocation.rs | 11 +++++ 2 files changed, 85 insertions(+), 14 deletions(-) diff --git a/src/grain/bytecode/verify.rs b/src/grain/bytecode/verify.rs index 547961f86..9d7e2d63e 100644 --- a/src/grain/bytecode/verify.rs +++ b/src/grain/bytecode/verify.rs @@ -252,10 +252,11 @@ pub fn verify( // // Over the whole buffer at once, because every chunk shares it and an // instruction boundary is a property of the bytes, not of who runs them. - let mut starts = vec![false; code.len() + 1]; + // One bit per byte. + let mut starts = vec![0u64; code.len() / 64 + 1]; let mut at = 0usize; while at < code.len() { - starts[at] = true; + starts[at / 64] |= 1 << (at % 64); let width = code::width(code, at).ok_or(VerifyError::Undecodable { at })?; check_indices(at, code, pools)?; at += width; @@ -296,13 +297,52 @@ struct State { handlers: usize, } +/// Whether an instruction starts at `at`, by pass one's bits. +fn starts_at(starts: &[u64], at: usize) -> bool { + starts + .get(at / 64) + .map_or(false, |bits| bits & (1 << (at % 64)) != 0) +} + +/// Every offset `op` can jump to, onto `out`. +fn jump_targets(op: &Op, pools: &Pools, out: &mut Vec) { + match *op { + Op::Jump(target) + | Op::JumpIfFalse { target } + | Op::JumpIfTrue { target } + | Op::SkipIfNotUnit { target } + | Op::PushHandler { target, .. } + | Op::IterNext { exit: target, .. } => out.push(target), + Op::Switch(index) => { + if let Some(table) = pools.switches.get(index as usize) { + out.extend( + table + .cases + .iter() + .flat_map(|cases| cases.values()) + .map(|&(target, _)| target) + .chain(table.ranges.iter().map(|range| range.target)) + .chain(core::iter::once(table.default)), + ); + } + } + _ => {} + } +} + /// Walk one chunk's reachable instructions, checking that every path into an /// instruction agrees on the stack depth. +/// +/// Only an instruction something jumps to can be reached from more than one +/// place: any other is reached only by falling through from the one before +/// it, so its state follows from that one's. Recording the state at the entry +/// and at every jump target is therefore enough to see every merge, and it +/// costs memory in the chunk's branches rather than its length. fn verify_chunk( caps: Caps, code: &[u8], chunk: &Chunk, - starts: &[bool], + starts: &[u64], pools: &Pools, ) -> Result { let (entry, end) = (chunk.entry() as usize, chunk.end() as usize); @@ -314,7 +354,18 @@ fn verify_chunk( }); } - let mut depth_at: Vec> = vec![None; code.len()]; + let mut targets = vec![entry as u32]; + let mut at = entry; + while at < end { + let op = code::decode(code, at).ok_or(VerifyError::Undecodable { at })?; + jump_targets(&op, pools, &mut targets); + at += code::width(code, at).ok_or(VerifyError::Undecodable { at })?; + } + targets.sort_unstable(); + targets.dedup(); + + // The state every path into each of `targets` arrives in, once one has. + let mut seen: Vec> = vec![None; targets.len()]; let mut work_list = vec![(entry, State::default())]; let mut high_water = 0usize; @@ -325,16 +376,18 @@ fn verify_chunk( // Merge point: either this is the first visit, or every earlier path // has to have arrived in the same state. - match depth_at[at] { - Some(seen) if seen == state => continue, - Some(seen) => { - return Err(VerifyError::DepthConflict { - at, - expected: seen.operands, - found: state.operands, - }) + if let Ok(slot) = targets.binary_search(&(at as u32)) { + match seen[slot] { + Some(seen) if seen == state => continue, + Some(seen) => { + return Err(VerifyError::DepthConflict { + at, + expected: seen.operands, + found: state.operands, + }) + } + None => seen[slot] = Some(state), } - None => depth_at[at] = Some(state), } let depth = state.operands; @@ -404,12 +457,19 @@ fn verify_chunk( target: target as u32, }); } - if !starts[target] { + if !starts_at(starts, target) { return Err(VerifyError::JumpIntoAnInstruction { at, target: target as u32, }); } + // Every edge followed here has to be one `jump_targets` found, or + // the instruction it lands on would go unchecked as a merge point. + // Refused rather than trusted, so a disagreement between the two + // can only reject a chunk, never pass one. + if targets.binary_search(&(target as u32)).is_err() { + return Err(VerifyError::Undecodable { at }); + } work_list.push((target, state)); Ok(()) }; diff --git a/tests/grain/allocation.rs b/tests/grain/allocation.rs index 3e12d1e14..6d1a5b133 100644 --- a/tests/grain/allocation.rs +++ b/tests/grain/allocation.rs @@ -299,6 +299,17 @@ fn allocation_footprint() { "the loaded chunk must point into the artifact, not into a copy of it", ); + // The verifier keeps a state per jump target, not per byte, so even + // this loop, a branch every few instructions, peaks at a few bytes + // per byte of code. + let verified = measure(|| loaded.value.verify().expect("must verify")); + assert!( + (verified.peak as usize) < code.len() * 32 / 10, + "verifying {} bytes of code peaked at {} bytes: a state per byte, 32 bytes each on this host, is back", + code.len(), + verified.peak, + ); + let ratio = tree.bytes as f64 / loaded.bytes as f64; ratios.push(ratio); println!("{:>7} {:>9} {:>9} {:>9} {:>8} {:>6.1}x", source.len(), stripped.len(), tree.bytes, loaded.bytes, loaded.count, ratio,); From cc791626537af7125cf74be43759dfc4ae49655c Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Wed, 21 Oct 2026 17:52:21 -0700 Subject: [PATCH 03/10] Keep a position per instruction that has one, not one per code byte --- src/grain/bytecode/positions.rs | 63 ++++++++++++++++++++++----------- tests/grain/allocation.rs | 8 +++++ 2 files changed, 51 insertions(+), 20 deletions(-) diff --git a/src/grain/bytecode/positions.rs b/src/grain/bytecode/positions.rs index d04148924..0e09cd0f4 100644 --- a/src/grain/bytecode/positions.rs +++ b/src/grain/bytecode/positions.rs @@ -15,12 +15,15 @@ use std::prelude::v1::*; /// instructions that had one. What it buys is that they can now be removed. /// `tests/format.rs` measures how much that removes. /// -/// In memory this is dense, because the lookup is not always cold. A loop -/// back-edge passes a position to `track_operation` on every iteration, and the -/// built-in operator path builds a `NativeCallContext` around one — both on the -/// hot path, both needing it only if something goes wrong. Indexing an array is -/// what makes that free. The compact delta form in [`pos`](crate::grain::pos) -/// is the wire form, expanded once at load. +/// In memory this holds only the instructions that have a position, in +/// address order. The lookup is not always cold — a loop back-edge passes a +/// position to `track_operation` on every iteration, and the built-in operator +/// path builds a `NativeCallContext` around one — so it is a binary search, +/// a dozen comparisons for a large script. An entry per code byte would make +/// it an index, at four bytes of heap for every byte of code, for a table that +/// is mostly empty: that is what a small device cannot spare. The compact +/// delta form in [`pos`](crate::grain::pos) is the wire form, expanded once at +/// load. /// /// [`Positions::Stripped`] is not a degraded mode to apologize for: it is what /// a device runs. Errors come back carrying an instruction address instead of a @@ -30,18 +33,32 @@ pub enum Positions { /// The table was never written, or was stripped before shipping. #[default] Stripped, - /// One entry per instruction, so a lookup is an index. - Dense(Box<[Position]>), + /// The address and position of every instruction that has one, by + /// address. + Sparse(Box<[(u32, Position)]>), } impl Positions { - /// Build from one position per instruction, dropping the table entirely if - /// none of them say anything. + /// Build from one position per code byte, keeping only those that say + /// something, and dropping the table entirely if none of them do. + #[cfg(any(test, not(feature = "no_ast")))] pub(crate) fn dense(positions: Vec) -> Self { - if positions.iter().all(|pos| pos.is_none()) { + Self::sparse( + positions + .into_iter() + .enumerate() + .filter(|(_, pos)| !pos.is_none()) + .map(|(pc, pos)| (pc as u32, pos)) + .collect(), + ) + } + + /// Build from the instructions that have a position, in address order. + fn sparse(entries: Vec<(u32, Position)>) -> Self { + if entries.is_empty() { return Self::Stripped; } - Self::Dense(positions.into_boxed_slice()) + Self::Sparse(entries.into_boxed_slice()) } /// The position recorded for `pc`, or `Position::NONE`. @@ -53,7 +70,10 @@ impl Positions { pub fn get(&self, pc: usize) -> Position { match self { Self::Stripped => Position::NONE, - Self::Dense(positions) => positions.get(pc).copied().unwrap_or(Position::NONE), + Self::Sparse(entries) => u32::try_from(pc) + .ok() + .and_then(|pc| entries.binary_search_by_key(&pc, |&(at, _)| at).ok()) + .map_or(Position::NONE, |i| entries[i].1), } } @@ -69,14 +89,14 @@ impl Positions { /// Instructions with no position are skipped, which is most of them. #[must_use] pub fn to_table(&self) -> Vec { - let Self::Dense(positions) = self else { + let Self::Sparse(entries) = self else { return crate::grain::pos::encode(core::iter::empty()); }; - crate::grain::pos::encode(positions.iter().enumerate().filter_map(|(pc, pos)| { + crate::grain::pos::encode(entries.iter().filter_map(|&(pc, pos)| { let line = pos.line()?; Some(( - pc as u32, + pc, Site { line: line as u32, column: pos.position().unwrap_or(0) as u32, @@ -85,7 +105,7 @@ impl Positions { })) } - /// Expand a compact table back to one entry per code byte. + /// Expand a compact table back to an entry per instruction it names. /// /// Every address has to name the *start* of an instruction in `code`. /// Counting how many landed in range is not enough: another program's table @@ -107,12 +127,15 @@ impl Positions { // Walked, not indexed: only an address the walk arrives at begins an // instruction. One landing mid-instruction is as wrong as one past the // end, and reads as a plausible position. - let mut positions = vec![Position::NONE; code.len()]; + let mut entries = Vec::with_capacity(count as usize); let mut matched = 0usize; let mut at = 0usize; while at < code.len() { if let Some(site) = crate::grain::pos::resolve(table, at as u32) { - positions[at] = site_to_position(site); + let pos = site_to_position(site); + if !pos.is_none() { + entries.push((at as u32, pos)); + } matched += 1; } match crate::grain::bytecode::code::width(code, at) { @@ -131,7 +154,7 @@ impl Positions { }); } - Ok(Self::dense(positions)) + Ok(Self::sparse(entries)) } } diff --git a/tests/grain/allocation.rs b/tests/grain/allocation.rs index 6d1a5b133..77c7f9b48 100644 --- a/tests/grain/allocation.rs +++ b/tests/grain/allocation.rs @@ -310,6 +310,14 @@ fn allocation_footprint() { verified.peak, ); + let full = program.write().expect("must be writable"); + let with_positions = measure(|| rhai::grain::Program::read(&full).expect("must load")); + // The position table keeps an entry per instruction that has a + // position, not one per code byte: under a byte of heap per byte of + // code, where a dense table cost four. + let table = (with_positions.bytes - loaded.bytes) as usize; + assert!(table < code.len(), "the position table for {} bytes of code retains {table} bytes", code.len()); + let ratio = tree.bytes as f64 / loaded.bytes as f64; ratios.push(ratio); println!("{:>7} {:>9} {:>9} {:>9} {:>8} {:>6.1}x", source.len(), stripped.len(), tree.bytes, loaded.bytes, loaded.count, ratio,); From 692aa5875a3309a9449b8ad538c20544ec5ede68 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Thu, 22 Oct 2026 14:39:55 -0700 Subject: [PATCH 04/10] Hold a Grain program's functions as a lazy source beside plugin manifests Grain functions and plugin manifests are now both sources in one LazyFunctions list, so a Module has one field for functions looked up only when called, and the Grain-only module helpers exist only with the grain feature. --- src/module/grain.rs | 65 +++++++++++------ src/module/lazy.rs | 170 +++++++++++++++++++++++++++----------------- src/module/mod.rs | 81 ++++++--------------- 3 files changed, 171 insertions(+), 145 deletions(-) diff --git a/src/module/grain.rs b/src/module/grain.rs index 0c79857aa..ec862c7c4 100644 --- a/src/module/grain.rs +++ b/src/module/grain.rs @@ -6,7 +6,8 @@ //! [`ScriptFuncDef`], which every later lookup shares. A function nothing //! calls costs one empty slot. -use super::{FnNamespace, Module, ModuleFlags}; +use super::lazy::LazySource; +use super::{FnNamespace, Module}; use crate::func::{EncapsulatedEnviron, RhaiFunc, ScriptFuncDef, ScriptFuncPayload}; use crate::grain::{Function, SharedProgram}; use crate::types::Span; @@ -194,8 +195,8 @@ impl GrainFunctions { } /// Register every function into `module` as [`Module::set_script_fn`] - /// does, leaving any it already has. - fn register_into(&self, module: &mut Module) { + /// does, in place of any it already has. + pub(super) fn register_into(&self, module: &mut Module) { for (i, _) in self.iter() { let hash = module.set_script_fn(self.def(i).clone()); if let Some(RhaiFunc::Script { env, .. }) = module.get_script_fn_by_hash_mut(hash) { @@ -206,6 +207,25 @@ impl GrainFunctions { } impl Module { + /// The functions of Grain programs this [`Module`] holds, looked up only + /// when called, from the lowest precedence to the highest. + pub(super) fn grain_functions(&self) -> impl DoubleEndedIterator { + self.lazy_functions.iter().flat_map(|f| f.grain()) + } + + /// Does this [`Module`] hold the functions of a Grain program, looked up + /// only when called? + #[must_use] + fn has_grain_functions(&self) -> bool { + self.grain_functions().next().is_some() + } + + /// The number of functions of Grain programs this [`Module`] holds. + #[must_use] + pub(super) fn num_grain_functions(&self) -> usize { + self.grain_functions().map(GrainFunctions::len).sum() + } + /// A [`Module`] whose functions are `functions`, looked up only when /// called. #[must_use] @@ -214,40 +234,45 @@ impl Module { if let Some(source) = functions.program.source() { module.set_id(source.clone()); } - module.set_grain_functions(functions); + module.add_grain_functions(functions); module.build_index(); module } - /// Hold `functions`, looked up only when called, in place of any Grain - /// functions held before. - pub(crate) fn set_grain_functions(&mut self, functions: GrainFunctions) { - self.grain_functions = Some(functions); - self.flags - .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); + /// Hold `functions`, looked up only when called, taking precedence over + /// any functions looked up only when called held before. + pub(crate) fn add_grain_functions(&mut self, functions: GrainFunctions) { + self.push_lazy_sources([LazySource::Grain(functions)]); } /// Register the Grain functions this holds as ordinary script functions. /// Functions already registered take precedence. pub(super) fn register_grain_functions(&mut self) { - if let Some(grain) = self.grain_functions.take() { - let functions = self.functions.take(); - grain.register_into(self); - if let Some(functions) = functions { - self.functions - .get_or_insert_with(|| super::new_hash_map(functions.len())) - .extend(functions); + if !self.has_grain_functions() { + return; + } + let sources = self.lazy_functions.take().map(|f| f.sources.into_vec()); + let functions = self.functions.take(); + let mut rest = Vec::new(); + for source in sources.into_iter().flatten() { + match source { + LazySource::Grain(grain) => grain.register_into(self), + LazySource::Manifest(..) => rest.push(source), } - self.flags - .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); } + if let Some(functions) = functions { + self.functions + .get_or_insert_with(|| super::new_hash_map(functions.len())) + .extend(functions); + } + self.push_lazy_sources(rest); } /// This [`Module`] with its Grain functions registered, cloning it if it /// holds any. #[must_use] pub(super) fn with_grain_registered(&self) -> std::borrow::Cow<'_, Self> { - if self.grain_functions.is_some() { + if self.has_grain_functions() { let mut module = self.clone(); module.register_grain_functions(); std::borrow::Cow::Owned(module) diff --git a/src/module/lazy.rs b/src/module/lazy.rs index f8b26bc4e..a08b39763 100644 --- a/src/module/lazy.rs +++ b/src/module/lazy.rs @@ -1,5 +1,9 @@ -//! Functions in plugin modules, looked up from their manifests only when called. +//! Functions looked up only when called: those in plugin module manifests, and those of Grain +//! programs. +#[cfg(feature = "grain")] +#[cfg(not(feature = "no_function"))] +use super::GrainFunctions; use super::{new_hash_map, FnNamespace, Module, ModuleFlags}; use crate::func::RhaiFunc; use crate::plugin::{FnManifestEntry, ModuleManifest}; @@ -7,11 +11,22 @@ use std::borrow::Cow; #[cfg(feature = "no_std")] use std::prelude::v1::*; -/// Plugin module manifests whose functions are looked up only when called. -#[derive(Debug, Clone)] +/// A source of functions looked up only when called. +#[derive(Clone)] +pub(super) enum LazySource { + /// A plugin module manifest, and whether its sub-modules are flattened. + Manifest(&'static ModuleManifest, bool), + /// The functions of a Grain program. + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + Grain(GrainFunctions), +} + +/// Sources of functions looked up only when called. +#[derive(Clone)] pub(super) struct LazyFunctions { - /// Manifests, and whether their sub-modules are flattened. Later manifests take precedence. - pub(super) manifests: Box<[(&'static ModuleManifest, bool)]>, + /// Later sources take precedence. Never empty. + pub(super) sources: Box<[LazySource]>, } impl LazyFunctions { @@ -20,42 +35,82 @@ impl LazyFunctions { /// /// Registering each function in turn leaves the one with the highest precedence in place. pub(super) fn for_each_fn(&self, f: &mut dyn FnMut(&'static FnManifestEntry)) { - for &(manifest, flatten) in &*self.manifests { - manifest.for_each_fn(flatten, f); + for source in &*self.sources { + match *source { + LazySource::Manifest(manifest, flatten) => manifest.for_each_fn(flatten, f), + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + LazySource::Grain(..) => (), + } } } - /// Find the function in the manifests with the highest precedence that matches a predicate. + /// The Grain functions, from the lowest precedence to the highest. + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + pub(super) fn grain(&self) -> impl DoubleEndedIterator { + self.sources.iter().filter_map(|source| match source { + LazySource::Grain(functions) => Some(functions), + LazySource::Manifest(..) => None, + }) + } + /// Find the function with the highest precedence that matches. + /// + /// `manifest` matches a function in a manifest; `grain` looks up a function of a Grain program. fn find_fn( &self, - f: &mut dyn FnMut(&FnManifestEntry) -> bool, - ) -> Option<&'static FnManifestEntry> { - self.manifests - .iter() - .rev() - .find_map(|&(manifest, flatten)| manifest.find_fn(flatten, f)) + manifest: &mut dyn FnMut(&FnManifestEntry) -> bool, + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + grain: &mut dyn FnMut(&GrainFunctions) -> Option, + ) -> Option { + self.sources.iter().rev().find_map(|source| match *source { + LazySource::Manifest(m, flatten) => m + .find_fn(flatten, manifest) + .map(|f| RhaiFunc::StaticPlugin { func: f.func() }), + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + LazySource::Grain(ref functions) => grain(functions), + }) + } + /// Is any function in the global namespace? + pub(super) fn has_global(&self) -> bool { + self.sources.iter().any(|source| match *source { + LazySource::Manifest(manifest, flatten) => manifest + .find_fn(flatten, &mut |f| f.namespace() == FnNamespace::Global) + .is_some(), + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + LazySource::Grain(ref functions) => functions.has_global(), + }) + } + /// Register every function into `module`, from the lowest precedence to the highest. + fn register_into(&self, module: &mut Module) { + for source in &*self.sources { + match *source { + LazySource::Manifest(manifest, flatten) => { + manifest.for_each_fn(flatten, &mut |f| f.register_into(module)); + } + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + LazySource::Grain(ref functions) => functions.register_into(module), + } + } } } impl Module { - /// Look up a function in plugin module manifests by hash, calculating the hash of each - /// function in turn. - /// - /// Also looks up the functions of a Grain program this [`Module`] holds. + /// Look up a function in plugin module manifests, or among the functions of a Grain program, + /// by hash, calculating the hash of each function in turn. #[must_use] pub(super) fn get_lazy_fn(&self, hash_native: u64) -> Option { - let found = self - .lazy_functions - .as_ref() - .and_then(|f| f.find_fn(&mut |f| f.calc_hash(&[]) == hash_native)) - .map(|f| RhaiFunc::StaticPlugin { func: f.func() }); - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - let found = found.or_else(|| { - self.grain_functions - .as_ref() - .and_then(|g| g.get(hash_native)) - }); - found + self.lazy_functions.as_ref().and_then(|f| { + f.find_fn( + &mut |f| f.calc_hash(&[]) == hash_native, + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + &mut |g| g.get(hash_native), + ) + }) } /// Create a new [`Module`] from a plugin module's [manifest][crate::plugin::ModuleManifest]. /// @@ -107,21 +162,18 @@ impl Module { self.dynamic_functions_filter.mark(hash_script); } }); - self.push_manifests([(manifest, flatten)]); + self.push_lazy_sources([LazySource::Manifest(manifest, flatten)]); } - /// Add manifests whose functions are looked up only when called, taking precedence over - /// existing ones, without updating the filter on functions with [`Dynamic`][crate::Dynamic] parameters. - pub(super) fn push_manifests( - &mut self, - manifests: impl IntoIterator, - ) { + /// Add sources of functions looked up only when called, taking precedence over existing ones, + /// without updating the filter on functions with [`Dynamic`][crate::Dynamic] parameters. + pub(super) fn push_lazy_sources(&mut self, sources: impl IntoIterator) { let existing = self .lazy_functions .take() - .map_or_else(Vec::new, |f| f.manifests.into_vec()); - let manifests: Box<[_]> = existing.into_iter().chain(manifests).collect(); - if !manifests.is_empty() { - self.lazy_functions = Some(LazyFunctions { manifests }); + .map_or_else(Vec::new, |f| f.sources.into_vec()); + let sources: Box<[_]> = existing.into_iter().chain(sources).collect(); + if !sources.is_empty() { + self.lazy_functions = Some(LazyFunctions { sources }); } self.flags @@ -135,13 +187,10 @@ impl Module { if !self.has_lazy_functions() { return self; } - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - self.register_grain_functions(); if let Some(lazy_functions) = self.lazy_functions.take() { - // Functions registered normally take precedence over functions in manifests + // Functions registered normally take precedence over functions looked up only when called let functions = self.functions.take(); - lazy_functions.for_each_fn(&mut |f| f.register_into(self)); + lazy_functions.register_into(self); if let Some(functions) = functions { self.functions .get_or_insert_with(|| new_hash_map(functions.len())) @@ -164,9 +213,7 @@ impl Module { #[inline] #[must_use] pub fn has_lazy_functions(&self) -> bool { - self.lazy_functions.is_some() - || self.has_grain_functions() - || self.modules.values().any(|m| m.has_lazy_functions()) + self.lazy_functions.is_some() || self.modules.values().any(|m| m.has_lazy_functions()) } /// Get this [`Module`] with all functions in plugin module manifests registered, cloning it /// if necessary. @@ -197,23 +244,16 @@ impl Module { hash: u64, ) -> Option { // Functions in the module take precedence over those in sub-modules - let found = module - .lazy_functions - .as_ref() - .and_then(|f| { - f.find_fn(&mut |f| { + let found = module.lazy_functions.as_ref().and_then(|f| { + f.find_fn( + &mut |f| { f.calc_hash(path) == hash || (f.namespace() == FnNamespace::Global && f.calc_hash(&[]) == hash) - }) - }) - .map(|f| RhaiFunc::StaticPlugin { func: f.func() }); - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - let found = found.or_else(|| { - module - .grain_functions - .as_ref() - .and_then(|g| g.get_qualified(path, hash)) + }, + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + &mut |g| g.get_qualified(path, hash), + ) }); found.or_else(|| { diff --git a/src/module/mod.rs b/src/module/mod.rs index f4a0f1802..65dd2a7ac 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -671,12 +671,8 @@ pub struct Module { all_type_iterators: BTreeMap>, /// Flags. flags: ModuleFlags, - /// Functions in plugin module manifests, looked up only when called. Never empty if [`Some`]. + /// Functions in plugin module manifests and Grain programs, looked up only when called. lazy_functions: Option, - /// Functions in a Grain program, looked up only when called. - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - grain_functions: Option, } impl Default for Module { @@ -809,9 +805,6 @@ impl Module { all_type_iterators: BTreeMap::new(), flags: ModuleFlags::INDEXED, lazy_functions: None, - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - grain_functions: None, } } @@ -947,11 +940,6 @@ impl Module { self.type_iterators.clear(); self.all_type_iterators.clear(); self.lazy_functions = None; - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - { - self.grain_functions = None; - } self.flags .remove(ModuleFlags::INDEXED | ModuleFlags::INDEXED_GLOBAL_FUNCTIONS); } @@ -1139,18 +1127,6 @@ impl Module { .map_or(true, StraightHashMap::is_empty) && self.all_type_iterators.is_empty() && self.lazy_functions.is_none() - && !self.has_grain_functions() - } - - /// Does this [`Module`] hold the functions of a Grain program, looked up only when called? - #[inline(always)] - #[must_use] - pub(crate) const fn has_grain_functions(&self) -> bool { - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - return self.grain_functions.is_some(); - #[cfg(any(not(feature = "grain"), feature = "no_function"))] - return false; } /// Is the [`Module`] indexed? @@ -1441,9 +1417,9 @@ impl Module { }); #[cfg(feature = "grain")] let found = found.or_else(|| { - self.grain_functions - .as_ref() - .and_then(|g| g.get_named(name, num_params)) + self.grain_functions() + .rev() + .find_map(|g| g.get_named(name, num_params)) }); found } @@ -2071,7 +2047,7 @@ impl Module { other }; if let Some(lazy_functions) = other.lazy_functions { - self.push_manifests(lazy_functions.manifests.into_vec()); + self.push_lazy_sources(lazy_functions.sources.into_vec()); } self.modules.extend(other.modules); self.variables.extend(other.variables); @@ -2115,7 +2091,7 @@ impl Module { self.combine_flatten(shared_take_or_clone(m)); } if let Some(lazy_functions) = other.lazy_functions { - self.push_manifests(lazy_functions.manifests.into_vec()); + self.push_lazy_sources(lazy_functions.sources.into_vec()); } self.variables.extend(other.variables); match self.functions { @@ -2151,9 +2127,9 @@ impl Module { let other = &*other.with_grain_registered(); // Existing manifests take precedence, so these go before them if let Some(ref lazy_functions) = other.lazy_functions { - let existing = self.lazy_functions.take().map(|f| f.manifests.into_vec()); - self.push_manifests(lazy_functions.manifests.iter().copied()); - self.push_manifests(existing.into_iter().flatten()); + let existing = self.lazy_functions.take().map(|f| f.sources.into_vec()); + self.push_lazy_sources(lazy_functions.sources.iter().cloned()); + self.push_lazy_sources(existing.into_iter().flatten()); } for (k, v) in &other.modules { if !self.modules.contains_key(k) { @@ -2307,24 +2283,18 @@ impl Module { #[inline(always)] #[must_use] pub fn count(&self) -> (usize, usize, usize) { + let num_functions = self.functions.as_ref().map_or(0, StraightHashMap::len); + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + let num_functions = num_functions + self.num_grain_functions(); + ( self.variables.len(), - self.functions.as_ref().map_or(0, StraightHashMap::len) + self.num_grain_functions(), + num_functions, self.type_iterators.len(), ) } - /// The number of functions of a Grain program this [`Module`] holds. - #[inline(always)] - #[must_use] - fn num_grain_functions(&self) -> usize { - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - return self.grain_functions.as_ref().map_or(0, |g| g.len()); - #[cfg(any(not(feature = "grain"), feature = "no_function"))] - return 0; - } - /// Get an iterator to the sub-modules in the [`Module`]. #[inline(always)] pub fn iter_sub_modules(&self) -> impl Iterator { @@ -2397,7 +2367,7 @@ impl Module { ) }); #[cfg(feature = "grain")] - return registered.chain(self.grain_functions.iter().flat_map(|g| g.iter_defs())); + return registered.chain(self.grain_functions().flat_map(|g| g.iter_defs())); #[cfg(not(feature = "grain"))] return registered; } @@ -2767,7 +2737,7 @@ impl Module { #[cfg(not(feature = "no_function"))] if !program.functions().is_empty() { - module.set_grain_functions(functions.exports(env)); + module.add_grain_functions(functions.exports(env)); } module.id = program.source().cloned(); @@ -2833,21 +2803,12 @@ impl Module { type_iterators.insert(type_id, func.clone()); } - // Functions in manifests are indexed on first lookup, but may be in the global namespace - if let Some(ref lazy_functions) = module.lazy_functions { - lazy_functions.for_each_fn(&mut |f| { - if f.namespace() == FnNamespace::Global { - contains_indexed_global_functions = true; - } - }); - } - - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] + // Functions looked up only when called are indexed on first lookup, but may be in the + // global namespace if module - .grain_functions + .lazy_functions .as_ref() - .map_or(false, |g| g.has_global()) + .map_or(false, lazy::LazyFunctions::has_global) { contains_indexed_global_functions = true; } From 4dfb038b0f76a96e652858a4e864a1a8a7bd3375 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Thu, 22 Oct 2026 14:48:55 -0700 Subject: [PATCH 05/10] Allow the position table two bytes per code byte under debugging Every statement emits a positioned marker under debugging, so the table outgrows the code it describes on a loop-heavy script. --- tests/grain/allocation.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/grain/allocation.rs b/tests/grain/allocation.rs index 77c7f9b48..e67b5f99f 100644 --- a/tests/grain/allocation.rs +++ b/tests/grain/allocation.rs @@ -314,9 +314,11 @@ fn allocation_footprint() { let with_positions = measure(|| rhai::grain::Program::read(&full).expect("must load")); // The position table keeps an entry per instruction that has a // position, not one per code byte: under a byte of heap per byte of - // code, where a dense table cost four. + // code, where a dense table cost four. Under `debugging` every + // statement also emits a positioned marker, so allow two. let table = (with_positions.bytes - loaded.bytes) as usize; - assert!(table < code.len(), "the position table for {} bytes of code retains {table} bytes", code.len()); + let bytes_per_code_byte = if cfg!(feature = "debugging") { 2 } else { 1 }; + assert!(table < code.len() * bytes_per_code_byte, "the position table for {} bytes of code retains {table} bytes", code.len()); let ratio = tree.bytes as f64 / loaded.bytes as f64; ratios.push(ratio); From c36155931bdda299cd98173286003e4bb17aad5a Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Fri, 23 Oct 2026 08:15:29 -0700 Subject: [PATCH 06/10] Build Module::get_script_fn_by_hash_mut only with grain Only Grain calls it, so it was dead code without the feature. --- src/module/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/module/mod.rs b/src/module/mod.rs index 65dd2a7ac..e5203c496 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -1370,7 +1370,7 @@ impl Module { } /// Get a mutable reference to a scripted function in the [`Module`] based on its hash. - /// Exported under the `internals` feature only. + #[cfg(feature = "grain")] #[cfg(not(feature = "no_function"))] #[inline(always)] #[must_use] From 642feafcd429a55034b133b51eb98554f79c7998 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Fri, 23 Oct 2026 08:15:30 -0700 Subject: [PATCH 07/10] Keep a combined or filled-in Grain module's functions lazy combine and combine_flatten remove the registered functions the incoming Grain functions take precedence over, so lookups fall through to them; fill_with skips the incoming functions this module's Grain functions already hold. Neither registers Grain functions up front any more. --- src/module/grain.rs | 25 +++++++++++++++++++++++++ src/module/mod.rs | 40 +++++++++++++++++----------------------- tests/grain/module.rs | 35 +++++++++++++++++++++++++++++++---- 3 files changed, 73 insertions(+), 27 deletions(-) diff --git a/src/module/grain.rs b/src/module/grain.rs index ec862c7c4..7b2ce44af 100644 --- a/src/module/grain.rs +++ b/src/module/grain.rs @@ -137,6 +137,11 @@ impl GrainFunctions { } } + /// The hashes [`Module::set_script_fn`] gives the functions. + pub(super) fn hashes(&self) -> impl Iterator + '_ { + self.iter().map(|(_, f)| self.hash(f, None)) + } + /// The function [`Module::set_script_fn`] would have put under `hash`. #[must_use] pub(super) fn get(&self, hash: u64) -> Option { @@ -226,6 +231,26 @@ impl Module { self.grain_functions().map(GrainFunctions::len).sum() } + /// Does this [`Module`] hold a Grain function under `hash`, looked up + /// only when called? + #[must_use] + pub(super) fn has_grain_fn(&self, hash: u64) -> bool { + self.grain_functions() + .any(|g| g.hashes().any(|h| h == hash)) + } + + /// Remove the functions registered in this [`Module`] that the Grain + /// functions `other` holds would take precedence over once combined + /// into it, so that lookups fall through to them. + pub(super) fn remove_fns_hidden_by_grain(&mut self, other: &Self) { + let Some(ref mut functions) = self.functions else { + return; + }; + for hash in other.grain_functions().flat_map(GrainFunctions::hashes) { + functions.remove(&hash); + } + } + /// A [`Module`] whose functions are `functions`, looked up only when /// called. #[must_use] diff --git a/src/module/mod.rs b/src/module/mod.rs index e5203c496..fa8b4f4e8 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -2041,11 +2041,7 @@ impl Module { pub fn combine(&mut self, other: Self) -> &mut Self { #[cfg(feature = "grain")] #[cfg(not(feature = "no_function"))] - let other = { - let mut other = other; - other.register_grain_functions(); - other - }; + self.remove_fns_hidden_by_grain(&other); if let Some(lazy_functions) = other.lazy_functions { self.push_lazy_sources(lazy_functions.sources.into_vec()); } @@ -2079,17 +2075,13 @@ impl Module { /// The other [`Module`] is _consumed_ to merge into this [`Module`]. /// Sub-modules are flattened onto the root [`Module`], with higher level overriding lower level. #[inline] - pub fn combine_flatten(&mut self, other: Self) -> &mut Self { - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - let other = { - let mut other = other; - other.register_grain_functions(); - other - }; - for m in other.modules.into_values() { + pub fn combine_flatten(&mut self, mut other: Self) -> &mut Self { + for m in std::mem::take(&mut other.modules).into_values() { self.combine_flatten(shared_take_or_clone(m)); } + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + self.remove_fns_hidden_by_grain(&other); if let Some(lazy_functions) = other.lazy_functions { self.push_lazy_sources(lazy_functions.sources.into_vec()); } @@ -2122,15 +2114,6 @@ impl Module { /// Only items not existing in this [`Module`] are added. #[inline] pub fn fill_with(&mut self, other: &Self) -> &mut Self { - #[cfg(feature = "grain")] - #[cfg(not(feature = "no_function"))] - let other = &*other.with_grain_registered(); - // Existing manifests take precedence, so these go before them - if let Some(ref lazy_functions) = other.lazy_functions { - let existing = self.lazy_functions.take().map(|f| f.sources.into_vec()); - self.push_lazy_sources(lazy_functions.sources.iter().cloned()); - self.push_lazy_sources(existing.into_iter().flatten()); - } for (k, v) in &other.modules { if !self.modules.contains_key(k) { self.modules.insert(k.clone(), v.clone()); @@ -2145,6 +2128,11 @@ impl Module { let others_len = functions.len(); for (&k, f) in functions { + #[cfg(feature = "grain")] + #[cfg(not(feature = "no_function"))] + if self.has_grain_fn(k) { + continue; + } let map = self .functions .get_or_insert_with(|| new_hash_map(FN_MAP_SIZE)); @@ -2152,6 +2140,12 @@ impl Module { map.entry(k).or_insert_with(|| f.clone()); } } + // Existing sources take precedence, so these go before them + if let Some(ref lazy_functions) = other.lazy_functions { + let existing = self.lazy_functions.take().map(|f| f.sources.into_vec()); + self.push_lazy_sources(lazy_functions.sources.iter().cloned()); + self.push_lazy_sources(existing.into_iter().flatten()); + } self.dynamic_functions_filter += &other.dynamic_functions_filter; for (&k, v) in &other.type_iterators { self.type_iterators.entry(k).or_insert_with(|| v.clone()); diff --git a/tests/grain/module.rs b/tests/grain/module.rs index a9da189ec..d27e5b90f 100644 --- a/tests/grain/module.rs +++ b/tests/grain/module.rs @@ -259,8 +259,8 @@ fn a_typed_method_is_found_globally() { assert_eq!(engine.eval::(r#"import "m" as m; let x = 21; x.doubled()"#).unwrap(), 42); } -/// Combining a Grain module into another registers its functions there, where -/// nothing looks them up in the program any more. +/// Combining a Grain module into another keeps its functions looked up only +/// when called, taking precedence over functions already there. #[test] #[cfg(not(feature = "no_function"))] fn a_combined_module_keeps_its_functions() { @@ -270,17 +270,44 @@ fn a_combined_module_keeps_its_functions() { let program = Compiler::new().compile(&ast); let grain = Module::eval_grain_as_new(Scope::new(), program, &engine).expect("must eval grain module"); - let mut combined = Module::new(); + let walker = engine.compile(r#"fn seven() { 0 } fn six() { 6 }"#).expect("must compile"); + let mut combined = Module::eval_ast_as_new(Scope::new(), &walker, &engine).expect("must eval module"); combined.set_native_fn("eight", || Ok(8 as INT)); combined.combine(grain); + assert!(combined.has_lazy_functions()); combined.build_index(); - assert_eq!(combined.count().1, 2); + assert_eq!(combined.count().1, 3); let mut resolver = StaticModuleResolver::new(); resolver.insert("m", combined); engine.set_module_resolver(resolver); assert_eq!(engine.eval::(r#"import "m" as m; m::seven() + m::eight()"#).unwrap(), 15); + assert_eq!(engine.eval::(r#"import "m" as m; m::six()"#).unwrap(), 6); +} + +/// Filling a module with a Grain module adds only the functions it does not +/// have, whether registered or looked up only when called. +#[test] +#[cfg(not(feature = "no_function"))] +fn a_filled_module_keeps_its_own_functions() { + let mut engine = corpus::engine(); + let program = |engine: &rhai::Engine, script: &str| Compiler::new().compile(&engine.compile(script).expect("must compile")); + + let mut filled = Module::eval_grain_as_new(Scope::new(), program(&engine, r#"fn seven() { 7 }"#), &engine).expect("must eval grain module"); + let other = engine.compile(r#"fn seven() { 0 } fn six() { 0 }"#).expect("must compile"); + filled.fill_with(&Module::eval_ast_as_new(Scope::new(), &other, &engine).expect("must eval module")); + let other = Module::eval_grain_as_new(Scope::new(), program(&engine, r#"fn six() { 6 } fn five() { 5 }"#), &engine).expect("must eval grain module"); + filled.fill_with(&other); + filled.build_index(); + + let mut resolver = StaticModuleResolver::new(); + resolver.insert("m", filled); + engine.set_module_resolver(resolver); + + assert_eq!(engine.eval::(r#"import "m" as m; m::seven()"#).unwrap(), 7); + assert_eq!(engine.eval::(r#"import "m" as m; m::six()"#).unwrap(), 0); + assert_eq!(engine.eval::(r#"import "m" as m; m::five()"#).unwrap(), 5); } /// Metadata lists a Grain module's functions, as it lists those of a plugin From c3dfa813c4e3ca1b78fb5fd2ced7da1e59713402 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Fri, 23 Oct 2026 08:18:19 -0700 Subject: [PATCH 08/10] Apply the filter when merging into a module with no functions Module::merge_filtered cloned every function across when the target had none, so AST::clone_functions_only_filtered kept every function. --- CHANGELOG.md | 1 + src/module/mod.rs | 9 ++++----- tests/functions.rs | 18 ++++++++++++++++++ 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5056cdff3..d470cb9a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ Bug fixes * Assigning a constant into a container variable (array, map etc.) now works properly without making the value read-only (thanks [`@cuervolu`](https://github.com/cuervolu) [`#1187`](https://github.com/rhaiscript/rhai/pull/1187)). * Under `sync`, a method call on a shared variable with a closure argument that captures the same variable now raises `ErrorDataRace`, as it does without `sync`, instead of panicking while resolving the function. `Dynamic::is_locked` now returns `true` under `sync` when the value cannot be read, which includes being locked for write by the current thread (thanks [`@yuvalrakavy`](https://github.com/yuvalrakavy) [`#1190`](https://github.com/rhaiscript/rhai/pull/1190)). * A panic inside a native Rust function that takes its first argument by value, called with a variable first argument, now unwinds normally and can be caught by `catch_unwind`. Previously the copied first argument was restored only on the normal-return path, so the unwind tripped an internal assertion and the process aborted (thanks [`@yuvalrakavy`](https://github.com/yuvalrakavy) [`#1189`](https://github.com/rhaiscript/rhai/pull/1189)). +* `AST::clone_functions_only_filtered` now applies its filter, as does `AST::merge_filtered` when the first `AST` has no functions. Previously every function was kept (thanks [`@ImTheSquid`](https://github.com/ImTheSquid) [`#1192`](https://github.com/rhaiscript/rhai/pull/1192)). Rhai Grain is now at feature parity with the AST interpreter ------------------------------------------------------------ diff --git a/src/module/mod.rs b/src/module/mod.rs index fa8b4f4e8..7e4b55553 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -2209,17 +2209,16 @@ impl Module { } if let Some(ref functions) = other.functions { - match self.functions { - Some(ref mut m) => m.extend( + self.functions + .get_or_insert_with(|| new_hash_map(functions.len())) + .extend( functions .iter() .filter(|(.., (f, m))| { _filter(m.namespace, m.access, f.is_script(), &m.name, m.num_params) }) .map(|(&k, f)| (k, f.clone())), - ), - None => self.functions.clone_from(&other.functions), - } + ); } self.dynamic_functions_filter += &other.dynamic_functions_filter; diff --git a/tests/functions.rs b/tests/functions.rs index bbc39d065..3c9a98491 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -838,3 +838,21 @@ fn test_functions_resolution_in_loop() { assert_eq!(engine.eval::("let t = 0; for i in 0..10 { t += calc(i); } t").unwrap(), 55); } + +#[test] +fn test_functions_clone_filtered() { + let engine = Engine::new(); + let ast = engine.compile("fn keep(x) { x } fn drop(x) { x } fn keep() { 0 }").unwrap(); + + let names = |ast: &rhai::AST| { + let mut names: Vec<_> = ast.iter_functions().map(|f| format!("{}/{}", f.name, f.params.len())).collect(); + names.sort(); + names + }; + + let filtered = ast.clone_functions_only_filtered(|_, _, _, name, _| name == "keep"); + assert_eq!(names(&filtered), ["keep/0", "keep/1"]); + + let merged = engine.compile("1").unwrap().merge_filtered(&ast, |_, _, _, _, params| params == 1); + assert_eq!(names(&merged), ["drop/1", "keep/1"]); +} From 24a0891b045141ff70bda9807b543b01fee4289a Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Fri, 23 Oct 2026 13:45:57 -0700 Subject: [PATCH 09/10] Find lazy Grain functions on the linked-script fast path, hash them once A function pointer linked to a script function short-circuits through Module::get_script_fn_by_hash, which saw only registered functions, so a callback into a Grain closure always took the slow path. It now looks in the Grain functions too, and returns None for a non-script function instead of panicking. Each lookup also hashed every function's name again, so a callback whose operators miss the Grain module paid for that on every call. The hashes are now made on the first lookup and shared by a module's library and its exports: 8 bytes per function. --- src/module/grain.rs | 34 ++++++++++++++++++++++++++++++---- src/module/mod.rs | 29 +++++++++++++++++++++++------ src/types/fn_ptr.rs | 8 +------- 3 files changed, 54 insertions(+), 17 deletions(-) diff --git a/src/module/grain.rs b/src/module/grain.rs index 7b2ce44af..e8635c916 100644 --- a/src/module/grain.rs +++ b/src/module/grain.rs @@ -31,6 +31,9 @@ pub(crate) struct GrainFunctions { /// every [`GrainFunctions`] made from the same one, so a module's library /// and its exports make each definition once between them. defs: Shared<[OnceBox>]>, + /// Each function's hash under no namespace, made on the first lookup and + /// shared as `defs` is. + hashes: Shared>>, } impl GrainFunctions { @@ -43,6 +46,7 @@ impl GrainFunctions { env: None, public_only: false, defs: defs.into(), + hashes: Shared::new(OnceBox::new()), } } @@ -55,6 +59,7 @@ impl GrainFunctions { env: Some(env), public_only: true, defs: self.defs.clone(), + hashes: self.hashes.clone(), } } @@ -100,6 +105,14 @@ impl GrainFunctions { hash } + /// Function `i`'s hash under no namespace, as [`hash`](Self::hash) gives it. + fn unqualified_hash(&self, i: usize) -> u64 { + self.hashes.get_or_init(|| { + let functions = self.program.functions(); + Box::new(functions.iter().map(|f| self.hash(f, None)).collect()) + })[i] + } + /// Function `i`'s definition, made on the first call. fn def(&self, i: usize) -> &Shared { self.defs[i].get_or_init(|| { @@ -139,25 +152,38 @@ impl GrainFunctions { /// The hashes [`Module::set_script_fn`] gives the functions. pub(super) fn hashes(&self) -> impl Iterator + '_ { - self.iter().map(|(_, f)| self.hash(f, None)) + self.iter().map(|(i, _)| self.unqualified_hash(i)) } /// The function [`Module::set_script_fn`] would have put under `hash`. #[must_use] pub(super) fn get(&self, hash: u64) -> Option { self.iter() - .find(|(_, f)| self.hash(f, None) == hash) + .find(|&(i, _)| self.unqualified_hash(i) == hash) .map(|(i, _)| self.func(i)) } + /// The definition of the function [`Module::set_script_fn`] would have + /// put under `hash`, and the environment it runs in. + #[must_use] + pub(super) fn get_script( + &self, + hash: u64, + ) -> Option<(&Shared, Option<&Shared>)> { + self.iter() + .find(|&(i, _)| self.unqualified_hash(i) == hash) + .map(|(i, _)| (self.def(i), self.env.as_ref())) + } + /// The function [`Module::build_index`] would have put under `hash`, for /// a module reached by `path`. #[must_use] pub(super) fn get_qualified(&self, path: &[&str], hash: u64) -> Option { self.iter() - .find(|(_, f)| { + .find(|&(i, f)| { self.hash(f, path.iter().copied()) == hash - || (Self::namespace(f) == FnNamespace::Global && self.hash(f, None) == hash) + || (Self::namespace(f) == FnNamespace::Global + && self.unqualified_hash(i) == hash) }) .map(|(i, _)| self.func(i)) } diff --git a/src/module/mod.rs b/src/module/mod.rs index 7e4b55553..8c775087d 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -1357,16 +1357,33 @@ impl Module { hash_script } - /// Get a reference to a scripted function in the [`Module`] based on its hash. - /// Exported under the `internals` feature only. + /// Get a reference to a scripted function in the [`Module`], and the environment it runs in, + /// based on its hash. #[cfg(not(feature = "no_function"))] - #[inline(always)] + #[inline] #[must_use] - pub(crate) fn get_script_fn_by_hash(&self, hash: u64) -> Option<&crate::func::RhaiFunc> { - self.functions + pub(crate) fn get_script_fn_by_hash( + &self, + hash: u64, + ) -> Option<( + &Shared, + Option<&Shared>, + )> { + let found = self + .functions .as_ref() .and_then(|f| f.get(&hash)) - .map(|(f, _)| f) + .and_then(|(f, _)| match f { + RhaiFunc::Script { fn_def, env } => Some((fn_def, env.as_ref())), + _ => None, + }); + #[cfg(feature = "grain")] + let found = found.or_else(|| { + self.grain_functions() + .rev() + .find_map(|g| g.get_script(hash)) + }); + found } /// Get a mutable reference to a scripted function in the [`Module`] based on its hash. diff --git a/src/types/fn_ptr.rs b/src/types/fn_ptr.rs index 6f96bd445..376b3f376 100644 --- a/src/types/fn_ptr.rs +++ b/src/types/fn_ptr.rs @@ -70,13 +70,7 @@ impl FnPtrType { Self::Script { num_params, hash } if *num_params == num_args && !global.lib.is_empty() => { - global.lib[0].get_script_fn_by_hash(*hash).map(|f| match f { - crate::func::RhaiFunc::Script { - ref fn_def, - ref env, - } => (fn_def, env.as_ref()), - _ => unreachable!(), - }) + global.lib[0].get_script_fn_by_hash(*hash) } _ => None, } From ef9a8dd7cdfcb50e7b98b9272ddbb9cdb0370502 Mon Sep 17 00:00:00 2001 From: Jack Hogan Date: Mon, 26 Oct 2026 10:09:41 -0700 Subject: [PATCH 10/10] Register callback functions up front unless lazy callbacks are asked for Vm::eval_with_callbacks registers every compiled function again, as it did before, since a typical script has too few functions for looking them up only when called to save much. Vm::eval_with_lazy_callbacks keeps them in the program until one is called back. --- CHANGELOG.md | 1 + src/grain/vm/callback.rs | 21 ++++++++++++++------- src/grain/vm/mod.rs | 32 +++++++++++++++++++++++++++++++- tests/grain/callback.rs | 17 ++++++++++++----- 4 files changed, 58 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d470cb9a1..dec1ce4a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,7 @@ Lazy loading of plugin module functions * Lazy-load plugin's are supported by both the `AST` interpreter and Rhai Grain. * Use the new `Module::from_manifest` (the lazy `exported_module!`) or `Module::combine_manifest` (the lazy `combine_with_exported_module!`) API to create a lazy-loaded module or add the manifest of lazy-loaded functions to a module. * `Module::register_lazy_functions` registers all lazy-loaded functions up-front instead, for the same speed as before. +* A `Module` created by `Module::eval_grain_as_new` now looks up the functions of the Rhai Grain program only when called. The new `Vm::eval_with_lazy_callbacks` does the same for functions called back from native Rust functions, while `Vm::eval_with_callbacks` still registers every function up-front ([`#1192`](https://github.com/rhaiscript/rhai/pull/1192)). Feature flag to replace `smartstring` with `compact_str` -------------------------------------------------------- diff --git a/src/grain/vm/callback.rs b/src/grain/vm/callback.rs index bd2c6b613..d7263a135 100644 --- a/src/grain/vm/callback.rs +++ b/src/grain/vm/callback.rs @@ -7,20 +7,27 @@ //! only the compiler and the call site share, and a `FnPtr` carries a string. //! //! So a program that hands a pointer out stacks a module of its compiled -//! functions for the length of the run. The module looks each function up in -//! the program only when Rhai asks for it, so a run pays for the functions it -//! calls back, not for every function it has. Direct dispatch is untouched — -//! this is somewhere for Rhai to look, not somewhere we look. +//! functions for the length of the run. Either every function is registered +//! in it up front, or, lazily, the module looks each one up in the program +//! only when Rhai asks for it, so a run pays for the functions it calls back, +//! not for every function it has. Direct dispatch is untouched — this is +//! somewhere for Rhai to look, not somewhere we look. #[cfg(not(feature = "no_function"))] use crate::{grain::program::SharedProgram, module::GrainFunctions, Module}; -/// A module of the compiled functions, for Rhai to resolve a pointer against. +/// A module of the compiled functions, for Rhai to resolve a pointer against, +/// with every function registered up front unless `lazy`. /// /// Built once per run rather than cached on the program: the definitions it /// makes hold the program, so anything the program held back would be a /// cycle. #[cfg(not(feature = "no_function"))] -pub(crate) fn wrappers(program: &SharedProgram) -> Module { - Module::from_grain_functions(GrainFunctions::new(program)) +pub(crate) fn wrappers(program: &SharedProgram, lazy: bool) -> Module { + let mut module = Module::from_grain_functions(GrainFunctions::new(program)); + if !lazy { + module.register_lazy_functions(); + module.build_index(); + } + module } diff --git a/src/grain/vm/mod.rs b/src/grain/vm/mod.rs index 1836963ed..5e23e661c 100644 --- a/src/grain/vm/mod.rs +++ b/src/grain/vm/mod.rs @@ -858,6 +858,10 @@ impl<'e> Vm<'e> { /// the [function pointer][crate::FnPtr] simply fails to resolve, as /// [`EvalAltResult::ErrorFunctionNotFound`], at the point the native tries /// to call it. + /// + /// Every compiled function is registered up front, for the run. + /// [`eval_with_lazy_callbacks`](Self::eval_with_lazy_callbacks) registers + /// none, and looks one up only when it is called. // // Named `eval_` rather than `run_` because it yields the [`Program`]'s value; // Rhai has no `Engine` method to mirror here, so the crate's own rule is @@ -871,7 +875,33 @@ impl<'e> Vm<'e> { return if program.functions().is_empty() { self.run_with(program, scope, None) } else { - let wrappers = callback::wrappers(program).into(); + let wrappers = callback::wrappers(program, false).into(); + self.run_with(program, scope, Some(wrappers)) + }; + #[cfg(feature = "no_function")] + return self.run_with(program, scope, None); + } + + /// Same as [`eval_with_callbacks`](Self::eval_with_callbacks), except that no + /// compiled function is registered up front: each one is looked up in the + /// [`Program`] only when it is called back. + /// + /// ### Cost + /// + /// This saves registering every function on each run, which matters for a + /// [`Program`] with many functions that calls back few of them. Each + /// function resolution that misses the [`Program`]'s functions, such as an + /// operator inside a callback, scans them instead of looking up one hash. + pub fn eval_with_lazy_callbacks( + &mut self, + scope: &mut Scope, + program: &SharedProgram, + ) -> RhaiResult { + #[cfg(not(feature = "no_function"))] + return if program.functions().is_empty() { + self.run_with(program, scope, None) + } else { + let wrappers = callback::wrappers(program, true).into(); self.run_with(program, scope, Some(wrappers)) }; #[cfg(feature = "no_function")] diff --git a/tests/grain/callback.rs b/tests/grain/callback.rs index 0d8ffdb41..b825c3bcc 100644 --- a/tests/grain/callback.rs +++ b/tests/grain/callback.rs @@ -15,12 +15,17 @@ use rhai::{Dynamic, Engine, Scope}; /// Run a source through the VM with the callback wrappers installed. fn run(engine: &Engine, source: &str) -> Result { + run_with(engine, source, false) +} + +/// Run a source through the VM with the callback wrappers installed, looked +/// up only when called if `lazy`. +fn run_with(engine: &Engine, source: &str, lazy: bool) -> Result { let ast = engine.compile(source).map_err(|err| format!("{err:?}"))?; let program = Compiler::new().compile(&ast).into_shared(); - Vm::new(engine) - .eval_with_callbacks(&mut Scope::new(), &program) - .map(|value| format!("{value:?}")) - .map_err(|err| format!("{err:?}")) + let mut vm = Vm::new(engine); + let result = if lazy { vm.eval_with_lazy_callbacks(&mut Scope::new(), &program) } else { vm.eval_with_callbacks(&mut Scope::new(), &program) }; + result.map(|value| format!("{value:?}")).map_err(|err| format!("{err:?}")) } /// The same source through the walker, which is what the answer has to be. @@ -31,7 +36,9 @@ fn walk(engine: &Engine, source: &str) -> Result { fn agree(source: &str) { let engine = corpus::engine(); - assert_eq!(walk(&engine, source), run(&engine, source), "{source}"); + let expected = walk(&engine, source); + assert_eq!(expected, run_with(&engine, source, false), "{source}"); + assert_eq!(expected, run_with(&engine, source, true), "lazy: {source}"); } /// The array is bound to a variable in every case here, and that is not