Memoizable methods

Author: laurmaedje

macros/src/lib.rs

@@ -15,7 +15,7 @@ mod track;
 
 use proc_macro::TokenStream as BoundaryStream;
 use proc_macro2::TokenStream;
-use quote::{quote, quote_spanned};
+use quote::{quote, quote_spanned, ToTokens};
 use syn::spanned::Spanned;
 use syn::{parse_quote, Error, Result};
 

macros/src/memoize.rs

@@ -13,38 +13,67 @@ pub fn expand(item: &syn::ItemFn) -> Result<proc_macro2::TokenStream> {
 struct Function {
     item: syn::ItemFn,
     name: syn::Ident,
-    args: Vec<syn::Ident>,
-    types: Vec<syn::Type>,
+    args: Vec<Argument>,
     output: syn::Type,
 }
 
+/// An argument to a memoized function.
+enum Argument {
+    Ident(syn::Ident),
+    Receiver(syn::Token![self]),
+}
+
+impl ToTokens for Argument {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        match self {
+            Self::Ident(ident) => ident.to_tokens(tokens),
+            Self::Receiver(token) => token.to_tokens(tokens),
+        }
+    }
+}
+
 /// Preprocess and validate a function.
 fn prepare(function: &syn::ItemFn) -> Result<Function> {
     let mut args = vec![];
-    let mut types = vec![];
 
     for input in &function.sig.inputs {
-        let typed = match input {
-            syn::FnArg::Typed(typed) => typed,
-            syn::FnArg::Receiver(_) => {
-                bail!(function, "methods are not supported")
+        match input {
+            syn::FnArg::Receiver(recv) => {
+                if recv.mutability.is_some() {
+                    bail!(recv, "memoized functions cannot have mutable parameters");
+                }
+
+                args.push(Argument::Receiver(recv.self_token));
+            }
+            syn::FnArg::Typed(typed) => {
+                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.as_ref().clone();
+                match ty {
+                    syn::Type::Reference(syn::TypeReference {
+                        mutability: Some(_),
+                        ..
+                    }) => {
+                        bail!(
+                            typed.ty,
+                            "memoized functions cannot have mutable parameters"
+                        )
+                    }
+                    _ => {}
+                }
+
+                args.push(Argument::Ident(name));
             }
-        };
-
-        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.as_ref().clone();
-        args.push(name);
-        types.push(ty);
+        }
     }
 
     let output = match &function.sig.output {
@@ -58,36 +87,49 @@ fn prepare(function: &syn::ItemFn) -> Result<Function> {
         item: function.clone(),
         name: function.sig.ident.clone(),
         args,
-        types,
         output,
     })
 }
 
 /// Rewrite a function's body to memoize it.
 fn process(function: &Function) -> Result<TokenStream> {
+    // Construct assertions that the arguments fulfill the necessary bounds.
+    let bounds = function.args.iter().map(|arg| {
+        quote_spanned! { function.item.span() =>
+            ::comemo::internal::assert_hashable_or_trackable(&#arg);
+        }
+    });
+
     // Construct a tuple from all arguments.
-    let args = &function.args;
+    let args = function.args.iter().map(|arg| match arg {
+        Argument::Ident(id) => id.to_token_stream(),
+        Argument::Receiver(token) => quote! {
+            ::comemo::internal::hash(&#token)
+        },
+    });
     let arg_tuple = quote! { (#(#args,)*) };
 
-    // Construct assertions that the arguments fulfill the necessary bounds.
-    let bounds = function.types.iter().map(|ty| {
-        quote! {
-            ::comemo::internal::assert_hashable_or_trackable::<#ty>();
-        }
+    // Construct a tuple for all parameters.
+    let params = function.args.iter().map(|arg| match arg {
+        Argument::Ident(id) => id.to_token_stream(),
+        Argument::Receiver(_) => quote! { _ },
     });
+    let param_tuple = quote! { (#(#params,)*) };
 
     // Construct the inner closure.
     let output = &function.output;
     let body = &function.item.block;
-    let closure = quote! { |#arg_tuple| -> #output #body };
+    let closure = quote! { |#param_tuple| -> #output #body };
 
     // Adjust the function's body.
     let mut wrapped = function.item.clone();
     let name = function.name.to_string();
     wrapped.block = parse_quote! { {
+        struct __ComemoUnique;
         #(#bounds;)*
-        ::comemo::internal::cached(
+        ::comemo::internal::memoized(
             #name,
+            ::core::any::TypeId::of::<__ComemoUnique>(),
             ::comemo::internal::Args(#arg_tuple),
             #closure,
         )

src/cache.rs

@@ -10,16 +10,16 @@ use crate::input::Input;
 use crate::internal::Family;
 
 /// Execute a function or use a cached result for it.
-pub fn cached<In, Out, F>(name: &str, input: In, func: F) -> Out
+pub fn memoized<In, Out, F>(name: &'static str, unique: TypeId, input: In, func: F) -> Out
 where
     In: Input,
     Out: Debug + Clone + 'static,
-    F: for<'f> Fn(<In::Tracked as Family<'f>>::Out) -> Out + 'static,
+    F: for<'f> FnOnce(<In::Tracked as Family<'f>>::Out) -> Out,
 {
     // Compute the hash of the input's key part.
     let hash = {
         let mut state = SipHasher::new();
-        TypeId::of::<F>().hash(&mut state);
+        unique.hash(&mut state);
         input.key(&mut state);
         state.finish128().as_u128()
     };

src/input.rs

@@ -6,7 +6,7 @@ use crate::internal::Family;
 use crate::track::{from_parts, to_parts, Track, Trackable, Tracked};
 
 /// Ensure a type is suitable as input.
-pub fn assert_hashable_or_trackable<In: Input>() {}
+pub fn assert_hashable_or_trackable<In: Input>(_: &In) {}
 
 /// An input to a cached function.
 ///

src/lib.rs

@@ -13,7 +13,7 @@ pub use comemo_macros::{memoize, track};
 /// These are implementation details. Do not rely on them!
 #[doc(hidden)]
 pub mod internal {
-    pub use crate::cache::cached;
+    pub use crate::cache::memoized;
     pub use crate::constraint::{hash, Join, MultiConstraint, SoloConstraint};
     pub use crate::input::{assert_hashable_or_trackable, Args};
     pub use crate::track::{to_parts, Trackable};

src/main.rs

@@ -3,9 +3,7 @@
 use comemo::{Track, Tracked};
 
 // TODO
-// - Bring over Prehashed
 // - Tracked return value from tracked method
-// - Memoized methods
 // - Reporting and evicting
 
 fn main() {
@@ -32,9 +30,11 @@ fn main() {
 
 /// Format the image's size humanly readable.
 fn describe(image: Tracked<Image>) -> &'static str {
-    ::comemo::internal::assert_hashable_or_trackable::<Tracked<Image>>();
-    ::comemo::internal::cached(
+    struct __ComemoUnique;
+    ::comemo::internal::assert_hashable_or_trackable(&image);
+    ::comemo::internal::memoized(
         "describe",
+        ::core::any::TypeId::of::<__ComemoUnique>(),
         ::comemo::internal::Args((image,)),
         |(image,)| {
             if image.width() > 50 || image.height() > 50 {

tests/kinds.rs

@@ -1,3 +1,5 @@
+use std::hash::Hash;
+
 use comemo::Track;
 
 #[test]
@@ -23,6 +25,11 @@ fn test_kinds() {
     unconditional(tracky); // [Miss] The combined length changed.
     conditional(tracky, "World"); // [Miss] "World" is now shorter.
     ignorant(tracky, "Ignorant"); // [Hit] Doesn't depend on `tester`.
+
+    Taker("Hello".into()).take(); // [Miss] Never called.
+    Taker("Hello".into()).copy(); // [Miss] Never called.
+    Taker("World".into()).take(); // [Miss] Different value.
+    Taker("Hello".into()).take(); // [Hit] Same value.
 }
 
 /// Always accesses data from both arguments.
@@ -37,14 +44,17 @@ fn unconditional(tester: Tracky) -> &'static str {
 
 /// Accesses data from both arguments conditionally.
 #[comemo::memoize]
-fn conditional(tester: Tracky, name: &str) -> String {
-    tester.double_ref(name).to_string()
+fn conditional<T>(tester: Tracky, name: T) -> String
+where
+    T: AsRef<str> + Hash,
+{
+    tester.double_ref(name.as_ref()).to_string()
 }
 
 /// Accesses only data from the second argument.
 #[comemo::memoize]
-fn ignorant(tester: Tracky, name: &str) -> String {
-    tester.arg_ref(name).to_string()
+fn ignorant(tester: Tracky, name: impl AsRef<str> + Hash) -> String {
+    tester.arg_ref(name.as_ref()).to_string()
 }
 
 /// Test with type alias.
@@ -81,3 +91,18 @@ impl Tester {
 /// A non-copy struct that is passed by value to a tracked method.
 #[derive(Clone, PartialEq)]
 struct Heavy(String);
+
+#[derive(Hash)]
+struct Taker(String);
+
+impl Taker {
+    #[comemo::memoize]
+    fn copy(&self) -> String {
+        self.0.clone()
+    }
+
+    #[comemo::memoize]
+    fn take(self) -> String {
+        self.0
+    }
+}

tests/layout.rs

@@ -0,0 +1,114 @@
+use comemo::{Prehashed, Track, Tracked};
+
+#[test]
+fn test_layout() {
+    let par = Paragraph(vec![
+        TextRun {
+            font: "Helvetica".into(),
+            text: "HELLO ".into(),
+        },
+        TextRun {
+            font: "Futura".into(),
+            text: "WORLD!".into(),
+        },
+    ]);
+
+    let mut fonts = Fonts::default();
+    fonts.insert("Helvetica", Style::Normal, vec![110; 75398]);
+    fonts.insert("Futura", Style::Italic, vec![55; 12453]);
+
+    // [Miss] The cache is empty.
+    par.layout(fonts.track());
+    fonts.insert("Verdana", Style::Normal, vec![99; 12554]);
+
+    // [Hit] Verdana isn't used.
+    par.layout(fonts.track());
+    fonts.insert("Helvetica", Style::Bold, vec![120; 98532]);
+
+    // [Miss] Helvetica changed.
+    par.layout(fonts.track());
+}
+
+/// A paragraph composed from text runs.
+#[derive(Debug, Hash)]
+struct Paragraph(Vec<TextRun>);
+
+impl Paragraph {
+    /// A memoized method.
+    #[comemo::memoize]
+    fn layout(&self, fonts: Tracked<Fonts>) -> String {
+        let mut result = String::new();
+        for run in &self.0 {
+            let font = fonts.select(&run.font).unwrap();
+            for c in run.text.chars() {
+                result.push(font.map(c));
+            }
+        }
+        result
+    }
+}
+
+/// A run of text with consistent font.
+#[derive(Debug, Hash)]
+struct TextRun {
+    font: String,
+    text: String,
+}
+
+/// Holds all fonts.
+///
+/// As font data is large and costly to hash, we use the `Prehashed` wrapper.
+/// Otherwise, every call to `Fonts::select` would hash the returned font from
+/// scratch.
+#[derive(Default)]
+struct Fonts(Vec<Prehashed<Font>>);
+
+impl Fonts {
+    /// Insert a new with name and data.
+    fn insert(&mut self, name: impl Into<String>, style: Style, data: Vec<u8>) {
+        let name = name.into();
+        self.0.retain(|font| font.name != name);
+        self.0.push(Prehashed::new(Font { name, style, data }))
+    }
+}
+
+#[comemo::track]
+impl Fonts {
+    /// Select a font by name.
+    fn select(&self, name: &str) -> Option<&Prehashed<Font>> {
+        self.0.iter().find(|font| font.name == name)
+    }
+}
+
+/// A large binary font.
+#[derive(Hash)]
+struct Font {
+    name: String,
+    data: Vec<u8>,
+    style: Style,
+}
+
+impl Font {
+    /// Map a character.
+    fn map(&self, c: char) -> char {
+        let base = match self.style {
+            Style::Normal => 0x41,
+            Style::Bold => 0x1D400,
+            Style::Italic => 0x1D434,
+        };
+
+        if c.is_ascii_alphabetic() {
+            std::char::from_u32(base + (c as u32 - 0x41)).unwrap()
+        } else {
+            c
+        }
+    }
+}
+
+/// A font style.
+#[derive(Hash)]
+enum Style {
+    Normal,
+    Italic,
+    Bold,
+}