State Macro

A rust macro for decluttering mutable state computations. Suppose you want to write this:

// some type like this
type State<S> = Rc<RefCell<S>>

fn dense(state: State<S>, x: Var) -> Var {
    // this expression is messy!
    let x = mat_mul(state, param(state), x) + param(state);
    let y = param(state);

This library lets you instead write

#[stateful(State<S>)]
fn dense(x: Var) -> (Var, Var) {
    let x = ::mat_mul(::param(), x) + ::param();
    let y = ::param();
    (x, y)
}

The idea is to make functions stateful by abusing Rust's syntax for the root namespace. This means you cannot directly refer to the root namespace inside a #stateful annotated function. You can use the with_state! macro to work around this problem if necessary.

Writing ::f is analogous to !f with Haskell's Monadic Bang notation.

How it works

#[stateful(T)] transforms the function by:

1. Adding a state parameter of type T as the function's first argument
2. Wrapping the whole function body in with_state!.

So the example above becomes:

fn dense(state: State<S>, x: Var) -> (Var, Var) {
    with_state! { state;
        let x = ::mat_mul(::param(), x) + ::param();
        let y = ::param();
        (x, y)
    }
}

Of course, you can also write stateful computations inline using the with_state! macro directly, which works by:

1. Parsing an expression (the state - here a variable "state") terminated by a ; token
2. Replacing all function calls for a name starting with :: with ones taking the expression from (1)

NOTE: the contents of with_state! after the first ; can be arbitrary lines of rust code; essentially the 'innards' of a block. But when inlining the macro, these will not be wrapped in a block. So the code above inlines as follows:

fn dense(state: State<S>, x: Var) -> (Var, Var) {
    // hijack :: prefix to mean "inject a state param".
    let x = mat_mul(state, param(state), x) + param(state);
    let y = param(state);
    (x, y)
}