Support multiple arguments

Author: laurmaedje

macros/src/lib.rs

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

macros/src/memoize.rs

@@ -3,6 +3,7 @@ use super::*;
 /// Memoize a function.
 pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
     let mut args = vec![];
+    let mut types = vec![];
     for input in &func.sig.inputs {
         let typed = match input {
             syn::FnArg::Typed(typed) => typed,
@@ -17,18 +18,29 @@ pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
         };
 
         args.push(name);
+        types.push(typed.ty.as_ref());
     }
 
+    let mut inner = func.clone();
+    let arg_tuple = quote! { (#(#args,)*) };
+    let type_tuple = quote! { (#(#types,)*) };
+    inner.sig.inputs = parse_quote! { #arg_tuple: #type_tuple };
+
+    let bounds = args.iter().zip(&types).map(|(arg, ty)| {
+        quote_spanned! {
+            arg.span() => ::comemo::internal::assert_hashable_or_trackable::<#ty>();
+        }
+    });
+
     let mut outer = func.clone();
     let name = &func.sig.ident;
-    let arg = &args[0];
-
     outer.block = parse_quote! { {
-        #func
+        #inner
+        #(#bounds;)*
         ::comemo::internal::CACHE.with(|cache|
             cache.query(
                 stringify!(#name),
-                #arg,
+                ::comemo::internal::Args(#arg_tuple),
                 #name,
             )
         )

src/cache.rs

@@ -36,7 +36,7 @@ impl Cache {
     where
         In: Input,
         Out: Debug + Clone + 'static,
-        F: for<'f> Fn(<In::Hooked as Family<'f>>::Out) -> Out,
+        F: for<'f> Fn(<In::Tracked as Family<'f>>::Out) -> Out,
     {
         // Compute the hash of the input's key part.
         let hash = {
@@ -48,8 +48,7 @@ impl Cache {
         let mut hit = true;
         let output = self.lookup::<In, Out>(hash, &input).unwrap_or_else(|| {
             let constraint = In::Constraint::default();
-            let input = input.hook_up(&constraint);
-            let value = func(input);
+            let value = func(input.track(&constraint));
             let constrained = Constrained { value: value.clone(), constraint };
             self.insert::<In, Out>(hash, constrained);
             hit = false;

src/input.rs

@@ -4,35 +4,38 @@ use std::marker::PhantomData;
 use crate::internal::Family;
 use crate::track::{from_parts, to_parts, Track, Trackable, Tracked};
 
-/// An argument to a cached function.
+/// Ensure a type is suitable as input.
+pub fn assert_hashable_or_trackable<T: Input>() {}
+
+/// An input to a cached function.
 pub trait Input {
-    /// Describes an instance of the argument.
+    /// Describes an instance of this input.
     type Constraint: Default + 'static;
 
-    /// The argument with constraints hooked in.
-    type Hooked: for<'f> Family<'f>;
+    /// The input with constraints hooked in.
+    type Tracked: for<'f> Family<'f>;
 
-    /// Hash the argument if it is a _key_ argument
+    /// Hash the _key_ parts of the input.
     fn key<H: Hasher>(&self, state: &mut H);
 
-    /// Validate the argument if it is a _tracked_ argument.
+    /// Validate the _tracked_ parts of the input.
     fn valid(&self, constraint: &Self::Constraint) -> bool;
 
-    /// Hook up the given constraints if this is a _tracked_ argument.
-    fn hook_up<'f>(
+    /// Hook up the given constraint to the _tracked_ parts of the input.
+    fn track<'f>(
         self,
         constraint: &'f Self::Constraint,
-    ) -> <Self::Hooked as Family<'f>>::Out
+    ) -> <Self::Tracked as Family<'f>>::Out
     where
         Self: 'f;
 }
 
 impl<T: Hash> Input for T {
-    /// No constraint for hashed arguments.
+    /// No constraint for hashed inputs.
     type Constraint = ();
 
     /// The hooked-up type is just `Self`.
-    type Hooked = IdFamily<Self>;
+    type Tracked = IdFamily<Self>;
 
     fn key<H: Hasher>(&self, state: &mut H) {
         Hash::hash(self, state);
@@ -42,45 +45,93 @@ impl<T: Hash> Input for T {
         true
     }
 
-    fn hook_up<'f>(self, _: &'f ()) -> Self
+    fn track<'f>(self, _: &'f ()) -> Self
     where
         Self: 'f,
     {
         self
     }
 }
 
+/// Identity type constructor.
+pub struct IdFamily<T>(PhantomData<T>);
+
+impl<T> Family<'_> for IdFamily<T> {
+    type Out = T;
+}
+
 impl<'a, T: Track> Input for Tracked<'a, T> {
-    /// Forwarded constraint from Trackable implementation.
+    /// Forward constraint from `Trackable` implementation.
     type Constraint = <T as Trackable>::Constraint;
 
     /// The hooked-up type is `Tracked<'f, T>`.
-    type Hooked = TrackedFamily<T>;
+    type Tracked = TrackedFamily<T>;
 
     fn key<H: Hasher>(&self, _: &mut H) {}
 
     fn valid(&self, constraint: &Self::Constraint) -> bool {
         Trackable::valid(to_parts(*self).0, constraint)
     }
 
-    fn hook_up<'f>(self, constraint: &'f Self::Constraint) -> Tracked<'f, T>
+    fn track<'f>(self, constraint: &'f Self::Constraint) -> Tracked<'f, T>
     where
         Self: 'f,
     {
         from_parts(to_parts(self).0, Some(constraint))
     }
 }
 
-/// Identity type constructor.
-pub struct IdFamily<T>(PhantomData<T>);
-
-impl<T> Family<'_> for IdFamily<T> {
-    type Out = T;
-}
-
 /// 'f -> Tracked<'f, T> type constructor.
 pub struct TrackedFamily<T>(PhantomData<T>);
 
 impl<'f, T: Track + 'f> Family<'f> for TrackedFamily<T> {
     type Out = Tracked<'f, T>;
 }
+
+/// Wrapper for multiple inputs.
+pub struct Args<T>(pub T);
+
+/// Lifetime to tuple of arguments type constructor.
+pub struct ArgsFamily<T>(PhantomData<T>);
+
+macro_rules! args_input {
+    ($($idx:tt: $letter:ident),*) => {
+        #[allow(unused_variables)]
+        impl<$($letter: Input),*> Input for Args<($($letter,)*)> {
+            type Constraint = ($($letter::Constraint,)*);
+            type Tracked = ArgsFamily<($($letter,)*)>;
+
+            fn key<H: Hasher>(&self, state: &mut H) {
+                $((self.0).$idx.key(state);)*
+            }
+
+            fn valid(&self, constraint: &Self::Constraint) -> bool {
+                true $(&& (self.0).$idx.valid(&constraint.$idx))*
+            }
+
+            fn track<'f>(
+                self,
+                constraint: &'f Self::Constraint,
+            ) -> <Self::Tracked as Family<'f>>::Out
+            where
+                Self: 'f,
+            {
+                ($((self.0).$idx.track(&constraint.$idx),)*)
+            }
+        }
+
+        #[allow(unused_parens)]
+        impl<'f, $($letter: Input),*> Family<'f> for ArgsFamily<($($letter,)*)> {
+            type Out = ($(<$letter::Tracked as Family<'f>>::Out,)*);
+        }
+    };
+}
+
+args_input! {}
+args_input! { 0: A }
+args_input! { 0: A, 1: B }
+args_input! { 0: A, 1: B, 2: C }
+args_input! { 0: A, 1: B, 2: C, 3: D }
+args_input! { 0: A, 1: B, 2: C, 3: D, 4: E }
+args_input! { 0: A, 1: B, 2: C, 3: D, 4: E, 5: F }
+args_input! { 0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G }

src/lib.rs

@@ -13,6 +13,7 @@ pub use comemo_macros::{memoize, track};
 pub mod internal {
     pub use crate::cache::CACHE;
     pub use crate::hash::HashConstraint;
+    pub use crate::input::{assert_hashable_or_trackable, Args};
     pub use crate::track::{from_parts, to_parts, Trackable};
 
     /// Helper trait for lifetime type families.

src/main.rs

@@ -29,16 +29,21 @@ fn main() {
 
 /// Format the image's size humanly readable.
 fn describe(image: Tracked<Image>) -> &'static str {
-    fn describe(image: Tracked<Image>) -> &'static str {
+    fn describe((image,): (Tracked<Image>,)) -> &'static str {
         if image.width() > 50 || image.height() > 50 {
             "The image is big!"
         } else {
             "The image is small!"
         }
     }
 
-    ::comemo::internal::CACHE
-        .with(|cache| cache.query(stringify!(describe), image, describe))
+    ::comemo::internal::CACHE.with(|cache| {
+        cache.query(
+            stringify!(describe),
+            ::comemo::internal::Args((image,)),
+            describe,
+        )
+    })
 }
 
 const _: () = {

tests/image.rs

@@ -4,22 +4,20 @@ use comemo::{Track, Tracked};
 fn test_image() {
     let mut image = Image::new(20, 40);
 
-    // [Miss] The cache is empty.
-    describe(image.track());
-
-    // [Hit] Everything stayed the same.
-    describe(image.track());
+    describe(image.track()); // [Miss] The cache is empty.
+    describe(image.track()); // [Hit] Everything stayed the same.
 
     image.resize(80, 30);
 
-    // [Miss] The image's width and height are different.
-    describe(image.track());
+    describe(image.track()); // [Miss] Width and height changed.
+    select(image.track(), "width"); // [Miss] First call.
+    select(image.track(), "height"); // [Miss]
 
     image.resize(80, 70);
     image.pixels.fill(255);
 
-    // [Hit] The last call only read the width and it stayed the same.
-    describe(image.track());
+    describe(image.track()); // [Hit] Width is > 50 stayed the same.
+    select(image.track(), "width"); // [Hit] Width stayed the same.
 }
 
 /// Format the image's size humanly readable.
@@ -32,6 +30,16 @@ fn describe(image: Tracked<Image>) -> &'static str {
     }
 }
 
+/// Select either width or height.
+#[comemo::memoize]
+fn select(image: Tracked<Image>, what: &str) -> u32 {
+    match what {
+        "width" => image.width(),
+        "height" => image.height(),
+        _ => panic!("there is nothing else!"),
+    }
+}
+
 #[comemo::track]
 impl Image {
     fn width(&self) -> u32 {

tests/simple.rs

@@ -0,0 +1,33 @@
+#[test]
+fn test_simple() {
+    empty(); // [Miss] The cache is empty.
+    empty(); // [Hit] Always a hit from now on.
+    empty(); // [Hit] Always a hit from now on.
+
+    double(2); // [Miss] The cache is empty.
+    double(4); // [Miss] Different number.
+    double(2); // [Hit] Same number as initially.
+
+    sum(2, 4); // [Miss] The cache is empty.
+    sum(2, 3); // [Miss] Different numbers.
+    sum(2, 3); // [Hit]  Same numbers
+    sum(4, 2); // [Miss] Different numbers.
+}
+
+/// Build a string.
+#[comemo::memoize]
+fn empty() -> String {
+    format!("The world is {}", "big")
+}
+
+/// Double a number.
+#[comemo::memoize]
+fn double(x: u32) -> u32 {
+    2 * x
+}
+
+/// Compute the sum of two numbers.
+#[comemo::memoize]
+fn sum(a: u32, b: u32) -> u32 {
+    a + b
+}