Less lifetimes and global cache

Author: laurmaedje

macros/src/lib.rs

@@ -13,7 +13,6 @@ mod memoize;
 mod track;
 
 use proc_macro::TokenStream;
-use proc_macro2::Span;
 use quote::quote;
 use syn::{parse_quote, Error, Result};
 

macros/src/memoize.rs

@@ -2,7 +2,7 @@ use super::*;
 
 /// Memoize a function.
 pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
-    let name = func.sig.ident.to_string();
+    let name = &func.sig.ident;
 
     let mut args = vec![];
     let mut types = vec![];
@@ -23,21 +23,12 @@ pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
         types.push(typed.ty.as_ref());
     }
 
-    let ret = match &func.sig.output {
-        syn::ReturnType::Default => {
-            bail!(func.sig, "function must have a return type")
-        }
-        syn::ReturnType::Type(.., ty) => ty.as_ref(),
-    };
-
-    let mut inner = func.clone();
-    inner.sig.ident = syn::Ident::new("inner", Span::call_site());
-
     if args.len() != 1 {
         bail!(func, "expected exactly one argument");
     }
 
     let arg = args[0];
+    /*
     let ty = types[0];
     let track = match ty {
         syn::Type::Path(path) => {
@@ -59,56 +50,14 @@ pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
         }
         _ => bail!(ty, "expected type of the form `Tracked<_>`"),
     };
-
-    let trackable = quote! {
-        <#track as ::comemo::internal::Trackable<'static>>
-    };
-
-    let body = quote! {
-        type Cache = ::core::cell::RefCell<
-            ::std::vec::Vec<(#trackable::Tracker, #ret)>
-        >;
-
-        thread_local! {
-            static CACHE: Cache = Default::default();
-        }
-
-        let mut hit = true;
-        let output = CACHE.with(|cache| {
-            cache
-                .borrow()
-                .iter()
-                .find(|(tracker, _)| {
-                    let (#arg, _) = ::comemo::internal::to_parts(#arg);
-                    #trackable::valid(#arg, tracker)
-                })
-                .map(|&(_, output)| output)
-        });
-
-        let output = output.unwrap_or_else(|| {
-            let tracker = ::core::default::Default::default();
-            let (#arg, _) = ::comemo::internal::to_parts(#arg);
-            let #arg = ::comemo::internal::from_parts(#arg, Some(&tracker));
-            let output = inner(#arg);
-            CACHE.with(|cache| cache.borrow_mut().push((tracker, output)));
-            hit = false;
-            output
-        });
-
-        println!(
-            "{} {} {}",
-            #name,
-            if hit { "[hit]: " } else { "[miss]:" },
-            output,
-        );
-
-        output
-    };
+    */
 
     let mut outer = func.clone();
     outer.block = parse_quote! { {
-        #inner
-        { #body }
+        #func
+        ::comemo::internal::CACHE.with(|cache|
+            cache.query(stringify!(#name), #name, #arg)
+        )
     } };
 
     Ok(quote! { #outer })

macros/src/track.rs

@@ -46,24 +46,29 @@ pub fn expand(block: &syn::ItemImpl) -> Result<proc_macro2::TokenStream> {
     let track_impl = quote! {
         use super::*;
 
-        impl<'a> ::comemo::Track<'a> for #ty {}
-        impl<'a> ::comemo::internal::Trackable<'a> for #ty {
+        impl ::comemo::Track for #ty {}
+        impl ::comemo::internal::Trackable for #ty {
             type Tracker = Tracker;
-            type Surface = Surface<'a>;
+            type Surface = SurfaceFamily;
 
             fn valid(&self, tracker: &Self::Tracker) -> bool {
                 #(#tracked_valids)&&*
             }
 
-            fn surface<'s>(tracked: &'s Tracked<'a, #ty>) -> &'s Self::Surface
+            fn surface<'a, 'r>(tracked: &'r Tracked<'a, #ty>) -> &'r Surface<'a>
             where
-                Self: Track<'a>,
+                Self: Track,
             {
                 // Safety: Surface is repr(transparent).
-                unsafe { &*(tracked as *const _ as *const Self::Surface) }
+                unsafe { &*(tracked as *const _ as *const _) }
             }
         }
 
+        pub enum SurfaceFamily {}
+        impl<'a> ::comemo::internal::Family<'a> for SurfaceFamily {
+            type Out = Surface<'a>;
+        }
+
         #[repr(transparent)]
         pub struct Surface<'a>(Tracked<'a, #ty>);
 

src/cache.rs

@@ -0,0 +1,73 @@
+use std::any::Any;
+use std::cell::RefCell;
+use std::fmt::Debug;
+
+use crate::track::{from_parts, to_parts, Track, Trackable, Tracked};
+
+thread_local! {
+    /// The global, dynamic cache shared by all memoized functions.
+    pub static CACHE: Cache = Cache::default();
+}
+
+/// An untyped cache.
+#[derive(Default)]
+pub struct Cache {
+    map: RefCell<Vec<Box<dyn Any>>>,
+}
+
+/// An entry in the cache.
+struct Entry<Tracker, R> {
+    tracker: Tracker,
+    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
+    where
+        F: Fn(Tracked<T>) -> R,
+        T: Track,
+        R: Debug + Clone + 'static,
+    {
+        let mut hit = true;
+        let output = self.lookup(tracked).unwrap_or_else(|| {
+            let tracker = T::Tracker::default();
+            let (inner, _) = to_parts(tracked);
+            let tracked = from_parts(inner, Some(&tracker));
+            let output = f(tracked);
+            self.insert::<T, R>(tracker, output.clone());
+            hit = false;
+            output
+        });
+
+        let label = if hit { "[hit]: " } else { "[miss]:" };
+        eprintln!("{name} {label} {output:?}");
+
+        output
+    }
+
+    /// Look for a matching entry in the cache.
+    fn lookup<T, R>(&self, tracked: Tracked<T>) -> Option<R>
+    where
+        T: Track,
+        R: Clone + 'static,
+    {
+        let (inner, _) = to_parts(tracked);
+        self.map
+            .borrow()
+            .iter()
+            .filter_map(|boxed| boxed.downcast_ref::<Entry<T::Tracker, R>>())
+            .find(|entry| Trackable::valid(inner, &entry.tracker))
+            .map(|entry| entry.output.clone())
+    }
+
+    /// Insert an entry into the cache.
+    fn insert<T, R>(&self, tracker: T::Tracker, output: R)
+    where
+        T: Track,
+        R: 'static,
+    {
+        let entry = Entry { tracker, output };
+        self.map.borrow_mut().push(Box::new(entry));
+    }
+}

src/hash.rs

@@ -0,0 +1,15 @@
+use std::hash::Hash;
+
+use siphasher::sip128::{Hasher128, SipHasher};
+use std::num::NonZeroU128;
+
+/// 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())
+}

src/internal.rs

@@ -1,72 +1,15 @@
 //! Tracked memoization.
 
-// These are implementation details. Do not rely on them!
-#[doc(hidden)]
-pub mod internal;
+mod cache;
+mod hash;
+mod track;
 
+pub use crate::track::{Track, Tracked};
 pub use comemo_macros::{memoize, track};
 
-use std::ops::Deref;
-
-/// A trackable type.
-///
-/// This is implemented by types that have an implementation block annoted with
-/// [`#[track]`](track).
-pub trait Track<'a>: internal::Trackable<'a> {
-    /// Start tracking a value.
-    fn track(&'a self) -> Tracked<'a, Self> {
-        Tracked { inner: self, tracker: None }
-    }
-}
-
-/// Tracks accesses to a value.
-///
-/// Encapsulates a reference to a value and tracks all accesses to it.
-/// The only methods accessible on `Tracked<T>` are those defined in an impl
-/// block for `T` annotated with [`#[track]`](track).
-///
-/// ```
-/// use comemo::Track;
-///
-/// let image = Image::random(20, 40);
-/// let sentence = describe(image.track());
-/// println!("{sentence}");
-/// ```
-pub struct Tracked<'a, T>
-where
-    T: Track<'a>,
-{
-    /// A reference to the tracked value.
-    inner: &'a T,
-    /// A tracker which stores constraints for T. It is filled by the tracked
-    /// methods on T's generated surface type.
-    ///
-    /// Starts out as `None` and is set to a stack-stored tracker in the
-    /// preamble of memoized functions.
-    tracker: Option<&'a T::Tracker>,
-}
-
-// The type `Tracked<T>` automatically dereferences to T's generated surface
-// type. This makes all tracked methods available, but leaves all other ones
-// unaccessible.
-impl<'a, T> Deref for Tracked<'a, T>
-where
-    T: Track<'a>,
-{
-    type Target = T::Surface;
-
-    fn deref(&self) -> &Self::Target {
-        T::surface(self)
-    }
-}
-
-impl<'a, T> Copy for Tracked<'a, T> where T: Track<'a> {}
-
-impl<'a, T> Clone for Tracked<'a, T>
-where
-    T: Track<'a>,
-{
-    fn clone(&self) -> Self {
-        *self
-    }
+/// These are implementation details. Do not rely on them!
+#[doc(hidden)]
+pub mod internal {
+    pub use crate::cache::CACHE;
+    pub use crate::track::{from_parts, to_parts, AccessTracker, Family, Trackable};
 }