How do I properly implement a caching struct in Rust for lazily-computed values?

Question

In Rust, I'm searching for a more concise and idiomatic way to define a cache for expensive calculations.

In Chapter 13 from the Second Edition of The Rust Programming Language, the authors leave an exercise for the reader to refactor a struct to return lazily-computed values.

After four days of trial and error, I came up with:

use std::collections::HashMap;

#[allow(dead_code)]
struct Cache<T> 
    where T: Fn(u32) -> u32
{
    calculation: T,
    internal: HashMap<u32, u32>
}

#[allow(dead_code)]
impl<T> Cache<T>
    where T: Fn(u32) -> u32
{
    fn new(calculation: T) -> Cache<T> {
        Cache {
            calculation,
            internal: HashMap::new(),
        }
    }

    fn set(&mut self, arg: u32, value: u32) -> u32 {
        self.internal.insert(arg, value);
        self.get(arg)
    }

    fn get(&mut self, arg: u32) -> u32 {
        self.internal[&arg]
    }

    fn value(&mut self, arg: u32) -> u32 {
        match self.internal.contains_key(&arg) {
            true => {
                self.get(arg)
            },
            false => {
                self.set(arg, (self.calculation)(arg))
            },
        }
    }
}

For the test function:

#[test]
fn call_with_different_values() {
    let mut c = Cache::new(|a| a);
    let _v1 = c.value(1);
    let v2 = c.value(2);
    assert_eq!(v2, 2);
}

Can I improve this impl? Is there a more acceptable way for this to be written?

Answer 1

• In set(), self.get(arg) can be simplified to value.

• get() only needs a non-mutable &self.

• get() and set() don’t really seem that helpful in the first place.

• The hash_map::Entry API can help you simplify value().

  fn value(&mut self, arg: u32) -> u32 {
      let Self { internal, calculation } = self;
      let entry = internal.entry(arg);
      *entry.or_insert_with(|| (calculation)(arg))
  }
  

Then the minor style things:

• You can use Self instead of repeating the struct name inside the impl.

• Trailing commas are good.

In all:

use std::collections::HashMap;

pub struct Cache<T>
    where T: Fn(u32) -> u32
{
    calculation: T,
    internal: HashMap<u32, u32>,
}

impl<T> Cache<T>
    where T: Fn(u32) -> u32
{
    pub fn new(calculation: T) -> Cache<T> {
        Self {
            calculation,
            internal: HashMap::new(),
        }
    }

    pub fn value(&mut self, arg: u32) -> u32 {
        let Self { internal, calculation } = self;
        let entry = internal.entry(arg);
        *entry.or_insert_with(|| (calculation)(arg))
    }
}

#[test]
fn call_with_different_values() {
    let mut c = Cache::new(|a| a);
    let _v1 = c.value(1);
    let v2 = c.value(2);
    assert_eq!(v2, 2);
}

playground