Enhances the memoization mechanism by introducing a traceable operator stack to manage dependencies and cache invalidation.

- Replaces simple integer keys with function pointers for better traceability.
- Adds `MemoOperator` enums to handle stack operations.
- Adds functionality to remove dependent entries from the cache when necessary.

Author: fxdmhtt

cache/src/lib.rs

@@ -1,31 +1,52 @@
 use lru::LruCache;
 use std::{any::Any, num::NonZeroUsize, rc::Rc};
 
+pub enum Trace {
+    Push,
+    Pop,
+}
+
 pub enum MemoOperator {
-    Memo,
+    Memo(Trace),
+    Pop,
 }
 
+pub type OperatorFunc = fn(MemoOperator);
+
 const CACHE_CAP: usize = 128;
 
-static mut CACHE: Option<LruCache<usize, Rc<dyn Any>>> = None;
+static mut CACHE: Option<LruCache<OperatorFunc, Rc<dyn Any>>> = None;
 
-fn cache() -> &'static mut LruCache<usize, Rc<dyn Any>> {
+static mut CALL_STACK: Option<Vec<OperatorFunc>> = None;
+
+fn cache() -> &'static mut LruCache<OperatorFunc, Rc<dyn Any>> {
     #[allow(static_mut_refs)]
     unsafe {
         CACHE.get_or_insert_with(|| LruCache::new(NonZeroUsize::new(CACHE_CAP).unwrap()))
     }
 }
 
-pub fn touch<T: 'static>(key: usize) -> Option<Rc<T>> {
+pub fn call_stack() -> &'static mut Vec<OperatorFunc> {
+    #[allow(static_mut_refs)]
+    unsafe {
+        CALL_STACK.get_or_insert_with(|| Vec::new())
+    }
+}
+
+pub fn touch<T: 'static>(key: OperatorFunc) -> Option<Rc<T>> {
     cache()
         .get(&key)
         .map(Rc::clone)
         .filter(|rc| rc.is::<T>())
         .map(|rc| unsafe { Rc::from_raw(Rc::into_raw(rc) as *const T) })
 }
 
-pub fn store_in_cache<T: 'static>(key: usize, val: T) -> Rc<T> {
+pub fn store_in_cache<T: 'static>(key: OperatorFunc, val: T) -> Rc<T> {
     let rc = Rc::new(val);
     cache().put(key, Rc::clone(&rc) as _);
     rc
 }
+
+pub fn remove_from_cache(key: OperatorFunc) {
+    cache().pop(&key);
+}

macros/src/lib.rs

@@ -9,6 +9,7 @@ pub fn memo(_attr: TokenStream, item: TokenStream) -> TokenStream {
     let vis = &func.vis;
     let sig = &func.sig;
     let block = &func.block;
+    let ident = &func.sig.ident;
 
     let output_ty = match &sig.output {
         ReturnType::Type(_, ty) => ty.clone(),
@@ -28,31 +29,53 @@ pub fn memo(_attr: TokenStream, item: TokenStream) -> TokenStream {
         .into();
     }
 
-    let _ident = format_ident!("_{}", sig.ident);
-    let mut _sig = sig.clone();
-    _sig.ident = _ident.clone();
-    _sig.inputs
+    let op_ident = format_ident!("{}_op", ident);
+    let mut op_sig = sig.clone();
+    op_sig.ident = op_ident.clone();
+    op_sig
+        .inputs
         .insert(0, parse_quote! { op: cache::MemoOperator });
+    op_sig.output = parse_quote! { -> () };
 
     let expanded = quote! {
         #vis #sig
         where #output_ty: Clone + 'static
         {
-            #_ident(cache::MemoOperator::Memo)
-        }
+            #op_ident(cache::MemoOperator::Memo(cache::Trace::Push));
 
-        #vis #_sig
-        where #output_ty: Clone + 'static
-        {
-            let key = #_ident as usize;
+            let key: cache::OperatorFunc = #op_ident;
             let rc = if let Some(rc) = cache::touch(key) {
                 rc
             } else {
                 let result: #output_ty = (|| #block)();
                 cache::store_in_cache(key, result)
             };
+
+            #op_ident(cache::MemoOperator::Memo(cache::Trace::Pop));
+
             (*rc).clone()
         }
+
+        #vis #op_sig
+        {
+            static mut dependents: Vec<cache::OperatorFunc> = Vec::new();
+            match op {
+                cache::MemoOperator::Memo(cache::Trace::Push) => {
+                    if let Some(last) = cache::call_stack().peek_mut() {
+                        unsafe { dependents.push(*last) };
+                    }
+                    cache::call_stack().push(#op_ident);
+                },
+                cache::MemoOperator::Memo(cache::Trace::Pop) => {
+                    cache::call_stack().pop();
+                },
+                cache::MemoOperator::Pop => {
+                    for dependent in unsafe { dependents.iter() } {
+                        cache::remove_from_cache(*dependent);
+                    }
+                },
+            }
+        }
     };
 
     expanded.into()

macros/tests/basic.rs

@@ -1,14 +1,22 @@
+#![feature(vec_peek_mut)]
+
 use macros::memo;
 
 #[memo]
 pub fn get_number() -> i32 {
-    println!("Calculate i32");
+    static mut INVOKED: bool = false;
+    assert!(!unsafe { INVOKED });
+    unsafe { INVOKED = true };
+
     42
 }
 
 #[memo]
 pub fn get_text() -> String {
-    println!("Calculate String");
+    static mut INVOKED: bool = false;
+    assert!(!unsafe { INVOKED });
+    unsafe { INVOKED = true };
+
     "hello".to_string()
 }