Code shared from the Rust Playground

playground.rs

use std::sync::Arc;
use std::sync::RwLock;
use std::hash::Hash;
use std::collections::HashMap;

fn main () {
    let m: Memoizer<i64,i64> = Memoizer::default();
    let yeet = m.memoize(5, simple_fn);
    assert_eq!(*yeet, 100);
    let yeet = m.memoize(5, simple_fn);
    assert_eq!(*yeet, 100);
    let yeet = m.memoize(5, simple_fn);
    assert_eq!(*yeet, 100);
    
    let m: Memoizer<&str, String> = Memoizer::default();
    let yeet = m.memoize("s", fun);
    assert_eq!(*yeet, String::from("s"));
    let yeet = m.memoize("s", fun);
    assert_eq!(*yeet, String::from("s"));
}

pub enum StorageOption<P> {
    DontStoreNew,
    StoreNoReplace,
    RemoveRefAndStore(P)
}

#[derive(Default)]
pub struct Store<P: Hash + Eq, O>(Arc<RwLock<HashMap< P, Arc<O> >>>);

impl <P,O> Store<P,O> 
where 
    P: Hash + Clone + Eq,
    O: Clone
{
    fn get(&self, p: &P) -> Option<Arc<O>> {
        let map = self.0.read().unwrap();
        return map.get(p).map(|o|o.clone())
    } 
    fn insert(&self, p: P, o: Arc<O>) {
        let mut map = self.0.write().unwrap();
        map.insert(p,o.clone());
    }
    
    fn remove(&self, p: &P) {
        let mut map = self.0.write().unwrap();
        map.remove(p);
    }
}

#[derive(Default)]
pub struct Memoizer<P,O>
where
    P: Hash + Eq
{
    store: Store<P, O>
}


impl <'a, P,O > Memoizer<P,O>
where 
    P: Hash + Clone + Eq,
    O: Clone
{
    fn cache_output(&self, p: P, cached_fun: fn(P) -> O) -> Arc<O> {
        let o = (cached_fun)(p.clone());
        let o = Arc::new(o);
        self.store.insert(p, o.clone());
        o
    }
    
    fn cache_output_with_ejection<CF,EF>(&self, p: P, cached_fun: CF, ejection_function: EF) -> Arc<O>
    where 
        CF: Fn(P) -> O,
        EF: Fn(&Store<P,O>, &O) -> StorageOption<P>
    {
        let o = (cached_fun)(p.clone());
        
        match (ejection_function)(&self.store, &o) {
            StorageOption::DontStoreNew => Arc::new(o),
            StorageOption::StoreNoReplace => {
                let o = Arc::new(o);
                self.store.insert(p, o.clone());
                o
            },
            StorageOption::RemoveRefAndStore(ref remove_p) => {
                self.store.remove(remove_p);
                let o = Arc::new(o);
                self.store.insert(p, o.clone());
                o
            }
        }
    }

    pub fn memoize(&self, p: P, cached_fun: fn(P) -> O) -> Arc<O> {
        self.store
            .get(&p)
            .unwrap_or_else(|| self.cache_output(p, cached_fun))
    }
    
    pub fn memoize_with_ejection<CF,EF>(&self, p: P, cached_fun: CF, ejection_fun: EF) -> Arc<O>
    where 
        CF: Fn(P) -> O,
        EF: Fn(&Store<P,O>, &O) -> StorageOption<P> // A function for assessing if the store is at capacity and if the new output should supplant an existing value
    {
        self.store
            .get(&p)
            .unwrap_or_else(|| self.cache_output_with_ejection(p, cached_fun, ejection_fun))
    }
}

fn simple_fn(p: i64) -> i64 {
    return p * 20
}

fn fun(p: &str) -> String {
    return p.to_string()
}