diff --git a/src/api.rs b/src/api.rs index ef23947d..d1fea857 100644 --- a/src/api.rs +++ b/src/api.rs @@ -36,7 +36,7 @@ impl<'e> Engine<'e> { args, }; - self.ext_functions.insert(spec, f); + self.functions.insert(spec, f); } /// Register a custom type for use with the `Engine`. @@ -848,13 +848,7 @@ impl<'e> Engine<'e> { functions: impl IntoIterator>, ) { for f in functions.into_iter() { - match self - .script_functions - .binary_search_by(|fn_def| fn_def.compare(&f.name, f.params.len())) - { - Ok(n) => self.script_functions[n] = f.clone(), - Err(n) => self.script_functions.insert(n, f.clone()), - } + self.fn_lib.add_or_replace_function(f.clone()); } } diff --git a/src/engine.rs b/src/engine.rs index fa928fce..59a6ff3c 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -12,6 +12,7 @@ use crate::stdlib::{ any::{type_name, TypeId}, borrow::Cow, boxed::Box, + cmp::Ordering, collections::HashMap, format, iter::once, @@ -53,6 +54,64 @@ pub struct FnSpec<'a> { pub args: Option>, } +/// A type that holds a library of script-defined functions. +/// +/// Since script-defined functions have `Dynamic` parameters, functions with the same name +/// and number of parameters are considered equivalent. +/// +/// Since the key is a combination of the function name (a String) plus the number of parameters, +/// we cannot use a `HashMap` because we don't want to clone the function name string just +/// to search for it. +/// +/// So instead this is implemented as a sorted list and binary searched. +#[derive(Debug)] +pub struct FunctionsLib(Vec>); + +impl FnDef { + /// Function to order two FnDef records, for binary search. + pub fn compare(&self, name: &str, params_len: usize) -> Ordering { + // First order by name + match self.name.as_str().cmp(name) { + // Then by number of parameters + Ordering::Equal => self.params.len().cmp(¶ms_len), + order => order, + } + } +} + +impl FunctionsLib { + /// Create a new `FunctionsLib`. + pub fn new() -> Self { + FunctionsLib(Vec::new()) + } + /// Clear the `FunctionsLib`. + pub fn clear(&mut self) { + self.0.clear(); + } + /// Does a certain function exist in the `FunctionsLib`? + pub fn has_function(&self, name: &str, params: usize) -> bool { + self.0.binary_search_by(|f| f.compare(name, params)).is_ok() + } + /// Add a function (or replace an existing one) in the `FunctionsLib`. + pub fn add_or_replace_function(&mut self, fn_def: Arc) { + match self + .0 + .binary_search_by(|f| f.compare(&fn_def.name, fn_def.params.len())) + { + Ok(n) => self.0[n] = fn_def, + Err(n) => self.0.insert(n, fn_def), + } + } + /// Get a function definition from the `FunctionsLib`. + pub fn get_function(&self, name: &str, params: usize) -> Option> { + if let Ok(n) = self.0.binary_search_by(|f| f.compare(name, params)) { + Some(self.0[n].clone()) + } else { + None + } + } +} + /// Rhai main scripting engine. /// /// ``` @@ -72,9 +131,9 @@ pub struct Engine<'e> { #[cfg(not(feature = "no_optimize"))] pub(crate) optimization_level: OptimizationLevel, /// A hashmap containing all compiled functions known to the engine - pub(crate) ext_functions: HashMap, Box>, + pub(crate) functions: HashMap, Box>, /// A hashmap containing all script-defined functions - pub(crate) script_functions: Vec>, + pub(crate) fn_lib: FunctionsLib, /// A hashmap containing all iterators known to the engine pub(crate) type_iterators: HashMap>, /// A hashmap mapping type names to pretty-print names @@ -101,8 +160,8 @@ impl Default for Engine<'_> { // Create the new scripting Engine let mut engine = Engine { - ext_functions: HashMap::new(), - script_functions: Vec::new(), + functions: HashMap::new(), + fn_lib: FunctionsLib::new(), type_iterators: HashMap::new(), type_names, on_print: Box::new(default_print), // default print/debug implementations @@ -152,7 +211,7 @@ impl Engine<'_> { }; // Search built-in's and external functions - if let Some(func) = self.ext_functions.get(&spec) { + if let Some(func) = self.functions.get(&spec) { // Run external function Ok(Some(func(args, pos)?)) } else { @@ -170,14 +229,9 @@ impl Engine<'_> { pos: Position, ) -> Result { // First search in script-defined functions (can override built-in) - if let Ok(n) = self - .script_functions - .binary_search_by(|f| f.compare(fn_name, args.len())) - { + if let Some(fn_def) = self.fn_lib.get_function(fn_name, args.len()) { let mut scope = Scope::new(); - let fn_def = self.script_functions[n].clone(); - scope.extend( // Put arguments into scope as variables fn_def @@ -191,12 +245,9 @@ impl Engine<'_> { // Convert return statement to return value return self .eval_stmt(&mut scope, &fn_def.body) - .or_else(|mut err| match err { + .or_else(|err| match err { EvalAltResult::Return(x, _) => Ok(x), - _ => { - err.set_position(pos); - Err(err) - } + err => Err(err.set_position(pos)), }); } @@ -213,7 +264,7 @@ impl Engine<'_> { } // Search built-in's and external functions - if let Some(func) = self.ext_functions.get(&spec) { + if let Some(func) = self.functions.get(&spec) { // Run external function let result = func(args, pos)?; @@ -985,11 +1036,7 @@ impl Engine<'_> { args: Some(vec![TypeId::of::()]), }; - engine.ext_functions.contains_key(&spec) - || engine - .script_functions - .binary_search_by(|f| f.compare(name, 1)) - .is_ok() + engine.functions.contains_key(&spec) || engine.fn_lib.has_function(name, 1) } match fn_name.as_str() { @@ -1005,13 +1052,12 @@ impl Engine<'_> { let result = args_expr_list .iter() .map(|expr| format!("{:#?}", expr)) - .collect::>() - .join("\n"); + .collect::>(); // Redirect call to `print` self.call_fn_raw( KEYWORD_PRINT, - &mut [result.into_dynamic().as_mut()], + &mut [result.join("\n").into_dynamic().as_mut()], None, pos, ) @@ -1061,10 +1107,7 @@ impl Engine<'_> { Ok(self .eval_ast_with_scope_raw(scope, true, &ast) - .map_err(|mut err| { - err.set_position(pos); - err - })?) + .map_err(|err| err.set_position(pos))?) } // Normal function call @@ -1297,7 +1340,7 @@ impl Engine<'_> { /// Clean up all script-defined functions within the `Engine`. pub fn clear_functions(&mut self) { - self.script_functions.clear(); + self.fn_lib.clear(); } } diff --git a/src/fn_register.rs b/src/fn_register.rs index acf0822d..df2d01d0 100644 --- a/src/fn_register.rs +++ b/src/fn_register.rs @@ -226,10 +226,8 @@ macro_rules! def_register { // Call the user-supplied function using ($clone) to // potentially clone the value, otherwise pass the reference. - f($(($clone)($par)),*).map(|r| Box::new(r) as Dynamic).map_err(|mut err| { - err.set_position(pos); - err - }) + f($(($clone)($par)),*).map(|r| Box::new(r) as Dynamic) + .map_err(|err| err.set_position(pos)) }; self.register_fn_raw(name, Some(vec![$(TypeId::of::<$par>()),*]), Box::new(fun)); } diff --git a/src/optimize.rs b/src/optimize.rs index 185fccc1..c8dc32f7 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -180,16 +180,15 @@ fn optimize_stmt<'a>(stmt: Stmt, state: &mut State<'a>, preserve_result: bool) - // Optimize each statement in the block let mut result: Vec<_> = block .into_iter() - .map(|stmt| { - if let Stmt::Const(name, value, pos) = stmt { - // Add constant into the state + .map(|stmt| match stmt { + // Add constant into the state + Stmt::Const(name, value, pos) => { state.push_constant(&name, *value); state.set_dirty(); Stmt::Noop(pos) // No need to keep constants - } else { - // Optimize the statement - optimize_stmt(stmt, state, preserve_result) } + // Optimize the statement + _ => optimize_stmt(stmt, state, preserve_result), }) .collect(); @@ -445,7 +444,7 @@ fn optimize_expr<'a>(expr: Expr, state: &mut State<'a>) -> Expr { && args.iter().all(|expr| expr.is_constant()) // all arguments are constants => { // First search in script-defined functions (can override built-in) - if state.engine.script_functions.binary_search_by(|f| f.compare(&id, args.len())).is_ok() { + if state.engine.fn_lib.has_function(&id, args.len()) { // A script-defined function overrides the built-in function - do not make the call return Expr::FunctionCall(id, args.into_iter().map(|a| optimize_expr(a, state)).collect(), def_value, pos); } diff --git a/src/parser.rs b/src/parser.rs index e14a1de1..92d0bb05 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -11,9 +11,7 @@ use crate::optimize::optimize_into_ast; use crate::stdlib::{ borrow::Cow, boxed::Box, - char, - cmp::Ordering, - fmt, format, + char, fmt, format, iter::Peekable, str::Chars, str::FromStr, @@ -175,18 +173,6 @@ pub struct FnDef { pub pos: Position, } -impl FnDef { - /// Function to order two FnDef records, for binary search. - pub fn compare(&self, name: &str, params_len: usize) -> Ordering { - // First order by name - match self.name.as_str().cmp(name) { - // Then by number of parameters - Ordering::Equal => self.params.len().cmp(¶ms_len), - order => order, - } - } -} - /// `return`/`throw` statement. #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)] pub enum ReturnType { diff --git a/src/result.rs b/src/result.rs index 06f5f075..793f99d5 100644 --- a/src/result.rs +++ b/src/result.rs @@ -236,7 +236,9 @@ impl EvalAltResult { } } - pub(crate) fn set_position(&mut self, new_position: Position) { + /// Consume the current `EvalAltResult` and return a new one + /// with the specified `Position`. + pub(crate) fn set_position(mut self, new_position: Position) -> Self { match self { #[cfg(not(feature = "no_std"))] Self::ErrorReadingScriptFile(_, _) => (), @@ -262,5 +264,7 @@ impl EvalAltResult { | Self::ErrorLoopBreak(ref mut pos) | Self::Return(_, ref mut pos) => *pos = new_position, } + + self } }