diff --git a/README.md b/README.md index bea8dd0..aa08845 100644 --- a/README.md +++ b/README.md @@ -130,6 +130,23 @@ As some hashers initializing functions other than `new()`, you can specifiy a `H #[memoize(CustomHasher: FxHashMap, HasherInit: FxHashMap::default())] ``` +Sometimes, you can't or don't want to store data as part of the cache. In those cases, you can use +the `Ignore` parameter in the `#[memoize]` macro to ignore an argument. Any `Ignore`d arguments no +longer need to be `Clone`-able, since they are not stored as part of the argument set, and changing +an `Ignore`d argument will not trigger calling the function again. You can `Ignore` multiple +arugments by specifying the `Ignore` parameter multiple times. + +```rust +// `Ignore: count_calls` lets our function take a `&mut u32` argument, which is normally not +// possible because it is not `Clone`-able. +#[memoize(Ignore: count_calls)] +fn add(a: u32, b: u32, count_calls: &mut u32) -> u32 { + // Keep track of the number of times the underlying function is called. + *count_calls += 1; + a + b +} +``` + ### Flushing If you memoize a function `f`, there will be a function called diff --git a/examples/ignore.rs b/examples/ignore.rs new file mode 100644 index 0000000..e9bfa21 --- /dev/null +++ b/examples/ignore.rs @@ -0,0 +1,41 @@ +use memoize::memoize; + +/// Wrapper struct for a [`u32`]. +/// +/// Note that A deliberately does not implement [`Clone`] or [`Hash`], to demonstrate that it can be +/// passed through. +struct C { + c: u32 +} + +#[memoize(Ignore: a, Ignore: c)] +fn add(a: u32, b: u32, c: C, d: u32) -> u32 { + a + b + c.c + d +} + +#[memoize(Ignore: call_count, SharedCache)] +fn add2(a: u32, b: u32, call_count: &mut u32) -> u32 { + *call_count += 1; + a + b +} + +fn main() { + // Note that the third argument is not `Clone` but can still be passed through. + assert_eq!(add(1, 2, C {c: 3}, 4), 10); + + assert_eq!(add(3, 2, C {c: 4}, 4), 10); + memoized_flush_add(); + + // Once cleared, all arguments is again used. + assert_eq!(add(3, 2, C {c: 4}, 4), 13); + + let mut count_unique_calls = 0; + assert_eq!(add2(1, 2, &mut count_unique_calls), 3); + assert_eq!(count_unique_calls, 1); + + // Calling `add2` again won't increment `count_unique_calls` + // because it's ignored as a parameter, and the other arguments + // are the same. + add2(1, 2, &mut count_unique_calls); + assert_eq!(count_unique_calls, 1); +} diff --git a/inner/src/lib.rs b/inner/src/lib.rs index 5820acb..c7d708b 100644 --- a/inner/src/lib.rs +++ b/inner/src/lib.rs @@ -11,6 +11,7 @@ mod kw { syn::custom_keyword!(SharedCache); syn::custom_keyword!(CustomHasher); syn::custom_keyword!(HasherInit); + syn::custom_keyword!(Ignore); syn::custom_punctuation!(Colon, :); } @@ -21,6 +22,7 @@ struct CacheOptions { shared_cache: bool, custom_hasher: Option, custom_hasher_initializer: Option, + ignore: Vec, } #[derive(Clone)] @@ -30,6 +32,7 @@ enum CacheOption { SharedCache, CustomHasher(Path), HasherInit(ExprCall), + Ignore(syn::Ident), } // To extend option parsing, add functionality here. @@ -77,6 +80,12 @@ impl parse::Parse for CacheOption { let cap: syn::ExprCall = input.parse().unwrap(); return Ok(CacheOption::HasherInit(cap)); } + if la.peek(kw::Ignore) { + input.parse::().unwrap(); + input.parse::().unwrap(); + let ignore_ident = input.parse::().unwrap(); + return Ok(CacheOption::Ignore(ignore_ident)); + } Err(la.error()) } } @@ -94,6 +103,7 @@ impl parse::Parse for CacheOptions { CacheOption::CustomHasher(hasher) => opts.custom_hasher = Some(hasher), CacheOption::HasherInit(init) => opts.custom_hasher_initializer = Some(init), CacheOption::SharedCache => opts.shared_cache = true, + CacheOption::Ignore(ident) => opts.ignore.push(ident), } } Ok(opts) @@ -208,11 +218,11 @@ mod store { /** * memoize is an attribute to create a memoized version of a (simple enough) function. * - * So far, it works on functions with one or more arguments which are `Clone`- and `Hash`-able, - * returning a `Clone`-able value. Several clones happen within the storage and recall layer, with - * the assumption being that `memoize` is used to cache such expensive functions that very few - * `clone()`s do not matter. `memoize` doesn't work on methods (functions with `[&/&mut/]self` - * receiver). + * So far, it works on non-method functions with one or more arguments returning a [`Clone`]-able + * value. Arguments that are cached must be [`Clone`]-able and [`Hash`]-able as well. Several clones + * happen within the storage and recall layer, with the assumption being that `memoize` is used to + * cache such expensive functions that very few `clone()`s do not matter. `memoize` doesn't work on + * methods (functions with `[&/&mut/]self` receiver). * * Calls are memoized for the lifetime of a program, using a statically allocated, Mutex-protected * HashMap. @@ -235,6 +245,10 @@ mod store { * If you need to use the un-memoized function, it is always available as `memoized_original_{fn}`, * in this case: `memoized_original_hello()`. * + * Parameters can be ignored by the cache using the `Ignore` parameter. `Ignore` can be specified + * multiple times, once per each parameter. `Ignore`d parameters do not need to implement [`Clone`] + * or [`Hash`]. + * * See the `examples` for concrete applications. * * *The following descriptions need the `full` feature enabled.* @@ -264,28 +278,61 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream { let flush_name = syn::Ident::new(format!("memoized_flush_{}", fn_name).as_str(), sig.span()); let map_name = format!("memoized_mapping_{}", fn_name); - // Extracted from the function signature. - let input_types: Vec>; - let input_names: Vec; - let return_type; - - match check_signature(sig) { - Ok((t, n)) => { - input_types = t; - input_names = n; - } - Err(e) => return e.to_compile_error().into(), - } - - let input_tuple_type = quote::quote! { (#(#input_types),*) }; - match &sig.output { - syn::ReturnType::Default => return_type = quote::quote! { () }, - syn::ReturnType::Type(_, ty) => return_type = ty.to_token_stream(), + if let Some(syn::FnArg::Receiver(_)) = sig.inputs.first() { + return quote::quote! { compile_error!("Cannot memoize methods!"); }.into(); } // Parse options from macro attributes let options: CacheOptions = syn::parse(attr.clone()).unwrap(); + // Extracted from the function signature. + let input_params = match check_signature(sig, &options) { + Ok(p) => p, + Err(e) => return e.to_compile_error().into(), + }; + + // Input types and names that are actually stored in the cache. + let memoized_input_types: Vec> = input_params + .iter() + .filter_map(|p| { + if p.is_memoized { + Some(p.arg_type.clone()) + } else { + None + } + }) + .collect(); + let memoized_input_names: Vec = input_params + .iter() + .filter_map(|p| { + if p.is_memoized { + Some(p.arg_name.clone()) + } else { + None + } + }) + .collect(); + + // For each input, expression to be passe through to the original function. + // Cached arguments are cloned, original arguments are forwarded as-is + let fn_forwarded_exprs: Vec<_> = input_params + .iter() + .map(|p| { + let ident = p.arg_name.clone(); + if p.is_memoized { + quote::quote! { #ident.clone() } + } else { + quote::quote! { #ident } + } + }) + .collect(); + + let input_tuple_type = quote::quote! { (#(#memoized_input_types),*) }; + let return_type = match &sig.output { + syn::ReturnType::Default => quote::quote! { () }, + syn::ReturnType::Type(_, ty) => ty.to_token_stream(), + }; + // Construct storage for the memoized keys and return values. let store_ident = syn::Ident::new(&map_name.to_uppercase(), sig.span()); let (cache_type, cache_init) = @@ -312,8 +359,9 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream { let memoized_id = &renamed_fn.sig.ident; // Construct memoizer function, which calls the original function. - let syntax_names_tuple = quote::quote! { (#(#input_names),*) }; - let syntax_names_tuple_cloned = quote::quote! { (#(#input_names.clone()),*) }; + let syntax_names_tuple = quote::quote! { (#(#memoized_input_names),*) }; + let syntax_names_tuple_cloned = quote::quote! { (#(#memoized_input_names.clone()),*) }; + let forwarding_tuple = quote::quote! { (#(#fn_forwarded_exprs),*) }; let (insert_fn, get_fn) = store::cache_access_methods(&options); let (read_memo, memoize) = match options.time_to_live { None => ( @@ -338,7 +386,7 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream { return ATTR_MEMOIZE_RETURN__ } } - let ATTR_MEMOIZE_RETURN__ = #memoized_id(#(#input_names.clone()),*); + let ATTR_MEMOIZE_RETURN__ = #memoized_id #forwarding_tuple; let mut ATTR_MEMOIZE_HM__ = #store_ident.lock().unwrap(); #memoize @@ -355,7 +403,7 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream { return ATTR_MEMOIZE_RETURN__; } - let ATTR_MEMOIZE_RETURN__ = #memoized_id(#(#input_names.clone()),*); + let ATTR_MEMOIZE_RETURN__ = #memoized_id #forwarding_tuple; #store_ident.with(|ATTR_MEMOIZE_HM__| { let mut ATTR_MEMOIZE_HM__ = ATTR_MEMOIZE_HM__.borrow_mut(); @@ -395,24 +443,40 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream { .into() } +/// An argument of the memoized function. +struct FnArgument { + /// Type of the argument. + arg_type: Box, + + /// Identifier (name) of the argument. + arg_name: syn::Ident, + + /// Whether or not this specific argument is included in the memoization. + is_memoized: bool, +} + fn check_signature( sig: &syn::Signature, -) -> Result<(Vec>, Vec), syn::Error> { + options: &CacheOptions, +) -> Result, syn::Error> { if sig.inputs.is_empty() { - return Ok((vec![], vec![])); - } - if let syn::FnArg::Receiver(_) = sig.inputs[0] { - return Err(syn::Error::new(sig.span(), "Cannot memoize methods!")); + return Ok(vec![]); } - let mut types = vec![]; - let mut names = vec![]; + let mut params = vec![]; + for a in &sig.inputs { if let syn::FnArg::Typed(ref arg) = a { - types.push(arg.ty.clone()); + let arg_type = arg.ty.clone(); if let syn::Pat::Ident(patident) = &*arg.pat { - names.push(patident.ident.clone()); + let arg_name = patident.ident.clone(); + let is_memoized = !options.ignore.contains(&arg_name); + params.push(FnArgument { + arg_type, + arg_name, + is_memoized, + }); } else { return Err(syn::Error::new( sig.span(), @@ -421,7 +485,7 @@ fn check_signature( } } } - Ok((types, names)) + Ok(params) } #[cfg(test)]