194: Coroutines and Generators via Effects

Difficulty: ⭐⭐⭐⭐⭐ Level: Expert Implement generators that lazily yield a sequence of values — like Python's `yield` — using Rust's stable Iterator trait, state-machine enums, and callback-based collection, all simulating OCaml 5's native effect-based approach.

The Problem This Solves

Generating sequences lazily is a common need: range values, Fibonacci numbers, lines from a file, events from a stream. The naive approach is to collect everything into a `Vec` first, but this breaks for infinite sequences and wastes memory for large ones. You want to produce values one at a time, pausing between each — that's what Python's `yield` keyword does. Python generators work by suspending execution: `yield x` pauses the function at that point, hands `x` to the caller, and resumes from the same point on the next `next()` call. This is a coroutine — a function that can pause and resume. Without language support, you'd have to manually encode this pause/resume logic as a state machine. Rust's nightly generators do this automatically; in stable Rust we do it ourselves. OCaml 5 implements this with algebraic effects: a `Yield_val` effect pauses the generator and hands the value to the handler, which can resume the generator for the next value. Rust stable offers multiple equivalent patterns: the `Iterator` trait captures the exact same semantics for simple cases, while state-machine enums let you encode more complex coroutines explicitly.

The Intuition

A coroutine is like a lazy letter-writer. Instead of writing the whole letter and sending it at once, they write one sentence, hand it to you, wait for you to read it, then write the next sentence. You control the pacing — you ask for the next sentence whenever you're ready. The letter-writer keeps their place between sentences. In Rust, `Iterator` is exactly this protocol: `next()` asks for one more value, the iterator remembers where it left off, and it returns `None` when done. A state-machine coroutine encodes "where the function is paused" as an enum variant. A callback-based generator (`fn generator<F>(f: F) -> Vec<i64>`) collects all yields eagerly — sacrificing laziness but making the API simple.

How It Works in Rust

// Approach 1: Iterator — the most idiomatic stable-Rust generator
struct RangeGenerator { current: i64, high: i64 }

impl Iterator for RangeGenerator {
 type Item = i64;
 fn next(&mut self) -> Option<i64> {
     if self.current > self.high { return None; }
     let v = self.current;
     self.current += 1;   // advance state — this is the "resume"
     Some(v)
 }
}

// Usage:
let sum: i64 = RangeGenerator::new(1, 100).sum();  // lazy, no Vec allocation
let squares: Vec<i64> = RangeGenerator::new(1, 5).map(|x| x * x).collect();

// Approach 2: State-machine coroutine — for multi-step logic between yields
#[derive(Debug)]
enum CoroutineState { Start, AfterFirst, AfterSecond, Done }

struct TwoStepCoroutine { state: CoroutineState, name: &'static str }

impl TwoStepCoroutine {
 fn resume(&mut self) -> Option<String> {
     match self.state {
         CoroutineState::Start => {
             self.state = CoroutineState::AfterFirst;
             Some(format!("{}: step 1", self.name))  // yield
         }
         CoroutineState::AfterFirst => {
             self.state = CoroutineState::AfterSecond;
             Some(format!("{}: step 2", self.name))  // yield
         }
         CoroutineState::AfterSecond => {
             self.state = CoroutineState::Done;
             Some(format!("{}: done", self.name))   // final yield
         }
         CoroutineState::Done => None,              // finished
     }
 }
}

// Interleave two coroutines — cooperative scheduling:
let mut task_a = TwoStepCoroutine::new("A");
let mut task_b = TwoStepCoroutine::new("B");
loop {
 let a = task_a.resume();
 let b = task_b.resume();
 if a.is_none() && b.is_none() { break; }
 // prints: A step1, B step1, A step2, B step2, A done, B done
}

// Approach 3: Callback-based (closest to OCaml's effect model)
fn generator<F: Fn(&mut dyn FnMut(i64))>(f: F) -> Vec<i64> {
 let mut values = Vec::new();
 f(&mut |v| values.push(v));  // yield_val closure = the Yield_val effect
 values
}

fn fibonacci_gen(count: usize, yield_val: &mut dyn FnMut(i64)) {
 let (mut a, mut b) = (0_i64, 1_i64);
 for _ in 0..count {
     yield_val(a);   // "perform (Yield_val a)" in OCaml
     let next = a + b; a = b; b = next;
 }
}

let fibs = generator(|y| fibonacci_gen(8, y));
// [0, 1, 1, 2, 3, 5, 8, 13]

In OCaml 5, `perform (Yield_val n)` suspends the generator and runs the handler's `continue k ()` to resume it. Rust's `Iterator::next()` is the same protocol — the caller drives resumption. The state-machine enum is how Rust's `async/await` and nightly generators are compiled internally.

What This Unlocks

Lazy sequences over large datasets — process lines from a file, events from a socket, or tree nodes in a traversal without loading everything into memory first.
Cooperative multitasking — interleave multiple coroutines by taking turns calling `resume()`; this is the foundation of async runtimes like Tokio before they had OS threads.
Infinite sequences — a coroutine that generates primes, Fibonacci numbers, or random values has no natural end; `Iterator` adapters (`take(N)`, `take_while(...)`) terminate the infinite stream lazily.

Key Differences

Concept	OCaml	Rust
Yielding a value	`perform (Yield_val n)` — suspends, gives `n` to handler	`yield_val(n)` callback (eager) or `Some(n)` from `Iterator::next` (lazy)
Resuming execution	`continue k ()` — native delimited continuation	`Iterator::next()` calls the `next()` method again; state machine advances
Suspend point encoding	Implicit — the effect captures the continuation	Explicit — enum variant records "where we are"
Infinite sequences	Natural — generator loops forever, handler pulls values	Iterator with `None` termination; infinite iterators via `.take()`
Nightly Rust generators	N/A	`#![feature(generators)]` gives Python-like `yield` syntax with compiler-generated state machines

// Coroutines and Generators — Stable Rust Simulation
//
// Rust's generators / coroutines are nightly-only.  This file shows two
// stable approaches:
//
//   1. State-machine enum: each "resume" advances the state.
//      Suitable when the yield points are known at compile time.
//
//   2. Closure-based generator: the generator holds a mutable iterator
//      internally and yields one value per call (Iterator adapter).
//
// The OCaml version uses effects (Yield_val) to implement a range generator.
// We replicate that with a stateful iterator.

// ── Approach 1: State-machine coroutine ───────────────────────────────────

#[derive(Debug)]
enum CoroutineState {
    Start,
    AfterFirst,
    AfterSecond,
    Done,
}

struct TwoStepCoroutine {
    state: CoroutineState,
    name: &'static str,
}

impl TwoStepCoroutine {
    fn new(name: &'static str) -> Self {
        TwoStepCoroutine { state: CoroutineState::Start, name }
    }

    /// Returns None when the coroutine is finished.
    fn resume(&mut self) -> Option<String> {
        match self.state {
            CoroutineState::Start => {
                self.state = CoroutineState::AfterFirst;
                Some(format!("{}: step 1", self.name))
            }
            CoroutineState::AfterFirst => {
                self.state = CoroutineState::AfterSecond;
                Some(format!("{}: step 2", self.name))
            }
            CoroutineState::AfterSecond => {
                self.state = CoroutineState::Done;
                Some(format!("{}: done", self.name))
            }
            CoroutineState::Done => None,
        }
    }
}

// ── Approach 2: Generator as an Iterator ─────────────────────────────────

struct RangeGenerator {
    current: i64,
    high: i64,
}

impl RangeGenerator {
    fn new(lo: i64, hi: i64) -> Self {
        RangeGenerator { current: lo, high: hi }
    }
}

impl Iterator for RangeGenerator {
    type Item = i64;

    fn next(&mut self) -> Option<i64> {
        if self.current > self.high {
            None
        } else {
            let v = self.current;
            self.current += 1;
            Some(v)
        }
    }
}

// ── Approach 3: Generator via Vec — collect all yields ────────────────────
//
// Models the OCaml `generator f` function that collects all yielded values.

fn generator<F: Fn(&mut dyn FnMut(i64))>(f: F) -> Vec<i64> {
    let mut values = Vec::new();
    f(&mut |v| values.push(v));
    values
}

fn range_gen(lo: i64, hi: i64, yield_val: &mut dyn FnMut(i64)) {
    let mut i = lo;
    while i <= hi {
        yield_val(i);
        i += 1;
    }
}

fn fibonacci_gen(count: usize, yield_val: &mut dyn FnMut(i64)) {
    let (mut a, mut b) = (0_i64, 1_i64);
    for _ in 0..count {
        yield_val(a);
        let next = a + b;
        a = b;
        b = next;
    }
}

// ── Approach 4: Resumable generator using a closure + channel ────────────
//
// For more complex coroutines that interleave with other code.

struct ClosureGenerator<T> {
    items: Box<dyn Iterator<Item = T>>,
}

impl<T> ClosureGenerator<T> {
    fn from_iter<I: Iterator<Item = T> + 'static>(iter: I) -> Self {
        ClosureGenerator { items: Box::new(iter) }
    }

    fn yield_next(&mut self) -> Option<T> {
        self.items.next()
    }
}

// ── Cooperative scheduler using generator protocol ────────────────────────

fn interleave_generators<A: Clone, B: Clone>(
    mut gen_a: impl Iterator<Item = A>,
    mut gen_b: impl Iterator<Item = B>,
) -> Vec<(Option<A>, Option<B>)> {
    let mut results = Vec::new();
    loop {
        let a = gen_a.next();
        let b = gen_b.next();
        if a.is_none() && b.is_none() { break; }
        results.push((a, b));
    }
    results
}

fn main() {
    // ── State-machine coroutine ────────────────────────────────────────
    println!("=== State-machine coroutine ===");
    let mut coro_a = TwoStepCoroutine::new("Task A");
    let mut coro_b = TwoStepCoroutine::new("Task B");
    loop {
        let a = coro_a.resume();
        let b = coro_b.resume();
        if a.is_none() && b.is_none() { break; }
        if let Some(msg) = a { println!("{}", msg); }
        if let Some(msg) = b { println!("{}", msg); }
    }

    println!();

    // ── Range generator via Iterator ──────────────────────────────────
    println!("=== Range generator (Iterator) ===");
    let vals: Vec<i64> = RangeGenerator::new(1, 5).collect();
    println!("range 1..5: {:?}", vals);

    let mapped: Vec<i64> = RangeGenerator::new(1, 5).map(|x| x * x).collect();
    println!("squares: {:?}", mapped);

    println!();

    // ── OCaml-style generator (collect all yields) ────────────────────
    println!("=== OCaml-style generator ===");
    let range_vals = generator(|y| range_gen(1, 5, y));
    println!("range 1..5: [{}]", range_vals.iter().map(|n| n.to_string()).collect::<Vec<_>>().join(";"));

    let fib_vals = generator(|y| fibonacci_gen(8, y));
    println!("fib(8):    [{}]", fib_vals.iter().map(|n| n.to_string()).collect::<Vec<_>>().join(";"));

    println!();

    // ── Closure generator ─────────────────────────────────────────────
    println!("=== Closure generator ===");
    let mut gen = ClosureGenerator::from_iter(0..5);
    print!("yields: ");
    while let Some(v) = gen.yield_next() {
        print!("{} ", v);
    }
    println!();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_range_generator_collects_all() {
        let vals = generator(|y| range_gen(1, 5, y));
        assert_eq!(vals, vec![1, 2, 3, 4, 5]);
    }

    #[test]
    fn test_range_generator_empty() {
        let vals = generator(|y| range_gen(5, 4, y));
        assert!(vals.is_empty());
    }

    #[test]
    fn test_fibonacci_first_eight() {
        let vals = generator(|y| fibonacci_gen(8, y));
        assert_eq!(vals, vec![0, 1, 1, 2, 3, 5, 8, 13]);
    }

    #[test]
    fn test_state_machine_coroutine_steps() {
        let mut coro = TwoStepCoroutine::new("T");
        assert!(coro.resume().unwrap().contains("step 1"));
        assert!(coro.resume().unwrap().contains("step 2"));
        assert!(coro.resume().unwrap().contains("done"));
        assert!(coro.resume().is_none());
    }

    #[test]
    fn test_range_iterator() {
        let v: Vec<i64> = RangeGenerator::new(3, 6).collect();
        assert_eq!(v, vec![3, 4, 5, 6]);
    }

    #[test]
    fn test_closure_generator_yields_in_order() {
        let mut gen = ClosureGenerator::from_iter(10..=13);
        assert_eq!(gen.yield_next(), Some(10));
        assert_eq!(gen.yield_next(), Some(11));
        assert_eq!(gen.yield_next(), Some(12));
        assert_eq!(gen.yield_next(), Some(13));
        assert_eq!(gen.yield_next(), None);
    }
}

effect Yield_val : int -> unit

let generator f =
  let resume = ref (fun () -> ()) in
  let values = ref [] in
  resume := (fun () ->
    match f () with
    | () -> ()
    | effect (Yield_val n) k ->
      values := n :: !values;
      resume := (fun () -> continue k ())
  );
  !resume ();
  List.rev !values

(* Range generator *)
let range lo hi () =
  let i = ref lo in
  while !i <= hi do
    perform (Yield_val !i);
    incr i
  done

let () =
  let vals = generator (range 1 5) in
  Printf.printf "range 1..5: [%s]\n"
    (vals |> List.map string_of_int |> String.concat ";")