Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ability to Ignore arguments #29

Merged
merged 5 commits into from
Nov 26, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
41 changes: 41 additions & 0 deletions examples/ignore.rs
Original file line number Diff line number Diff line change
@@ -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);
}
136 changes: 100 additions & 36 deletions inner/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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, :);
}

Expand All @@ -21,6 +22,7 @@ struct CacheOptions {
shared_cache: bool,
custom_hasher: Option<Path>,
custom_hasher_initializer: Option<ExprCall>,
ignore: Vec<syn::Ident>,
}

#[derive(Clone)]
Expand All @@ -30,6 +32,7 @@ enum CacheOption {
SharedCache,
CustomHasher(Path),
HasherInit(ExprCall),
Ignore(syn::Ident),
}

// To extend option parsing, add functionality here.
Expand Down Expand Up @@ -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::<kw::Ignore>().unwrap();
input.parse::<kw::Colon>().unwrap();
let ignore_ident = input.parse::<syn::Ident>().unwrap();
return Ok(CacheOption::Ignore(ignore_ident));
}
Err(la.error())
}
}
Expand All @@ -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)
Expand Down Expand Up @@ -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.
Expand All @@ -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.*
Expand Down Expand Up @@ -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<Box<syn::Type>>;
let input_names: Vec<syn::Ident>;
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<Box<syn::Type>> = input_params
.iter()
.filter_map(|p| {
if p.is_memoized {
Some(p.arg_type.clone())
} else {
None
}
})
.collect();
let memoized_input_names: Vec<syn::Ident> = 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) =
Expand All @@ -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 => (
Expand All @@ -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
Expand All @@ -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();
Expand Down Expand Up @@ -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<syn::Type>,

/// 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<Box<syn::Type>>, Vec<syn::Ident>), syn::Error> {
options: &CacheOptions,
) -> Result<Vec<FnArgument>, 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(),
Expand All @@ -421,7 +485,7 @@ fn check_signature(
}
}
}
Ok((types, names))
Ok(params)
}

#[cfg(test)]
Expand Down
Loading