Basic memoization

Author: laurmaedje

Cargo.toml

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

macros/src/lib.rs

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

macros/src/memoize.rs

@@ -5,7 +5,6 @@ pub fn expand(mut func: syn::ItemFn) -> Result<proc_macro2::TokenStream> {
     let name = func.sig.ident.to_string();
 
     let mut args = vec![];
-    let mut asserts = vec![];
     for input in &func.sig.inputs {
         let typed = match input {
             syn::FnArg::Typed(typed) => typed,
@@ -19,15 +18,68 @@ pub fn expand(mut func: syn::ItemFn) -> Result<proc_macro2::TokenStream> {
             _ => bail!(typed.pat, "only simple identifiers are supported"),
         };
 
-        asserts.push(quote_spanned! { ident.span() => assert_hash(&#ident); });
         args.push(ident);
     }
 
-    func.block.stmts.insert(0, parse_quote! { {
-        println!("calling {}", #name);
-        fn assert_hash<T: std::hash::Hash>(_: &T) {}
-        // #(#asserts)*
-    } });
+    let mut inner = func.clone();
+    inner.sig.ident = syn::Ident::new("inner", Span::call_site());
+
+    let asserts = args.iter().map(|ident| {
+        quote_spanned! {
+            ident.span() =>
+            argument_must_be_hash_or_tracked(&#ident);
+        }
+    });
+
+    let hashes = args.iter().map(|ident| {
+        quote! {
+            #ident.hash(&mut state);
+        }
+    });
+
+    func.block = parse_quote! { { #inner {
+        fn argument_must_be_hash_or_tracked<T: std::hash::Hash>(_: &T) {}
+        #(#asserts)*
+
+        use std::collections::HashMap;
+        use std::hash::Hash;
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        use std::sync::Mutex;
+        use comemo::once_cell::sync::Lazy;
+        use comemo::siphasher::sip128::{Hasher128, SipHasher};
+
+        static NR: AtomicUsize = AtomicUsize::new(1);
+        static MAP: Lazy<Mutex<HashMap<u128, String>>> =
+            Lazy::new(|| Mutex::new(HashMap::new()));
+
+        let nr = NR.fetch_add(1, Ordering::SeqCst);
+        let hash = {
+            let mut state = SipHasher::new();
+            #(#hashes)*
+            state.finish128().as_u128()
+        };
+
+        let mut hit = true;
+        let result = MAP
+            .lock()
+            .unwrap()
+            .entry(hash)
+            .or_insert_with(|| {
+                hit = false;
+                inner(#(#args),*)
+            })
+            .clone();
+
+        println!(
+            "{} {} {} {}",
+            #name,
+            nr,
+            if hit { "[hit]: " } else { "[miss]:" },
+            result,
+        );
+
+        result
+    } } };
 
     Ok(quote! { #func })
 }

src/lib.rs

@@ -4,6 +4,10 @@ Tracked memoization.
 
 pub use comemo_macros::{memoize, track};
 
+// This is an implementation detail, which shouldn't directly be used.
+#[doc(hidden)]
+pub use {once_cell, siphasher};
+
 use std::ops::Deref;
 
 /// Tracks accesses to a value.
@@ -30,15 +34,14 @@ pub trait Track: Sized {
         Tracked(self)
     }
 
+    // Everything below is an implementation detail.
+    // You shouldn't use that stuff directly.
+
     /// The tracked API surface of the type.
-    ///
-    /// This is an implementation detail, which shouldn't directly be used.
     #[doc(hidden)]
     type Surface;
 
     /// Access the tracked API surface.
-    ///
-    /// This is an implementation detail, which shouldn't directly be used.
     #[doc(hidden)]
     fn surface(&self) -> &Self::Surface;
 }