Tracked methods with arguments

Author: laurmaedje

Cargo.toml

@@ -6,3 +6,6 @@ edition = "2021"
 [dependencies]
 comemo-macros = { path = "./macros" }
 siphasher = "0.3"
+
+[dev-dependencies]
+unscanny = "0.1"

macros/src/lib.rs

@@ -13,8 +13,7 @@ mod memoize;
 mod track;
 
 use proc_macro::TokenStream;
-use quote::{quote, quote_spanned};
-use syn::spanned::Spanned;
+use quote::quote;
 use syn::{parse_quote, Error, Result};
 
 /// Memoize a pure function.

macros/src/memoize.rs

@@ -12,28 +12,35 @@ pub fn expand(mut func: syn::ItemFn) -> Result<proc_macro2::TokenStream> {
             }
         };
 
-        let name = match &*typed.pat {
-            syn::Pat::Ident(ident) => ident,
-            _ => bail!(typed.pat, "only simple identifiers are supported"),
+        let name = match typed.pat.as_ref() {
+            syn::Pat::Ident(syn::PatIdent {
+                by_ref: None,
+                mutability: None,
+                ident,
+                subpat: None,
+                ..
+            }) => ident,
+            pat => bail!(pat, "only simple identifiers are supported"),
         };
 
+        let ty = typed.ty.as_ref();
         args.push(name);
-        types.push(typed.ty.as_ref());
+        types.push(ty);
     }
 
     // Construct a tuple from all arguments.
     let arg_tuple = quote! { (#(#args,)*) };
 
     // Construct assertions that the arguments fulfill the necessary bounds.
-    let bounds = args.iter().zip(&types).map(|(arg, ty)| {
-        quote_spanned! {
-            arg.span() => ::comemo::internal::assert_hashable_or_trackable::<#ty>();
+    let bounds = types.iter().map(|ty| {
+        quote! {
+            ::comemo::internal::assert_hashable_or_trackable::<#ty>();
         }
     });
 
     // Construct the inner closure.
     let body = &func.block;
-    let inner = quote! { |#arg_tuple| #body };
+    let closure = quote! { |#arg_tuple| #body };
 
     // Adjust the function's body.
     let name = func.sig.ident.to_string();
@@ -42,7 +49,7 @@ pub fn expand(mut func: syn::ItemFn) -> Result<proc_macro2::TokenStream> {
         ::comemo::internal::cached(
             #name,
             ::comemo::internal::Args(#arg_tuple),
-            #inner,
+            #closure,
         )
     } };
 

macros/src/track.rs

@@ -10,39 +10,68 @@ pub fn expand(block: syn::ItemImpl) -> Result<proc_macro2::TokenStream> {
         methods.push(method(&item)?);
     }
 
-    let tracked_fields = methods.iter().map(|method| {
-        let name = &method.sig.ident;
-        let ty = match &method.sig.output {
-            syn::ReturnType::Default => unreachable!(),
-            syn::ReturnType::Type(_, ty) => ty.as_ref(),
-        };
-        quote! { #name: ::comemo::internal::HashConstraint<#ty>, }
+    let tracked_valids = methods.iter().map(|method| {
+        let name = &method.name;
+        let args = &method.args;
+        if args.is_empty() {
+            quote! { constraint.#name.valid(&self.#name()) }
+        } else {
+            quote! {
+                constraint.#name
+                    .valid(|(#(#args,)*)| self.#name(#(#args.clone(),)*))
+            }
+        }
     });
 
     let tracked_methods = methods.iter().map(|method| {
-        let name = &method.sig.ident;
-        let mut method = (*method).clone();
-        if matches!(method.vis, syn::Visibility::Inherited) {
-            method.vis = parse_quote! { pub(super) };
+        let mut wrapper = method.item.clone();
+        if matches!(wrapper.vis, syn::Visibility::Inherited) {
+            wrapper.vis = parse_quote! { pub(super) };
         }
-        method.block = parse_quote! { {
-            let (inner, constraint) = ::comemo::internal::to_parts(self.0);
-            let output = inner.#name();
+
+        let name = &method.name;
+        let args = &method.args;
+        let set = if args.is_empty() {
+            quote! { constraint.#name.set(&output) }
+        } else {
+            quote! { constraint.#name.set((#(#args,)*), &output) }
+        };
+
+        // Construct assertions that the arguments fulfill the necessary bounds.
+        let bounds = method.types.iter().map(|ty| {
+            quote! {
+                ::comemo::internal::assert_clone_and_partial_eq::<#ty>();
+            }
+        });
+
+        wrapper.block = parse_quote! { {
+            #(#bounds;)*
+            let (value, constraint) = ::comemo::internal::to_parts(self.0);
+            let output = value.#name(#(#args.clone(),)*);
             if let Some(constraint) = &constraint {
-                constraint.#name.set(&output);
+                #set;
             }
             output
         } };
-        method
+
+        wrapper
     });
 
-    let tracked_valids = methods.iter().map(|method| {
-        let name = &method.sig.ident;
-        quote! {
-            constraint.#name.valid(&self.#name())
+    let tracked_fields = methods.iter().map(|method| {
+        let name = &method.name;
+        let types = &method.types;
+        if types.is_empty() {
+            quote! { #name: ::comemo::internal::HashConstraint, }
+        } else {
+            quote! { #name: ::comemo::internal::FuncConstraint<(#(#types,)*)>, }
         }
     });
 
+    let join_calls = methods.iter().map(|method| {
+        let name = &method.name;
+        quote! { self.#name.join(&inner.#name); }
+    });
+
     let track_impl = quote! {
         use super::*;
 
@@ -77,6 +106,12 @@ pub fn expand(block: syn::ItemImpl) -> Result<proc_macro2::TokenStream> {
         pub struct Constraint {
             #(#tracked_fields)*
         }
+
+        impl ::comemo::internal::Join for Constraint {
+            fn join(&self, inner: &Self) {
+                #(#join_calls)*
+            }
+        }
     };
 
     Ok(quote! {
@@ -85,8 +120,15 @@ pub fn expand(block: syn::ItemImpl) -> Result<proc_macro2::TokenStream> {
     })
 }
 
+struct Method {
+    item: syn::ImplItemMethod,
+    name: syn::Ident,
+    args: Vec<syn::Ident>,
+    types: Vec<syn::Type>,
+}
+
 /// Extract and validate a method.
-fn method(item: &syn::ImplItem) -> Result<&syn::ImplItemMethod> {
+fn method(item: &syn::ImplItem) -> Result<Method> {
     let method = match item {
         syn::ImplItem::Method(method) => method,
         _ => bail!(item, "only methods are supported"),
@@ -98,6 +140,27 @@ fn method(item: &syn::ImplItem) -> Result<&syn::ImplItemMethod> {
         _ => bail!(method.vis, "only private and public methods are supported"),
     }
 
+    if let Some(unsafety) = method.sig.unsafety {
+        bail!(unsafety, "unsafe methods are not supported");
+    }
+
+    if let Some(asyncness) = method.sig.asyncness {
+        bail!(asyncness, "async methods are not supported");
+    }
+
+    if let Some(constness) = method.sig.constness {
+        bail!(constness, "const methods are not supported");
+    }
+
+    for param in method.sig.generics.params.iter() {
+        match param {
+            syn::GenericParam::Const(_) | syn::GenericParam::Type(_) => {
+                bail!(param, "method must not be generic")
+            }
+            syn::GenericParam::Lifetime(_) => {}
+        }
+    }
+
     let mut inputs = method.sig.inputs.iter();
     let receiver = match inputs.next() {
         Some(syn::FnArg::Receiver(recv)) => recv,
@@ -108,20 +171,46 @@ fn method(item: &syn::ImplItem) -> Result<&syn::ImplItemMethod> {
         bail!(receiver, "must take self by shared reference");
     }
 
-    if inputs.next().is_some() {
-        bail!(
-            method.sig,
-            "currently, only methods without extra arguments are supported"
-        );
+    let mut args = vec![];
+    let mut types = vec![];
+    for input in inputs {
+        let typed = match input {
+            syn::FnArg::Typed(typed) => typed,
+            syn::FnArg::Receiver(_) => continue,
+        };
+
+        let name = match typed.pat.as_ref() {
+            syn::Pat::Ident(syn::PatIdent {
+                by_ref: None,
+                mutability: None,
+                ident,
+                subpat: None,
+                ..
+            }) => ident.clone(),
+            pat => bail!(pat, "only simple identifiers are supported"),
+        };
+
+        let ty = (*typed.ty).clone();
+        match ty {
+            syn::Type::ImplTrait(_) => bail!(ty, "method must not be generic"),
+            _ => {}
+        }
+
+        args.push(name);
+        types.push(ty);
     }
 
-    let output = &method.sig.output;
-    match output {
+    match method.sig.output {
         syn::ReturnType::Default => {
             bail!(method.sig, "method must have a return type")
         }
         syn::ReturnType::Type(..) => {}
     }
 
-    Ok(method)
+    Ok(Method {
+        item: method.clone(),
+        name: method.sig.ident.clone(),
+        args,
+        types,
+    })
 }

src/cache.rs

@@ -1,10 +1,11 @@
 use std::any::{Any, TypeId};
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::fmt::Debug;
 use std::hash::Hash;
 
 use siphasher::sip128::{Hasher128, SipHasher};
 
+use crate::constraint::Join;
 use crate::input::Input;
 use crate::internal::Family;
 
@@ -26,41 +27,51 @@ where
     let mut hit = true;
     let output = CACHE.with(|cache| {
         cache.lookup::<In, Out>(hash, &input).unwrap_or_else(|| {
+            DEPTH.with(|v| v.set(v.get() + 1));
             let constraint = In::Constraint::default();
-            let value = func(input.track(&constraint));
-            let constrained = Constrained { value: value.clone(), constraint };
-            cache.insert::<In, Out>(hash, constrained);
+            let (tracked, outer) = input.retrack(&constraint);
+            let output = func(tracked);
+            outer.join(&constraint);
+            cache.insert::<In, Out>(hash, Constrained {
+                output: output.clone(),
+                constraint,
+            });
             hit = false;
-            value
+            DEPTH.with(|v| v.set(v.get() - 1));
+            output
         })
     });
 
+    let depth = DEPTH.with(|v| v.get());
     let label = if hit { "[hit]" } else { "[miss]" };
-    eprintln!("{name:<9} {label:<7} {output:?}");
+    eprintln!("{depth} {name:<9} {label:<7} {output:?}");
 
     output
 }
 
 thread_local! {
     /// The global, dynamic cache shared by all memoized functions.
-    pub static CACHE: Cache = Cache::default();
+    static CACHE: Cache = Cache::default();
+
+    /// The current depth of the memoized call stack.
+    static DEPTH: Cell<usize> = Cell::new(0);
 }
 
 /// An untyped cache.
 #[derive(Default)]
-pub struct Cache {
+struct Cache {
     map: RefCell<Vec<Entry>>,
 }
 
 /// An entry in the cache.
 struct Entry {
     hash: u128,
-    output: Box<dyn Any>,
+    constrained: Box<dyn Any>,
 }
 
 /// A value with a constraint.
 struct Constrained<T, C> {
-    value: T,
+    output: T,
     constraint: C,
 }
 
@@ -77,21 +88,21 @@ impl Cache {
             .filter(|entry| entry.hash == hash)
             .map(|entry| {
                 entry
-                    .output
+                    .constrained
                     .downcast_ref::<Constrained<Out, In::Constraint>>()
                     .expect("comemo: a hash collision occurred")
             })
             .find(|output| input.valid(&output.constraint))
-            .map(|output| output.value.clone())
+            .map(|output| output.output.clone())
     }
 
     /// Insert an entry into the cache.
-    fn insert<In, Out>(&self, hash: u128, output: Constrained<Out, In::Constraint>)
+    fn insert<In, Out>(&self, hash: u128, constrained: Constrained<Out, In::Constraint>)
     where
         In: Input,
         Out: 'static,
     {
-        let entry = Entry { hash, output: Box::new(output) };
+        let entry = Entry { hash, constrained: Box::new(constrained) };
         self.map.borrow_mut().push(entry);
     }
 }