From f8b747a40e82ecf76d9017b0636f74148a4afccc Mon Sep 17 00:00:00 2001 From: Mason Smith Date: Wed, 8 Nov 2023 12:25:01 -0800 Subject: [PATCH 1/5] implementation of Ignore parameter --- examples/ignore.rs | 42 ++++++++++++++++ inner/src/lib.rs | 122 +++++++++++++++++++++++++++++++++------------ 2 files changed, 133 insertions(+), 31 deletions(-) create mode 100644 examples/ignore.rs diff --git a/examples/ignore.rs b/examples/ignore.rs new file mode 100644 index 0000000..a2f8e94 --- /dev/null +++ b/examples/ignore.rs @@ -0,0 +1,42 @@ +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 +} + +#[cfg(feature = "full")] +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..4160b99 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) @@ -264,28 +274,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 syn::FnArg::Receiver(_) = sig.inputs[0] { + 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 +355,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 +382,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 +399,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 +439,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 +481,7 @@ fn check_signature( } } } - Ok((types, names)) + Ok(params) } #[cfg(test)] From 7e256f58e1004abf5cee95f2d39cb09fe181dc17 Mon Sep 17 00:00:00 2001 From: Mason Smith Date: Wed, 8 Nov 2023 13:58:46 -0800 Subject: [PATCH 2/5] update docs to note Ignore arguments --- inner/src/lib.rs | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/inner/src/lib.rs b/inner/src/lib.rs index 4160b99..eb48014 100644 --- a/inner/src/lib.rs +++ b/inner/src/lib.rs @@ -218,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. @@ -245,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.* From a420c703b0ef785cc805514d9cc517f0d1b81b7a Mon Sep 17 00:00:00 2001 From: Mason Smith Date: Thu, 9 Nov 2023 09:32:18 -0800 Subject: [PATCH 3/5] review feedback: remove 'full' feature gate from ignore example --- examples/ignore.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/ignore.rs b/examples/ignore.rs index a2f8e94..e9bfa21 100644 --- a/examples/ignore.rs +++ b/examples/ignore.rs @@ -19,7 +19,6 @@ fn add2(a: u32, b: u32, call_count: &mut u32) -> u32 { a + b } -#[cfg(feature = "full")] 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); From 9157bf6f99f88e9ebdc2c563b3cdc080d228df0b Mon Sep 17 00:00:00 2001 From: Mason Smith Date: Thu, 9 Nov 2023 09:43:33 -0800 Subject: [PATCH 4/5] update README.md to include 'Ignore' information --- README.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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 From d6cdf4e6c9e273ec75529d3852849b89026180c4 Mon Sep 17 00:00:00 2001 From: Mason Smith Date: Mon, 13 Nov 2023 13:15:03 -0800 Subject: [PATCH 5/5] Fix for functions that take no arguments --- inner/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inner/src/lib.rs b/inner/src/lib.rs index eb48014..c7d708b 100644 --- a/inner/src/lib.rs +++ b/inner/src/lib.rs @@ -278,7 +278,7 @@ 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); - if let syn::FnArg::Receiver(_) = sig.inputs[0] { + if let Some(syn::FnArg::Receiver(_)) = sig.inputs.first() { return quote::quote! { compile_error!("Cannot memoize methods!"); }.into(); }