Memoization in Rust - demo of Cell and Rc

memo.rs

use std::cell::Cell;
use std::rc::Rc;

// Memoization - lazily compute some value and cache it once it's been computed.
//
// This uses a Cell containing an Option containing an Rc.
//
// A Cell is a structure that provides interior mutability - it's possible to mutate the cell's value without requiring
// the cell itself, or the struct that contains it (in this case the struct Memo) to be mutable. There are to types of
// cell: Cell and RefCell. A Cell owns the value that's inside it. If that value is non-copy, then the only way to
// access the value in the cell is by swapping it out with another value.
//
// The Option is used to keep track of whether the value has already been computed or not, and the computed value is
// wrapped in an Rc (shared, reference counted pointer) to make it possible to have multiple references to the value;
// the cell needs to contain a reference it, but we also want to return a reference to the value from the Memo::get()
// function.
//
// The Memo::get() function works as follows:
//
// First we need to get the Option inside the cell, so we (temporarily) swap it with None. Then we match on the Option:
// if it was Some, the computation was already done and we just get the Rc it contains; if it was None, we call the
// computation function and wrap the result in a new Rc. Before we return the Rc, we set the value of the cell to Some
// with a clone of the Rc (which increments the Rc's reference count).

pub struct Memo<T, F: Fn() -> T> {
    cell: Cell<Option<Rc<T>>>,
    func: F,
}

impl<T, F: Fn() -> T> Memo<T, F> {
    pub fn new(func: F) -> Memo<T, F> {
        Memo {
            cell: Cell::new(None),
            func,
        }
    }

    pub fn get(&self) -> Rc<T> {
        let rc = match self.cell.replace(None) {
            Some(rc) => rc,
            None => Rc::new((self.func)()),
        };

        self.cell.set(Some(rc.clone()));

        rc
    }
}

fn main() {
    let a = 37846;
    let b = 19338;

    let memo = Memo::new(|| {
        println!("Computing...");
        a * b
    });

    // Computation is only executed once, at the first get() call

    println!("Ready");
    println!("{}", memo.get());
    println!("{}", memo.get());
}