How to store state within a module between function calls?

Viewed 327

I have the following functions in a module:

pub fn square(s: u32) -> u64 {
    if s < 1 || s > 64 {
        panic!("Square must be between 1 and 64")
    }

    total_for_square(s) - total_for_square(s - 1)
}

fn total_for_square(s: u32) -> u64 {
    if s == 64 {
        return u64::max_value();
    }

    2u64.pow(s) - 1
}

pub fn total() -> u64 {
    u64::max_value()
}

This works fine when calling individual functions directly. However, I want to optimize it and cache values to total_for_square to speed up future look ups (storing in a HashMap). How should I approach where to store the HashMap so it's available between calls? I know I could refactor to put all of this in a struct, but in this case, I cannot change the API.

In other, higher level languages I have used, I would just have a variable in the same scope as the functions. However, it's not clear if that is possible in Rust on the module level.

2 Answers

In other, higher level languages I have used, I would just have a variable in the same scope as the functions.

You can use something similar in Rust but it's syntactically more complicated: you need to create a global for your cache using lazy_static or once_cell for instance.

The cache will need to be thread-safe though, so either a regular map sitting behind a Mutex or RwLock, or some sort of concurrent map.

Although given you only have 64 inputs, you could just precompute the entire thing and return precomputed values directly.

The cached crate comes in handy:

use cached::proc_macro::cached;

#[cached]
fn total_for_square(s: u32) -> u64 {
    if s == 64 {
        return u64::MAX;
    }

    2u64.pow(s) - 1
}

Indeed, you only need to write two lines, and the crate will take care of everything. Internally, the cached values are stored in a hash map.

(Note that u64::max_value() has been superseded by u64::MAX)


Side note: in this specific case, the simplest solution is probably to modify square so that it returns s * s.

Related