060: State Monad

The Problem This Solves

// Without State monad — state threads manually through every call:
fn count_three(state: i32) -> ((i32, i32, i32), i32) {
 let a = state;       let state = state + 1;
 let b = state;       let state = state + 1;
 let c = state;       let state = state + 1;
 ((a, b, c), state)
}
// Works! But imagine 15 steps where state is used 5 times each...
// Every intermediate variable is named `state`, rebinding it each time.
// The business logic drowns in plumbing.

The Intuition

state₀ → [Machine A] → (resultA, state₁)
state₁ → [Machine B] → (resultB, state₂)
state₂ → [Machine C] → (resultC, state₃)

• State monad — a wrapper around functions of type `S -> (A, S)` with a `bind` that threads state automatically
• `get` — returns the current state as the result (state is unchanged)
• `put(s)` — replaces the current state with `s` (result is `()`)
• `modify(f)` — applies `f` to the current state (like `get` + `put`)
• `run(initial_state)` — actually execute the whole chain starting with `initial_state`

How It Works in Rust

// The core type: a wrapper around a state-threading function
struct State<S, A> {
 run: Box<dyn FnOnce(S) -> (A, S)>,
}

impl<S: 'static, A: 'static> State<S, A> {
 fn new(f: impl FnOnce(S) -> (A, S) + 'static) -> Self {
     State { run: Box::new(f) }
 }

 // Execute the whole chain with an initial state
 fn run(self, s: S) -> (A, S) {
     (self.run)(s)
 }

 // Wrap a plain value — state passes through unchanged
 fn pure(a: A) -> Self {
     State::new(move |s| (a, s))
 }

 // Chain: run self, pass result + new state to f, run f
 fn and_then<B: 'static>(self, f: impl FnOnce(A) -> State<S, B> + 'static) -> State<S, B> {
     State::new(move |s| {
         let (a, s2) = self.run(s);   // run first step, get result + updated state
         f(a).run(s2)                  // feed result + state into next step
     })
 }
}

// Primitives for working with state:

fn get<S: Clone + 'static>() -> State<S, S> {
 // Copy the state into the result — state unchanged
 State::new(|s: S| (s.clone(), s))
}

fn put<S: 'static>(new_s: S) -> State<S, ()> {
 // Replace state with new_s — result is ()
 State::new(move |_| ((), new_s))
}

fn modify<S: 'static>(f: impl FnOnce(S) -> S + 'static) -> State<S, ()> {
 // Apply f to state — result is ()
 State::new(move |s| ((), f(s)))
}

// A counter step: read current value, increment state, return old value
fn tick() -> State<i32, i32> {
 get::<i32>().and_then(|n| {
     put(n + 1).map(move |()| n)  // increment state, return old n
 })
}

// Chain three ticks — state threads automatically
fn count3() -> State<i32, (i32, i32, i32)> {
 tick().and_then(|a|
     tick().and_then(move |b|
         tick().map(move |c| (a, b, c))))
}

// Run it:
let ((a, b, c), final_state) = count3().run(0);
// a=0, b=1, c=2, final_state=3

// Stack operations — state is Vec<i32>
fn push(x: i32) -> State<Vec<i32>, ()> {
 modify(move |mut stack: Vec<i32>| { stack.push(x); stack })
}

fn pop() -> State<Vec<i32>, Option<i32>> {
 State::new(|mut stack: Vec<i32>| {
     let val = stack.pop();
     (val, stack)
 })
}

// Chain push/pop operations:
let ops = push(1)
 .and_then(|()| push(2))
 .and_then(|()| push(3))
 .and_then(|()| pop());

let (top, remaining_stack) = ops.run(vec![]);
// top = Some(3), remaining_stack = [1, 2]

What This Unlocks

• Incremental ID generation: A stateful counter wrapped in State can be composed into any sequence of operations that need unique IDs without passing `&mut next_id` everywhere.
• Parser combinators: A parser is naturally `State<&str, ParseResult>` — each parser step consumes some input (mutates state) and returns a result. State monad lets you compose parsers without manually tracking the remaining input.
• Game scripting: A sequence of game actions (move, attack, pick up item) that all read/update game state can be composed as `State<GameState, ()>` operations, enabling AI behavior trees to be written as pure data.

Key Differences