Allow tracked or hashed argument

Author: laurmaedje

macros/src/memoize.rs

@@ -2,10 +2,7 @@ use super::*;
 
 /// Memoize a function.
 pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
-    let name = &func.sig.ident;
-
     let mut args = vec![];
-    let mut types = vec![];
     for input in &func.sig.inputs {
         let typed = match input {
             syn::FnArg::Typed(typed) => typed,
@@ -14,47 +11,26 @@ pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
             }
         };
 
-        let ident = match &*typed.pat {
+        let name = match &*typed.pat {
             syn::Pat::Ident(ident) => ident,
             _ => bail!(typed.pat, "only simple identifiers are supported"),
         };
 
-        args.push(ident);
-        types.push(typed.ty.as_ref());
+        args.push(name);
     }
 
-    if args.len() != 1 {
-        bail!(func, "expected exactly one argument");
-    }
-
-    let arg = args[0];
-    let ty = types[0];
-    let _inner = match ty {
-        syn::Type::Path(path) => {
-            let segs = &path.path.segments;
-            if segs.len() != 1 {
-                bail!(ty, "expected exactly one path segment")
-            }
-            let args = match &segs[0].arguments {
-                syn::PathArguments::AngleBracketed(args) => &args.args,
-                _ => bail!(ty, "expected `Tracked<_>` type"),
-            };
-            if args.len() != 1 {
-                bail!(args, "expected exactly one generic argument")
-            }
-            match &args[0] {
-                syn::GenericArgument::Type(ty) => ty,
-                ty => bail!(ty, "expected type argument"),
-            }
-        }
-        _ => bail!(ty, "expected type of the form `Tracked<_>`"),
-    };
-
     let mut outer = func.clone();
+    let name = &func.sig.ident;
+    let arg = &args[0];
+
     outer.block = parse_quote! { {
         #func
         ::comemo::internal::CACHE.with(|cache|
-            cache.query(stringify!(#name), #name, #arg)
+            cache.query(
+                stringify!(#name),
+                #arg,
+                #name,
+            )
         )
     } };
 

macros/src/track.rs

@@ -52,7 +52,7 @@ pub fn expand(block: &syn::ItemImpl) -> Result<proc_macro2::TokenStream> {
             type Surface = SurfaceFamily;
 
             fn valid(&self, constraint: &Self::Constraint) -> bool {
-                #(#tracked_valids)&&*
+                true #(&& #tracked_valids)*
             }
 
             fn surface<'a, 'r>(tracked: &'r Tracked<'a, #ty>) -> &'r Surface<'a>

src/cache.rs

@@ -1,6 +1,10 @@
 use std::any::Any;
 use std::cell::RefCell;
 use std::fmt::Debug;
+use std::hash::{Hash, Hasher};
+use std::marker::PhantomData;
+
+use siphasher::sip128::{Hasher128, SipHasher};
 
 use crate::track::{from_parts, to_parts, Track, Trackable, Tracked};
 
@@ -12,32 +16,45 @@ thread_local! {
 /// An untyped cache.
 #[derive(Default)]
 pub struct Cache {
-    map: RefCell<Vec<Box<dyn Any>>>,
+    map: RefCell<Vec<Entry>>,
 }
 
 /// An entry in the cache.
-struct Entry<C, R> {
+struct Entry {
+    hash: u128,
+    output: Box<dyn Any>,
+}
+
+/// A value with a constraint.
+struct Constrained<T, C> {
+    value: T,
     constraint: C,
-    output: R,
 }
 
 impl Cache {
     /// Execute `f` or use a cached result for it.
-    pub fn query<F, T, R>(&self, name: &'static str, f: F, tracked: Tracked<T>) -> R
+    pub fn query<In, Out, F>(&self, name: &str, input: In, func: F) -> Out
     where
-        F: Fn(Tracked<T>) -> R,
-        T: Track,
-        R: Debug + Clone + 'static,
+        In: Input,
+        Out: Debug + Clone + 'static,
+        F: for<'f> Fn(<In::Focus as Family<'f>>::Out) -> Out,
     {
+        // Compute the hash of the input's key part.
+        let hash = {
+            let mut state = SipHasher::new();
+            input.hash(&mut state);
+            state.finish128().as_u128()
+        };
+
         let mut hit = true;
-        let output = self.lookup::<T, R>(tracked).unwrap_or_else(|| {
-            let constraint = T::Constraint::default();
-            let (inner, _) = to_parts(tracked);
-            let tracked = from_parts(inner, Some(&constraint));
-            let output = f(tracked);
-            self.insert::<T, R>(constraint, output.clone());
+        let output = self.lookup::<In, Out>(hash, &input).unwrap_or_else(|| {
+            let constraint = In::Constraint::default();
+            let input = input.focus(&constraint);
+            let value = func(input);
+            let constrained = Constrained { value: value.clone(), constraint };
+            self.insert::<In, Out>(hash, constrained);
             hit = false;
-            output
+            value
         });
 
         let label = if hit { "[hit]" } else { "[miss]" };
@@ -47,27 +64,99 @@ impl Cache {
     }
 
     /// Look for a matching entry in the cache.
-    fn lookup<T, R>(&self, tracked: Tracked<T>) -> Option<R>
+    fn lookup<In, Out>(&self, hash: u128, input: &In) -> Option<Out>
     where
-        T: Track,
-        R: Clone + 'static,
+        In: Input,
+        Out: Clone + 'static,
     {
-        let (inner, _) = to_parts(tracked);
         self.map
             .borrow()
             .iter()
-            .filter_map(|boxed| boxed.downcast_ref::<Entry<T::Constraint, R>>())
-            .find(|entry| Trackable::valid(inner, &entry.constraint))
-            .map(|entry| entry.output.clone())
+            .filter(|entry| entry.hash == hash)
+            .map(|entry| {
+                entry
+                    .output
+                    .downcast_ref::<Constrained<Out, In::Constraint>>()
+                    .expect("comemo: hash collision")
+            })
+            .find(|output| input.valid(&output.constraint))
+            .map(|output| output.value.clone())
     }
 
     /// Insert an entry into the cache.
-    fn insert<T, R>(&self, constraint: T::Constraint, output: R)
+    fn insert<In, Out>(&self, hash: u128, output: Constrained<Out, In::Constraint>)
     where
-        T: Track,
-        R: 'static,
+        In: Input,
+        Out: 'static,
     {
-        let entry = Entry { constraint, output };
-        self.map.borrow_mut().push(Box::new(entry));
+        let entry = Entry { hash, output: Box::new(output) };
+        self.map.borrow_mut().push(entry);
     }
 }
+
+pub trait Input {
+    type Constraint: Default + 'static;
+    type Focus: for<'f> Family<'f>;
+
+    fn hash<H: Hasher>(&self, state: &mut H);
+    fn valid(&self, constraint: &Self::Constraint) -> bool;
+    fn focus<'f>(
+        self,
+        constraint: &'f Self::Constraint,
+    ) -> <Self::Focus as Family<'f>>::Out
+    where
+        Self: 'f;
+}
+
+impl<T: Hash> Input for T {
+    type Constraint = ();
+    type Focus = HashFamily<T>;
+
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        Hash::hash(self, state);
+    }
+
+    fn valid(&self, _: &()) -> bool {
+        true
+    }
+
+    fn focus<'f>(self, _: &'f ()) -> Self
+    where
+        Self: 'f,
+    {
+        self
+    }
+}
+
+pub struct HashFamily<T>(PhantomData<T>);
+impl<T> Family<'_> for HashFamily<T> {
+    type Out = T;
+}
+
+impl<'a, T: Track> Input for Tracked<'a, T> {
+    type Constraint = <T as Trackable>::Constraint;
+    type Focus = TrackedFamily<T>;
+
+    fn hash<H: Hasher>(&self, _: &mut H) {}
+
+    fn valid(&self, constraint: &Self::Constraint) -> bool {
+        Trackable::valid(to_parts(*self).0, constraint)
+    }
+
+    fn focus<'f>(self, constraint: &'f Self::Constraint) -> Tracked<'f, T>
+    where
+        Self: 'f,
+    {
+        from_parts(to_parts(self).0, Some(constraint))
+    }
+}
+
+pub struct TrackedFamily<T>(PhantomData<T>);
+impl<'f, T: Track + 'f> Family<'f> for TrackedFamily<T> {
+    type Out = Tracked<'f, T>;
+}
+
+pub trait Family<'a> {
+    /// The surface with lifetime.
+    type Out;
+}

src/hash.rs

@@ -1,36 +1,31 @@
 use std::cell::Cell;
 use std::hash::Hash;
 use std::marker::PhantomData;
-use std::num::NonZeroU128;
 
 use siphasher::sip128::{Hasher128, SipHasher};
 
-/// Produce a non zero 128-bit hash of a value.
-pub fn siphash<T: Hash>(value: &T) -> NonZeroU128 {
-    let mut state = SipHasher::new();
-    value.hash(&mut state);
-    state
-        .finish128()
-        .as_u128()
-        .try_into()
-        .unwrap_or(NonZeroU128::new(u128::MAX).unwrap())
-}
-
 /// Defines a constraint for a value through its hash.
 #[derive(Default)]
 pub struct HashConstraint<T: Hash> {
-    cell: Cell<Option<NonZeroU128>>,
+    cell: Cell<Option<u128>>,
     marker: PhantomData<T>,
 }
 
 impl<T: Hash> HashConstraint<T> {
     /// Set the constraint for the value.
     pub fn set(&self, value: &T) {
-        self.cell.set(Some(siphash(value)));
+        self.cell.set(Some(hash(value)));
     }
 
     /// Whether the value fulfills the constraint.
     pub fn valid(&self, value: &T) -> bool {
-        self.cell.get().map_or(true, |v| v == siphash(value))
+        self.cell.get().map_or(true, |v| v == hash(value))
     }
 }
+
+/// Produce a non zero 128-bit hash of a value.
+pub fn hash<T: Hash>(value: &T) -> u128 {
+    let mut state = SipHasher::new();
+    value.hash(&mut state);
+    state.finish128().as_u128()
+}

src/lib.rs

@@ -12,5 +12,5 @@ pub use comemo_macros::{memoize, track};
 pub mod internal {
     pub use crate::cache::CACHE;
     pub use crate::hash::HashConstraint;
-    pub use crate::track::{to_parts, Family, Trackable};
+    pub use crate::track::{from_parts, to_parts, Family, Trackable};
 }

src/main.rs

@@ -38,7 +38,7 @@ fn describe(image: Tracked<Image>) -> &'static str {
     }
 
     ::comemo::internal::CACHE
-        .with(|cache| cache.query(stringify!(describe), describe, image))
+        .with(|cache| cache.query(stringify!(describe), image, describe))
 }
 
 const _: () = {

src/track.rs

@@ -75,7 +75,7 @@ pub trait Track: Trackable {
 }
 
 /// Non-exposed parts of the `Track` trait.
-pub trait Trackable {
+pub trait Trackable: 'static {
     /// Describes an instance of type.
     type Constraint: Default + 'static;