Encoding ML-style modules in Rust

fn str_len(s: &str) -> usize {
    println!("Called expensive function str_len");
    s.len()
}

// Declare that we return an `FnMut` rather than an `Fn`
// because the returned closure may modify mutable state
// (and thus isn't safe to call in parallel -- not relevant
// in this example).
fn make_cached_str_len() -> impl (FnMut(&str) -> usize) {
    let mut cache = HashMap::new();
    // Use the `move` keyword to move the mutable `cache`
    // into the returned closure object. Unlike a variable defined
    // in the closure body itself, this cache will be shared across
    // *all* invocations of the new function.
    move |s: &str| -> usize {
        // For simplicity, we clone the `&str` as the cache key
        // rather than worry about the lifetime annotations here.
        // If we didn't do this, it's possible that our internal `cache`
        // would hold a reference to a `str` even after that `str` was freed.
        let key: String = s.to_owned();
        // Look up the entry in our cache,
        // or create a new one if none exists.
        let result: usize = *cache.entry(key).or_insert_with_key(|key| str_len(key));
        result
    }
}

fn main() {
    let mut cached_str_len = make_cached_str_len();
    let key = "foo";

    // Prints:
    // Called expensive function str_len
    // Cached str len 1/2: 3
    println!("Cached str len 1/2: {}", cached_str_len(key));

    // Cached str len 2/2: 3
    println!("Cached str len 2/2: {}", cached_str_len(key));
}

// Create a function which is generic over the key and value types,
// and accepts a "base" function which takes a key and returns a value.
//
// In order to store the key in the cache, we had to add the additional
// bounds `Eq + Hash` to `K`. To simplify the lifetimes, we added
// the lifetime bound `'static` to `K` and the bound `Clone` to `V`
// (and cloned the value before returning it).
fn cache_f<K: Eq + Hash + 'static, V: Clone>(f: impl Fn(&K) -> V) -> impl (FnMut(K) -> V) {
    let mut cache = HashMap::new();
    move |key: K| -> V {
        let result: &V = cache.entry(key).or_insert_with_key(|key| f(key));
        result.clone()
    }
}

fn main() {
    // To make some types work, add an additional `&`
    // to the parameter type.
    let mut cached_str_len = cache_f(|s: &&str| str_len(s));
    let key = "foo";

    // Prints:
    // Called expensive function str_len
    // Cached str len 1/2: 3
    println!("Cached str len 1/2: {}", cached_str_len(key));

    // Cached str len 2/2: 3
    println!("Cached str len 2/2: {}", cached_str_len(key));
}

// `impl T` does not compile (yet).
type CachedF<K: Eq + Hash + 'static, V: Clone> = impl FnMut(K) -> V;
type CachedStrLenF = CachedF<&'static str, usize>;

fn accepts_cached_str_len(cached_str_len: CachedStrLenF) {
    let key = "foo";
    println!("Cached str len 1/2: {}", cached_str_len(key));
    println!("Cached str len 2/2: {}", cached_str_len(key));
}

fn main() {
    let mut cached_str_len = cache_f(|s: &&str| str_len(s));
    accepts_cached_str_len(cached_str_len);
}

fn accepts_cached_str_len(mut cached_str_len: impl FnMut(&'static str) -> usize) {
    let key = "foo";
    println!("Cached str len 1/2: {}", cached_str_len(key));
    println!("Cached str len 2/2: {}", cached_str_len(key));
}

fn main() {
    let cached_str_len = cache_f(|s: &&str| str_len(s));
    accepts_cached_str_len(cached_str_len);
}

trait Backend {
    // We declare "associated types" for this trait. For any type 
    // which implements `Backend`, it must declare corresponding `Key`
    // and `Value` types, and include an implementation of `compute`
    // adhering to those types.
    type Key;
    type Value;
    fn compute(&self, key: &Self::Key) -> Self::Value;
}

// This type doesn't need state, but you could imagine
// embedding e.g. a database connection in this struct.
//
// Normally, types without any data aren't very useful, but
// this is one of the key ideas of the "encoding of ML-style
// modules". Instead of state (or in addition to state), the
// type instead has associated types and functions,
// declared below.
//
// This type corresponds to the "closed-over" variables
// of a closure object.
struct StrLenBackend;

impl Backend for StrLenBackend {
    type Key = &'static str;
    type Value = usize;
    fn compute(&self, key: &Self::Key) -> Self::Value {
        println!("Computing length of key {key:?}");
        key.len()
    }
}

struct CachedBackend<B: Backend> {
    backend: B,
    cache: HashMap<B::Key, B::Value>,
}

impl<B: Backend> CachedBackend<B>
where
    // We require that the wrapped backend adheres to these bounds
    // in order to call any of the functions in this `impl`
    // (including `new`, which is how you would get a new value
    // of this type).
    //
    // Previously, we had to add the appropriate bounds to 
    // each generic type in each place that it appeared.
    // Using a `trait`, we can instead declare the bounds
    // just once on the `impl`.
    B::Key: Eq + Hash,
{
    fn new(backend: B) -> Self {
        Self {
            backend,
            cache: Default::default(),
        }
    }

    fn compute_cached(&mut self, key: B::Key) -> &B::Value {
        self.cache
            .entry(key)
            .or_insert_with_key(|key| self.backend.compute(key.clone()))
    }
}

type CachedStrLenBackend = CachedBackend<StrLenBackend>;

fn accepts_cached_str_len(mut cached_str_len: CachedStrLenBackend) {
    let key = "foo";

    // Prints:
    // Computing length of key "foo"
    // Cached len 1/2: 3
    println!("Cached len 1/2: {}", cached_str_len.compute_cached(key));

    // Cached len 2/2: 3
    println!("Cached len 2/2: {}", cached_str_len.compute_cached(key));
}

fn main() {
    let backend = StrLenBackend;
    let cached_str_len = CachedBackend::new(backend);
    accepts_cached_str_len(cached_str_len);
}

trait Backend {
    type Key;
    type Value;

    // We no longer take a value of type `&self`.
    fn compute(key: &Self::Key) -> Self::Value;
}

// In Rust, it's not possible to make a value of the type
// of an empty enum.
enum StrLenBackend {}

impl Backend for StrLenBackend {
    type Key = &'static str;
    type Value = usize;
    fn compute(key: &Self::Key) -> Self::Value {
        println!("Computing length of key {key:?}");
        key.len()
    }
}

// Note that we no longer have a member of type `B`.
struct CachedBackend<B: Backend> {
    cache: HashMap<B::Key, B::Value>,
}

impl<B: Backend> CachedBackend<B>
where
    B::Key: Eq + Hash,
{
    fn new() -> Self {
        Self {
            cache: Default::default(),
        }
    }

    fn compute_cached(&mut self, key: B::Key) -> &B::Value {
        self.cache
            .entry(key)
            .or_insert_with_key(|key| B::compute(key.clone()))
    }
}

type CachedStrLenBackend = CachedBackend<StrLenBackend>;

fn accepts_cached_str_len(mut cached_str_len: CachedStrLenBackend) {
    let key = "foo";

    // Prints:
    // Computing length of key "foo"
    // Cached len 1/2: 3
    println!("Cached len 1/2: {}", cached_str_len.compute_cached(key));

    // Cached len 2/2: 3
    println!("Cached len 2/2: {}", cached_str_len.compute_cached(key));
}

fn main() {
    // This weird syntax is equivalent to `CachedBackend<StrLenBackend>::new()`
    // in some other languages. It's written with an extra `::` due to
    // technical restrictions on the syntax.
    //
    // Note that, unlike before, we never pass a *value* into
    // the `CachedBackend`, only a type.
    let cached_str_len = CachedBackend::<StrLenBackend>::new();

    // We could also write the type annotation on the left side and let
    // type inference figure it out:
    let _also_works: CachedBackend<StrLenBackend> = CachedBackend::new();

    // Or we could use our previously-defined type alias:
    let _also_works = CachedStrLenBackend::new();

    accepts_cached_str_len(cached_str_len);
}