Basic fully macro-generated memoization

Author: laurmaedje

Cargo.toml

@@ -6,4 +6,3 @@ edition = "2021"
 [dependencies]
 comemo-macros = { path = "./macros" }
 siphasher = "0.3"
-once_cell = "1"

examples/image.rs

@@ -31,6 +31,17 @@ fn describe(image: Tracked<Image>) -> &'static str {
     }
 }
 
+#[comemo::track]
+impl Image {
+    fn width(&self) -> u32 {
+        self.width
+    }
+
+    fn height(&self) -> u32 {
+        self.height
+    }
+}
+
 /// A raster image.
 struct Image {
     width: u32,
@@ -53,14 +64,3 @@ impl Image {
         // Resize the actual image ...
     }
 }
-
-#[comemo::track]
-impl Image {
-    fn width(&self) -> u32 {
-        self.width
-    }
-
-    fn height(&self) -> u32 {
-        self.height
-    }
-}

macros/src/memoize.rs

@@ -1,10 +1,11 @@
 use super::*;
 
-/// Make a type trackable.
+/// Memoize a function.
 pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
     let name = func.sig.ident.to_string();
 
     let mut args = vec![];
+    let mut types = vec![];
     for input in &func.sig.inputs {
         let typed = match input {
             syn::FnArg::Typed(typed) => typed,
@@ -19,41 +20,96 @@ pub fn expand(func: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
         };
 
         args.push(ident);
+        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());
 
-    let cts = args.iter().map(|arg| {
-        quote! {
-            Validate::constraint(&#arg)
+    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) => {
+            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();
-    outer.block = parse_quote! { { #inner {
-        use std::sync::atomic::{AtomicUsize, Ordering};
-        use comemo::Validate;
+    let trackable = quote! {
+        <#track as ::comemo::internal::Trackable<'static>>
+    };
+
+    let body = quote! {
+        type Cache = ::core::cell::RefCell<
+            ::std::vec::Vec<(#trackable::Tracker, #ret)>
+        >;
 
-        static NR: AtomicUsize = AtomicUsize::new(1);
-        let nr = NR.fetch_add(1, Ordering::SeqCst);
-        let cts = (#(#cts,)*);
+        thread_local! {
+            static CACHE: Cache = Default::default();
+        }
 
-        println!("{:?}", cts);
+        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 mut hit = false;
-        let result = inner(#(#args),*);
+        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,
-            nr,
             if hit { "[hit]: " } else { "[miss]:" },
-            result,
+            output,
         );
 
-        result
-    } } };
+        output
+    };
+
+    let mut outer = func.clone();
+    outer.block = parse_quote! { {
+        #inner
+        { #body }
+    } };
 
     Ok(quote! { #outer })
 }

macros/src/track.rs

@@ -10,64 +10,76 @@ 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::AccessTracker<#ty>, }
+    });
+
     let tracked_methods = methods.iter().map(|method| {
-        let mut method = (*method).clone();
         let name = &method.sig.ident;
+        let mut method = (*method).clone();
+        if matches!(method.vis, syn::Visibility::Inherited) {
+            method.vis = parse_quote! { pub(super) };
+        }
         method.block = parse_quote! { {
-            let output = self.0.#name();
-            let slot = &mut self.constraint().#name;
-            let ct = Validate::constraint(&output);
-            if slot.is_none() {
-                assert_eq!(*slot, Some(ct), "comemo: method is not pure");
+            let (inner, tracker) = ::comemo::internal::to_parts(self.0);
+            let output = inner.#name();
+            if let Some(tracker) = &tracker {
+                tracker.#name.track(&output);
             }
-            *slot = Some(ct);
             output
         } };
         method
     });
 
-    let tracked_fields = methods.iter().map(|method| {
+    let tracked_valids = 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: Option<<#ty as Validate>::Constraint>, }
+        quote! {
+            tracker.#name.valid(&self.#name())
+        }
     });
 
     let track_impl = quote! {
-        use comemo::Validate;
+        use super::*;
 
-        struct Surface<'a>(&'a #ty);
-
-        impl Surface<'_> {
-            #(#tracked_methods)*
+        impl<'a> ::comemo::Track<'a> for #ty {}
+        impl<'a> ::comemo::internal::Trackable<'a> for #ty {
+            type Tracker = Tracker;
+            type Surface = Surface<'a>;
 
-            fn constraint(&self) -> &mut Constraint {
-                todo!()
+            fn valid(&self, tracker: &Self::Tracker) -> bool {
+                #(#tracked_valids)&&*
             }
-        }
 
-        impl<'a> From<&'a #ty> for Surface<'a> {
-            fn from(val: &'a #ty) -> Self {
-                Self(val)
+            fn surface<'s>(tracked: &'s Tracked<'a, #ty>) -> &'s Self::Surface
+            where
+                Self: Track<'a>,
+            {
+                // Safety: Surface is repr(transparent).
+                unsafe { &*(tracked as *const _ as *const Self::Surface) }
             }
         }
 
-        #[derive(Debug, Default)]
-        struct Constraint {
-            #(#tracked_fields)*
+        #[repr(transparent)]
+        pub struct Surface<'a>(Tracked<'a, #ty>);
+
+        impl Surface<'_> {
+            #(#tracked_methods)*
         }
 
-        impl<'a> comemo::Track<'a> for #ty {
-            type Surface = Surface<'a>;
-            type Constraint = Constraint;
+        #[derive(Default)]
+        pub struct Tracker {
+            #(#tracked_fields)*
         }
     };
 
     Ok(quote! {
         #block
-        const _: () = { #track_impl };
+        const _: () = { mod private { #track_impl } };
     })
 }
 
@@ -78,6 +90,12 @@ fn method(item: &syn::ImplItem) -> Result<&syn::ImplItemMethod> {
         _ => bail!(item, "only methods are supported"),
     };
 
+    match method.vis {
+        syn::Visibility::Inherited => {}
+        syn::Visibility::Public(_) => {}
+        _ => bail!(method.vis, "only private and public methods are supported"),
+    }
+
     let mut inputs = method.sig.inputs.iter();
     let receiver = match inputs.next() {
         Some(syn::FnArg::Receiver(recv)) => recv,

src/internal.rs

@@ -0,0 +1,66 @@
+use std::cell::Cell;
+use std::hash::Hash;
+use std::marker::PhantomData;
+use std::num::NonZeroU128;
+
+use siphasher::sip128::{Hasher128, SipHasher};
+
+use super::{Track, Tracked};
+
+/// Destructure a tracker into its parts.
+pub fn to_parts<'a, T>(tracked: Tracked<'a, T>) -> (&'a T, Option<&'a T::Tracker>)
+where
+    T: Track<'a>,
+{
+    (tracked.inner, tracked.tracker)
+}
+
+/// Create a tracker from its parts.
+pub fn from_parts<'a, T>(inner: &'a T, tracker: Option<&'a T::Tracker>) -> Tracked<'a, T>
+where
+    T: Track<'a>,
+{
+    Tracked { inner, tracker }
+}
+
+/// Non-exposed parts of the `Track` trait.
+pub trait Trackable<'a>: Sized + 'a {
+    /// Keeps track of accesses to the value.
+    type Tracker: Default;
+
+    /// The tracked API surface of this type.
+    type Surface;
+
+    /// Whether an instance fulfills the given tracker's constraints.
+    fn valid(&self, tracker: &Self::Tracker) -> bool;
+
+    /// Cast a reference from `Tracked` to this type's surface.
+    fn surface<'s>(tracked: &'s Tracked<'a, Self>) -> &'s Self::Surface
+    where
+        Self: Track<'a>;
+}
+
+/// Tracks accesses to a value.
+#[derive(Default)]
+pub struct AccessTracker<T: Hash>(Cell<Option<NonZeroU128>>, PhantomData<T>);
+
+impl<T: Hash> AccessTracker<T> {
+    pub fn track(&self, value: &T) {
+        self.0.set(Some(siphash(value)));
+    }
+
+    pub fn valid(&self, value: &T) -> bool {
+        self.0.get().map_or(true, |v| v == siphash(value))
+    }
+}
+
+/// Produce a non zero 128-bit hash of a value.
+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())
+}