Rust Book Ch 13 Closures with Generic Parameters and HashMap for memoization

Question

I have been reading through the rust book. I have made it to chapter 13, at the end of the section called (Storing Closures Using Generic Parameters and the Fn Traits) a couple of improvements are described and left as kind of an exercise.

The task at hand is to change the cache implementation to support not only u32 but any data type and also allow the results to be stored in a hashmap thus allowing us to memoize multiple queries.

I went the route of the everything is copied (since I tried to avoid lifetimes for my first implementation) and I just wanted the feedback from more experienced people on how they would tackle and what pitfall does this version I wrote have.

use std::collections::HashMap;
use std::cmp::Eq;
use std::hash::Hash;

struct Cacher<I, O, T>
where 
    T: Fn(I) -> O,
{
    calc: T,
    cache: HashMap<I, O>
}


impl<I: Eq + Hash + Copy, O, T> Cacher<I, O, T> 
where
    T: Fn(I) -> O
{
    fn new(calc: T) -> Cacher<I, O, T> {
        Cacher {
            calc,
            cache: HashMap::new()
        }
    }

    fn get(&mut self, n: &I) -> &O {
        self.cache.entry(*n).or_insert((self.calc)(*n))
    }
}

fn main() {
    let mut workout = Cacher::new(|x| x + 1);
    println!("{}", workout.get(&4));
    println!("{}", workout.get(&5));
    println!("{}", workout.get(&3));
    println!("{}", workout.get(&4));
}

Another question is where trait are necessary in this case, I realized the compiler doesn't complain if I only put them in the implementation block (but wondered why doesn't it think they are required at the level of the struct since Generic type I has to be both Hashable and Comparable when creating the hashmap right ?)

Answer 1

Even though I'm relatively new to Rust myself, I'll give this a shot since no one else has responded so far.

Is Cacher caching?

Let's test that Cacher is working as intended (wouldn't call calc on the same value twice) by adding a log:

// --snip--
let mut workout = Cacher::new(|x| {
    println!("calc {}", x);
    x + 1
});
// --snip--

The output is:

calc 4
5
calc 5
6
calc 3
4
calc 4
5

Why is calc called twice on 4? Let's check the implementation of Cacher::get:

fn get(&mut self, n: &I) -> &O {
    self.cache.entry(*n).or_insert((self.calc)(*n))
}

The problem is that the expression (self.calc)(*n) is evaluated and the result value given to or_insert whenever Cacher::get is called, even though the result is discarded if the key *n already exists in the hashmap. This can be fixed by using or_insert_with:

fn get(&mut self, n: &I) -> &O {
    let calc = &self.calc;
    self.cache.entry(*n).or_insert_with(|| calc(*n))
}

Now the output is as expected:

calc 4
5
calc 5
6
calc 3
4
5

NB: The borrow checker wouldn't let us simply write || (self.calc)(*n). The reason is self.cache is borrowed mutably within Cacher::get, and the whole self (implying self.cache) would be borrowed within the closure. That would violate the rule that a mutable borrow is exclusive. We circumvent this by explicitly splitting the borrow outside of the closure. The closure then captures only calc.

Some links to read up on this:

• https://doc.rust-lang.org/nomicon/borrow-splitting.html
• https://stackoverflow.com/questions/61699010/rust-not-allowing-mutable-borrow-when-splitting-properly
• https://www.reddit.com/r/rust/comments/cufiq7/rust_mutably_borrows_the_whole_struct_even_though/

Trait Bounds

Let's get to your question on trait bounds. First off, it does make a real difference where we put them:

• When the trait bounds are on the struct, then the compiler will not let us instantiate the generic without satisfying the trait bounds, at all.

• When the trait bounds are on the impl block, then we are allowed to instantiate the generic with types that do not satisfy the trait bounds, but the methods defined in the impl block will be available only to instances where the trait bounds are satisfied.

In fact, if you have a look at the implementation of HashMap, you'll notice that HashMap has trait bounds only on impl blocks for particular sets of methods that require the respective traits. Thus we may create, but not insert into, a HashMap whose key type doesn't satisfy Eq + Hash:

// This is alright:
struct X {}
let mut map = HashMap::<X, u32>::new();

// This won't compile:
map.insert(X {}, 1);

This being a factual difference, what should we do here? A question on StackOverflow has some interesting answers discussing this question: https://stackoverflow.com/questions/49229332/should-trait-bounds-be-duplicated-in-struct-and-impl. In this case, I'd probably stick with HashMap's approach and be as specific as possible (see code below).

A Note on Tooling

Rust's ecosystem brings some nice tools. Clippy is a linter that does, in fact, hint at the aforementioned bug:

$ cargo clippy
warning: use of `or_insert` followed by a function call
  --> src/main.rs:26:30
   |
26 |         self.cache.entry(*n).or_insert((self.calc)(*n))
   |                              ^^^^^^^^^^^^^^^^^^^^^^^^^^ help: try this: `or_insert_with(|| (self.calc)(*n))`
   |
   = note: `#[warn(clippy::or_fun_call)]` on by default
   = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#or_fun_call

Also, there's cargo fmt to enforce a standard code style. In this case, it sorts the imports, adds a trailing comma in the struct and removes a duplicated blank line.

Final Code

Here's the code I ended up with, including some further remarks:

use std::cmp::Eq;
use std::collections::HashMap;
use std::hash::Hash;

// We may even leave out the trait bound for `T`
// (Example: Perhaps at some point you want to add another implementation
//           where calculation isn't done via a closure/function,
//           but via some special object. This is general enough to allow that.)
struct Cacher<I, O, T> {
    calc: T,
    cache: HashMap<I, O>,
}

// Here I'm splitting into two `impl` blocks.
// A `Cacher` may be created without knowing anything about the types.
impl<I, O, T> Cacher<I, O, T> {
    fn new(calc: T) -> Cacher<I, O, T> {
        Cacher {
            calc,
            cache: HashMap::new(),
        }
    }
}

// Now that the hashmap is used, and `calc` is used,
// the trait bounds are indeed required
impl<I, O, T> Cacher<I, O, T>
where
    // For clarity and consistency, I put all trait bounds into the where block
    I: Eq + Hash + Copy,
    T: Fn(I) -> O,
{
    fn get(&mut self, n: &I) -> &O {
        let calc = &self.calc;
        self.cache.entry(*n).or_insert_with(|| calc(*n))
    }
}

fn main() {
    let mut workout = Cacher::new(|x| x + 1);
    println!("{}", workout.get(&4));
    println!("{}", workout.get(&5));
    println!("{}", workout.get(&3));
    println!("{}", workout.get(&4));
}