Enable &str and String parameters in module functions.

This commit is contained in:
Stephen Chung 2020-09-19 18:12:23 +08:00
parent 2ff3a1fde5
commit 39546b7053
2 changed files with 76 additions and 26 deletions

View File

@ -4,10 +4,11 @@ use crate::any::{Dynamic, Variant};
use crate::calc_fn_hash; use crate::calc_fn_hash;
use crate::engine::Engine; use crate::engine::Engine;
use crate::fn_native::{CallableFunction as Func, FnCallArgs, IteratorFn, SendSync}; use crate::fn_native::{CallableFunction as Func, FnCallArgs, IteratorFn, SendSync};
use crate::fn_register::by_value as cast_arg;
use crate::parser::{FnAccess, FnAccess::Public, ScriptFnDef}; use crate::parser::{FnAccess, FnAccess::Public, ScriptFnDef};
use crate::result::EvalAltResult; use crate::result::EvalAltResult;
use crate::token::{Position, Token}; use crate::token::{Position, Token};
use crate::utils::{StaticVec, StraightHasherBuilder}; use crate::utils::{ImmutableString, StaticVec, StraightHasherBuilder};
#[cfg(not(feature = "no_function"))] #[cfg(not(feature = "no_function"))]
use crate::fn_native::Shared; use crate::fn_native::Shared;
@ -32,7 +33,6 @@ use crate::stdlib::{
collections::HashMap, collections::HashMap,
fmt, format, fmt, format,
iter::empty, iter::empty,
mem,
num::NonZeroUsize, num::NonZeroUsize,
ops::{Deref, DerefMut}, ops::{Deref, DerefMut},
string::{String, ToString}, string::{String, ToString},
@ -396,9 +396,21 @@ impl Module {
arg_types.len() arg_types.len()
}; };
let hash_fn = calc_fn_hash(empty(), &name, args_len, arg_types.iter().cloned()); let params = arg_types
.into_iter()
.cloned()
.map(|id| {
if id == TypeId::of::<&str>() {
TypeId::of::<ImmutableString>()
} else if id == TypeId::of::<String>() {
TypeId::of::<ImmutableString>()
} else {
id
}
})
.collect();
let params = arg_types.into_iter().cloned().collect(); let hash_fn = calc_fn_hash(empty(), &name, args_len, arg_types.iter().cloned());
self.functions self.functions
.insert(hash_fn, (name, access, params, func.into())); .insert(hash_fn, (name, access, params, func.into()));
@ -518,7 +530,7 @@ impl Module {
func: impl Fn(A) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(A) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
func(mem::take(args[0]).cast::<A>()).map(Dynamic::from) func(cast_arg::<A>(&mut args[0])).map(Dynamic::from)
}; };
let arg_types = [TypeId::of::<A>()]; let arg_types = [TypeId::of::<A>()];
self.set_fn(name, Public, &arg_types, Func::from_pure(Box::new(f))) self.set_fn(name, Public, &arg_types, Func::from_pure(Box::new(f)))
@ -592,8 +604,8 @@ impl Module {
func: impl Fn(A, B) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(A, B) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let a = mem::take(args[0]).cast::<A>(); let a = cast_arg::<A>(&mut args[0]);
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[1]);
func(a, b).map(Dynamic::from) func(a, b).map(Dynamic::from)
}; };
@ -623,7 +635,7 @@ impl Module {
func: impl Fn(&mut A, B) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(&mut A, B) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[1]);
let a = &mut args[0].write_lock::<A>().unwrap(); let a = &mut args[0].write_lock::<A>().unwrap();
func(a, b).map(Dynamic::from) func(a, b).map(Dynamic::from)
@ -709,9 +721,9 @@ impl Module {
func: impl Fn(A, B, C) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(A, B, C) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let a = mem::take(args[0]).cast::<A>(); let a = cast_arg::<A>(&mut args[0]);
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[1]);
let c = mem::take(args[2]).cast::<C>(); let c = cast_arg::<C>(&mut args[2]);
func(a, b, c).map(Dynamic::from) func(a, b, c).map(Dynamic::from)
}; };
@ -746,8 +758,8 @@ impl Module {
func: impl Fn(&mut A, B, C) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(&mut A, B, C) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[2]);
let c = mem::take(args[2]).cast::<C>(); let c = cast_arg::<C>(&mut args[3]);
let a = &mut args[0].write_lock::<A>().unwrap(); let a = &mut args[0].write_lock::<A>().unwrap();
func(a, b, c).map(Dynamic::from) func(a, b, c).map(Dynamic::from)
@ -780,8 +792,8 @@ impl Module {
func: impl Fn(&mut A, B, C) -> FuncReturn<()> + SendSync + 'static, func: impl Fn(&mut A, B, C) -> FuncReturn<()> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[1]);
let c = mem::take(args[2]).cast::<C>(); let c = cast_arg::<C>(&mut args[2]);
let a = &mut args[0].write_lock::<A>().unwrap(); let a = &mut args[0].write_lock::<A>().unwrap();
func(a, b, c).map(Dynamic::from) func(a, b, c).map(Dynamic::from)
@ -858,10 +870,10 @@ impl Module {
func: impl Fn(A, B, C, D) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(A, B, C, D) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let a = mem::take(args[0]).cast::<A>(); let a = cast_arg::<A>(&mut args[0]);
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[1]);
let c = mem::take(args[2]).cast::<C>(); let c = cast_arg::<C>(&mut args[2]);
let d = mem::take(args[3]).cast::<D>(); let d = cast_arg::<D>(&mut args[3]);
func(a, b, c, d).map(Dynamic::from) func(a, b, c, d).map(Dynamic::from)
}; };
@ -902,9 +914,9 @@ impl Module {
func: impl Fn(&mut A, B, C, D) -> FuncReturn<T> + SendSync + 'static, func: impl Fn(&mut A, B, C, D) -> FuncReturn<T> + SendSync + 'static,
) -> u64 { ) -> u64 {
let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| { let f = move |_: &Engine, _: &Module, args: &mut FnCallArgs| {
let b = mem::take(args[1]).cast::<B>(); let b = cast_arg::<B>(&mut args[1]);
let c = mem::take(args[2]).cast::<C>(); let c = cast_arg::<C>(&mut args[2]);
let d = mem::take(args[3]).cast::<D>(); let d = cast_arg::<D>(&mut args[3]);
let a = &mut args[0].write_lock::<A>().unwrap(); let a = &mut args[0].write_lock::<A>().unwrap();
func(a, b, c, d).map(Dynamic::from) func(a, b, c, d).map(Dynamic::from)

View File

@ -1,7 +1,7 @@
#![cfg(not(feature = "no_module"))] #![cfg(not(feature = "no_module"))]
use rhai::{ use rhai::{
module_resolvers::StaticModuleResolver, Dynamic, Engine, EvalAltResult, Module, ParseError, module_resolvers::StaticModuleResolver, Dynamic, Engine, EvalAltResult, ImmutableString,
ParseErrorType, Scope, INT, Module, ParseError, ParseErrorType, Scope, INT,
}; };
#[test] #[test]
@ -79,12 +79,12 @@ fn test_module_resolver() -> Result<(), Box<EvalAltResult>> {
}); });
#[cfg(not(feature = "no_float"))] #[cfg(not(feature = "no_float"))]
module.set_fn_4_mut( module.set_fn_4_mut(
"sum_of_three_args".to_string(), "sum_of_three_args".to_string(),
|target: &mut INT, a: INT, b: INT, c: f64| { |target: &mut INT, a: INT, b: INT, c: f64| {
*target = a + b + c as INT; *target = a + b + c as INT;
Ok(()) Ok(())
} },
); );
resolver.insert("hello", module); resolver.insert("hello", module);
@ -316,3 +316,41 @@ fn test_module_export() -> Result<(), Box<EvalAltResult>> {
Ok(()) Ok(())
} }
#[test]
fn test_module_str() -> Result<(), Box<EvalAltResult>> {
fn test_fn(_input: ImmutableString) -> Result<INT, Box<EvalAltResult>> {
Ok(42)
}
fn test_fn2(_input: &str) -> Result<INT, Box<EvalAltResult>> {
Ok(42)
}
fn test_fn3(_input: String) -> Result<INT, Box<EvalAltResult>> {
Ok(42)
}
let mut engine = rhai::Engine::new();
let mut module = Module::new();
module.set_fn_1("test", test_fn);
module.set_fn_1("test2", test_fn2);
module.set_fn_1("test3", test_fn3);
let mut static_modules = rhai::module_resolvers::StaticModuleResolver::new();
static_modules.insert("test", module);
engine.set_module_resolver(Some(static_modules));
assert_eq!(
engine.eval::<INT>(r#"import "test" as test; test::test("test");"#)?,
42
);
assert_eq!(
engine.eval::<INT>(r#"import "test" as test; test::test2("test");"#)?,
42
);
assert_eq!(
engine.eval::<INT>(r#"import "test" as test; test::test3("test");"#)?,
42
);
Ok(())
}