Skip to content

Commit

Permalink
Add 'memoized_size_{}' funcion to retrieve the size of the cache
Browse files Browse the repository at this point in the history
  • Loading branch information
gzsombor committed Dec 11, 2024
1 parent 84e37f7 commit ad406a0
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 0 deletions.
3 changes: 3 additions & 0 deletions examples/empty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,11 @@ fn hello() -> bool {
fn main() {
// `hello` is only called once here.
assert!(hello());
assert_eq!(memoized_size_hello(), 1);
assert!(hello());
assert_eq!(memoized_size_hello(), 1);
memoized_flush_hello();
assert_eq!(memoized_size_hello(), 0);
// and again here.
assert!(hello());
}
2 changes: 2 additions & 0 deletions examples/shared.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,5 +12,7 @@ fn main() {
assert!(!hello("World".to_string(), 0));
// Sometimes one might need the original function.
assert!(!memoized_original_hello("World".to_string(), 0));
assert_eq!(memoized_size_hello(), 1);
memoized_flush_hello();
assert_eq!(memoized_size_hello(), 0);
}
16 changes: 16 additions & 0 deletions inner/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,7 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream {
let fn_name = &sig.ident.to_string();
let renamed_name = format!("memoized_original_{}", fn_name);
let flush_name = syn::Ident::new(format!("memoized_flush_{}", fn_name).as_str(), sig.span());
let size_name = syn::Ident::new(format!("memoized_size_{}", fn_name).as_str(), sig.span());
let map_name = format!("memoized_mapping_{}", fn_name);

if let Some(syn::FnArg::Receiver(_)) = sig.inputs.first() {
Expand Down Expand Up @@ -430,9 +431,24 @@ pub fn memoize(attr: TokenStream, item: TokenStream) -> TokenStream {
}
};

let size_func = if options.shared_cache {
quote::quote! {
#vis fn #size_name() -> usize {
#store_ident.lock().unwrap().len()
}
}
} else {
quote::quote! {
#vis fn #size_name() -> usize {
#store_ident.with(|ATTR_MEMOIZE_HM__| ATTR_MEMOIZE_HM__.borrow().len())
}
}
};

quote::quote! {
#renamed_fn
#flusher
#size_func
#store

#[allow(unused_variables, unused_mut)]
Expand Down

0 comments on commit ad406a0

Please sign in to comment.