086 Intermediate

086 — Custom Iterator with State

Functional Programming

Tutorial

The Problem

Implement three stateful custom iterators: a Counter that steps by a given value indefinitely, a Fib Fibonacci iterator, and a Collatz iterator that terminates at 1. Demonstrate that implementing next unlocks the full iterator adapter chain. Compare with OCaml's closure-based counters and Seq lazy sequences.

🎯 Learning Outcomes

• Hold mutable state in struct fields rather than closures for complex iterators

• Return Some(val) forever from an infinite iterator; use .take(n) to bound it

• Implement a finite iterator by tracking a done_ flag and returning None

• Use the Collatz sequence as a real-world example of a self-terminating iterator

• Map Rust's mutable struct iterator to OCaml's mutable ref closure and Seq thunk

• Compose multiple custom iterators using .zip, .map, .sum

Code Example

//! 086: Custom Iterator with State — stateful iterators built by implementing
//! [`Iterator::next`].
//!
//! The OCaml original expresses stateful sequences three different ways:
//!
//! * `make_counter` captures a `ref` cell in a closure and mutates it on each
//!   call.
//! * `fibonacci` returns a `Seq.t`, i.e. a thunk producing `Seq.Cons(a, …)`
//!   at every step.
//! * `collatz_seq` also uses `Seq.t`, but signals termination with `Seq.Nil`
//!   when the current value reaches 1.
//!
//! Rust unifies all three under a single abstraction: a `struct` holding the
//! state plus an `impl Iterator` with a `next` method. Infinite iterators
//! always return `Some(_)` (the caller bounds them with [`Iterator::take`]);
//! finite iterators return `None` to signal termination. Implementing `next`
//! is sufficient to inherit the full iterator adapter chain (`map`, `filter`,
//! `collect`, `zip`, …).

/// Arithmetic progression: yields `start + step`, `start + 2*step`, … forever.
///
/// The OCaml `make_counter start step` closure captures a mutable `ref`
/// initialised to `start - step` and pre-increments before returning the
/// value. This struct achieves the same effect by storing `current` and
/// advancing it on every call to [`Iterator::next`].
///
/// Because the sequence is infinite, `next` always returns `Some(_)`. Use
/// [`Iterator::take`] to bound it.
#[derive(Debug, Clone, Copy)]
pub struct Counter {
    current: i32,
    step: i32,
}

impl Counter {
    /// Build a counter whose first yielded value is `start + step`.
    pub fn new(start: i32, step: i32) -> Self {
        Self {
            current: start,
            step,
        }
    }
}

impl Iterator for Counter {
    type Item = i32;

    fn next(&mut self) -> Option<Self::Item> {
        self.current += self.step;
        Some(self.current)
    }
}

/// Infinite Fibonacci sequence: `0, 1, 1, 2, 3, 5, 8, …`.
///
/// The OCaml version is a `Seq.t` whose thunk returns
/// `Seq.Cons(a, aux b (a + b))`. The Rust struct keeps the same two-element
/// state `(a, b)` and shifts it on every call. Bound the iterator with
/// [`Iterator::take`].
#[derive(Debug, Clone, Copy)]
pub struct Fib {
    a: u64,
    b: u64,
}

impl Fib {
    /// Start the sequence at `(0, 1)`.
    pub fn new() -> Self {
        Self { a: 0, b: 1 }
    }
}

impl Default for Fib {
    fn default() -> Self {
        Self::new()
    }
}

impl Iterator for Fib {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        let value = self.a;
        let next = self.a + self.b;
        self.a = self.b;
        self.b = next;
        Some(value)
    }
}

/// Collatz (3n+1) sequence from a chosen start, terminating at 1.
///
/// Each step halves `n` when it is even, or replaces it with `3n + 1` when
/// odd. The sequence is finite: once 1 is yielded, subsequent calls to
/// `next` return [`None`]. Mirrors OCaml's `collatz_seq`, which wraps the
/// final element in `Seq.Cons(1, fun () -> Seq.Nil)`.
#[derive(Debug, Clone, Copy)]
pub struct Collatz {
    n: u64,
    finished: bool,
}

impl Collatz {
    /// Build a Collatz iterator starting at `start`. The first yielded value
    /// is `start` itself.
    pub fn new(start: u64) -> Self {
        Self {
            n: start,
            finished: false,
        }
    }
}

impl Iterator for Collatz {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            return None;
        }
        let value = self.n;
        if self.n == 1 {
            self.finished = true;
        } else if self.n.is_multiple_of(2) {
            self.n /= 2;
        } else {
            self.n = 3 * self.n + 1;
        }
        Some(value)
    }
}

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

    #[test]
    fn counter_matches_ocaml_make_counter() {
        // OCaml: let counter = make_counter 0 2 in counter () = Some 2, etc.
        let v: Vec<i32> = Counter::new(0, 2).take(3).collect();
        assert_eq!(v, vec![2, 4, 6]);
    }

    #[test]
    fn counter_supports_negative_step() {
        let v: Vec<i32> = Counter::new(10, -3).take(4).collect();
        assert_eq!(v, vec![7, 4, 1, -2]);
    }

    #[test]
    fn fibonacci_first_eight_terms() {
        // OCaml: take_seq 8 (fibonacci ()) = [0; 1; 1; 2; 3; 5; 8; 13]
        let fibs: Vec<u64> = Fib::new().take(8).collect();
        assert_eq!(fibs, vec![0, 1, 1, 2, 3, 5, 8, 13]);
    }

    #[test]
    fn collatz_from_six_terminates_at_one() {
        // OCaml: List.of_seq (collatz_seq 6) = [6; 3; 10; 5; 16; 8; 4; 2; 1]
        let v: Vec<u64> = Collatz::new(6).collect();
        assert_eq!(v, vec![6, 3, 10, 5, 16, 8, 4, 2, 1]);
    }

    #[test]
    fn collatz_from_one_is_single_element() {
        let v: Vec<u64> = Collatz::new(1).collect();
        assert_eq!(v, vec![1]);
    }

    #[test]
    fn custom_iterators_compose_with_adapters() {
        // Implementing `next` is enough to get the whole adapter chain.
        let sum_of_even_fibs: u64 = Fib::new().take(10).filter(|n| n.is_multiple_of(2)).sum();
        assert_eq!(sum_of_even_fibs, 2 + 8 + 34);

        let zipped: Vec<(i32, u64)> = Counter::new(0, 1).zip(Fib::new()).take(4).collect();
        assert_eq!(zipped, vec![(1, 0), (2, 1), (3, 1), (4, 2)]);
    }
}

(* 086: Custom Iterator with State *)

(* Approach 1: Counter using mutable ref *)
let make_counter start step =
  let n = ref (start - step) in
  fun () -> n := !n + step; Some !n

(* Approach 2: Fibonacci via Seq *)
let fibonacci () =
  let rec aux a b () = Seq.Cons (a, aux b (a + b)) in
  aux 0 1

let take_seq n s =
  let rec aux n s acc =
    if n <= 0 then List.rev acc
    else match s () with
      | Seq.Nil -> List.rev acc
      | Seq.Cons (x, rest) -> aux (n - 1) rest (x :: acc)
  in
  aux n s []

(* Approach 3: Collatz sequence iterator *)
let collatz_seq start =
  let rec aux n () =
    if n = 1 then Seq.Cons (1, fun () -> Seq.Nil)
    else Seq.Cons (n, aux (if n mod 2 = 0 then n / 2 else 3 * n + 1))
  in
  aux start

(* Tests *)
let () =
  let counter = make_counter 0 2 in
  assert (counter () = Some 2);
  assert (counter () = Some 4);
  assert (counter () = Some 6);
  let fibs = take_seq 8 (fibonacci ()) in
  assert (fibs = [0; 1; 1; 2; 3; 5; 8; 13]);
  let collatz = List.of_seq (collatz_seq 6) in
  assert (collatz = [6; 3; 10; 5; 16; 8; 4; 2; 1]);
  Printf.printf "✓ All tests passed\n"

Key Differences

Aspect	Rust	OCaml
State storage	Struct fields	Mutable `ref` or closure capture
Infinite	Return `Some(…)` always	`Seq.Cons(…, thunk)` always
Finite	Return `None` at end	`Seq.Nil` at end
Lazy	Pull-based (call `next`)	Thunk-based (call `seq ()`)
Adapter chain	Free from `Iterator` trait	Manual `seq_map`/`seq_filter`
Termination signal	`done_` flag or pattern on value	`Seq.Nil`

The Collatz iterator demonstrates a common pattern: an iterator that terminates when some condition on the internal state is met. This is cleaner than computing the full sequence upfront and is memory-efficient for long sequences.

OCaml Approach

OCaml uses mutable ref cells for stateful counters: let n = ref (start - step) in fun () -> n := !n + step; Some !n. Fibonacci uses the Seq module: let rec aux a b () = Seq.Cons(a, aux b (a+b)). The take_seq helper materialises the first n elements. Collatz is also expressed as a Seq with termination when n = 1. OCaml's approach requires explicit Seq.Cons/Seq.Nil construction while Rust's Option<T> return is simpler.

Full Source

//! 086: Custom Iterator with State — stateful iterators built by implementing
//! [`Iterator::next`].
//!
//! The OCaml original expresses stateful sequences three different ways:
//!
//! * `make_counter` captures a `ref` cell in a closure and mutates it on each
//!   call.
//! * `fibonacci` returns a `Seq.t`, i.e. a thunk producing `Seq.Cons(a, …)`
//!   at every step.
//! * `collatz_seq` also uses `Seq.t`, but signals termination with `Seq.Nil`
//!   when the current value reaches 1.
//!
//! Rust unifies all three under a single abstraction: a `struct` holding the
//! state plus an `impl Iterator` with a `next` method. Infinite iterators
//! always return `Some(_)` (the caller bounds them with [`Iterator::take`]);
//! finite iterators return `None` to signal termination. Implementing `next`
//! is sufficient to inherit the full iterator adapter chain (`map`, `filter`,
//! `collect`, `zip`, …).

/// Arithmetic progression: yields `start + step`, `start + 2*step`, … forever.
///
/// The OCaml `make_counter start step` closure captures a mutable `ref`
/// initialised to `start - step` and pre-increments before returning the
/// value. This struct achieves the same effect by storing `current` and
/// advancing it on every call to [`Iterator::next`].
///
/// Because the sequence is infinite, `next` always returns `Some(_)`. Use
/// [`Iterator::take`] to bound it.
#[derive(Debug, Clone, Copy)]
pub struct Counter {
    current: i32,
    step: i32,
}

impl Counter {
    /// Build a counter whose first yielded value is `start + step`.
    pub fn new(start: i32, step: i32) -> Self {
        Self {
            current: start,
            step,
        }
    }
}

impl Iterator for Counter {
    type Item = i32;

    fn next(&mut self) -> Option<Self::Item> {
        self.current += self.step;
        Some(self.current)
    }
}

/// Infinite Fibonacci sequence: `0, 1, 1, 2, 3, 5, 8, …`.
///
/// The OCaml version is a `Seq.t` whose thunk returns
/// `Seq.Cons(a, aux b (a + b))`. The Rust struct keeps the same two-element
/// state `(a, b)` and shifts it on every call. Bound the iterator with
/// [`Iterator::take`].
#[derive(Debug, Clone, Copy)]
pub struct Fib {
    a: u64,
    b: u64,
}

impl Fib {
    /// Start the sequence at `(0, 1)`.
    pub fn new() -> Self {
        Self { a: 0, b: 1 }
    }
}

impl Default for Fib {
    fn default() -> Self {
        Self::new()
    }
}

impl Iterator for Fib {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        let value = self.a;
        let next = self.a + self.b;
        self.a = self.b;
        self.b = next;
        Some(value)
    }
}

/// Collatz (3n+1) sequence from a chosen start, terminating at 1.
///
/// Each step halves `n` when it is even, or replaces it with `3n + 1` when
/// odd. The sequence is finite: once 1 is yielded, subsequent calls to
/// `next` return [`None`]. Mirrors OCaml's `collatz_seq`, which wraps the
/// final element in `Seq.Cons(1, fun () -> Seq.Nil)`.
#[derive(Debug, Clone, Copy)]
pub struct Collatz {
    n: u64,
    finished: bool,
}

impl Collatz {
    /// Build a Collatz iterator starting at `start`. The first yielded value
    /// is `start` itself.
    pub fn new(start: u64) -> Self {
        Self {
            n: start,
            finished: false,
        }
    }
}

impl Iterator for Collatz {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            return None;
        }
        let value = self.n;
        if self.n == 1 {
            self.finished = true;
        } else if self.n.is_multiple_of(2) {
            self.n /= 2;
        } else {
            self.n = 3 * self.n + 1;
        }
        Some(value)
    }
}

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

    #[test]
    fn counter_matches_ocaml_make_counter() {
        // OCaml: let counter = make_counter 0 2 in counter () = Some 2, etc.
        let v: Vec<i32> = Counter::new(0, 2).take(3).collect();
        assert_eq!(v, vec![2, 4, 6]);
    }

    #[test]
    fn counter_supports_negative_step() {
        let v: Vec<i32> = Counter::new(10, -3).take(4).collect();
        assert_eq!(v, vec![7, 4, 1, -2]);
    }

    #[test]
    fn fibonacci_first_eight_terms() {
        // OCaml: take_seq 8 (fibonacci ()) = [0; 1; 1; 2; 3; 5; 8; 13]
        let fibs: Vec<u64> = Fib::new().take(8).collect();
        assert_eq!(fibs, vec![0, 1, 1, 2, 3, 5, 8, 13]);
    }

    #[test]
    fn collatz_from_six_terminates_at_one() {
        // OCaml: List.of_seq (collatz_seq 6) = [6; 3; 10; 5; 16; 8; 4; 2; 1]
        let v: Vec<u64> = Collatz::new(6).collect();
        assert_eq!(v, vec![6, 3, 10, 5, 16, 8, 4, 2, 1]);
    }

    #[test]
    fn collatz_from_one_is_single_element() {
        let v: Vec<u64> = Collatz::new(1).collect();
        assert_eq!(v, vec![1]);
    }

    #[test]
    fn custom_iterators_compose_with_adapters() {
        // Implementing `next` is enough to get the whole adapter chain.
        let sum_of_even_fibs: u64 = Fib::new().take(10).filter(|n| n.is_multiple_of(2)).sum();
        assert_eq!(sum_of_even_fibs, 2 + 8 + 34);

        let zipped: Vec<(i32, u64)> = Counter::new(0, 1).zip(Fib::new()).take(4).collect();
        assert_eq!(zipped, vec![(1, 0), (2, 1), (3, 1), (4, 2)]);
    }
}

(* 086: Custom Iterator with State *)

(* Approach 1: Counter using mutable ref *)
let make_counter start step =
  let n = ref (start - step) in
  fun () -> n := !n + step; Some !n

(* Approach 2: Fibonacci via Seq *)
let fibonacci () =
  let rec aux a b () = Seq.Cons (a, aux b (a + b)) in
  aux 0 1

let take_seq n s =
  let rec aux n s acc =
    if n <= 0 then List.rev acc
    else match s () with
      | Seq.Nil -> List.rev acc
      | Seq.Cons (x, rest) -> aux (n - 1) rest (x :: acc)
  in
  aux n s []

(* Approach 3: Collatz sequence iterator *)
let collatz_seq start =
  let rec aux n () =
    if n = 1 then Seq.Cons (1, fun () -> Seq.Nil)
    else Seq.Cons (n, aux (if n mod 2 = 0 then n / 2 else 3 * n + 1))
  in
  aux start

(* Tests *)
let () =
  let counter = make_counter 0 2 in
  assert (counter () = Some 2);
  assert (counter () = Some 4);
  assert (counter () = Some 6);
  let fibs = take_seq 8 (fibonacci ()) in
  assert (fibs = [0; 1; 1; 2; 3; 5; 8; 13]);
  let collatz = List.of_seq (collatz_seq 6) in
  assert (collatz = [6; 3; 10; 5; 16; 8; 4; 2; 1]);
  Printf.printf "✓ All tests passed\n"

✓ Tests Rust test suite

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

    #[test]
    fn counter_matches_ocaml_make_counter() {
        // OCaml: let counter = make_counter 0 2 in counter () = Some 2, etc.
        let v: Vec<i32> = Counter::new(0, 2).take(3).collect();
        assert_eq!(v, vec![2, 4, 6]);
    }

    #[test]
    fn counter_supports_negative_step() {
        let v: Vec<i32> = Counter::new(10, -3).take(4).collect();
        assert_eq!(v, vec![7, 4, 1, -2]);
    }

    #[test]
    fn fibonacci_first_eight_terms() {
        // OCaml: take_seq 8 (fibonacci ()) = [0; 1; 1; 2; 3; 5; 8; 13]
        let fibs: Vec<u64> = Fib::new().take(8).collect();
        assert_eq!(fibs, vec![0, 1, 1, 2, 3, 5, 8, 13]);
    }

    #[test]
    fn collatz_from_six_terminates_at_one() {
        // OCaml: List.of_seq (collatz_seq 6) = [6; 3; 10; 5; 16; 8; 4; 2; 1]
        let v: Vec<u64> = Collatz::new(6).collect();
        assert_eq!(v, vec![6, 3, 10, 5, 16, 8, 4, 2, 1]);
    }

    #[test]
    fn collatz_from_one_is_single_element() {
        let v: Vec<u64> = Collatz::new(1).collect();
        assert_eq!(v, vec![1]);
    }

    #[test]
    fn custom_iterators_compose_with_adapters() {
        // Implementing `next` is enough to get the whole adapter chain.
        let sum_of_even_fibs: u64 = Fib::new().take(10).filter(|n| n.is_multiple_of(2)).sum();
        assert_eq!(sum_of_even_fibs, 2 + 8 + 34);

        let zipped: Vec<(i32, u64)> = Counter::new(0, 1).zip(Fib::new()).take(4).collect();
        assert_eq!(zipped, vec![(1, 0), (2, 1), (3, 1), (4, 2)]);
    }
}

Deep Comparison

Core Insight

Custom iterators encapsulate state. Each call to next() advances the internal state machine. This is more explicit than OCaml's closure-based sequences.

OCaml Approach

• Closure captures mutable ref: let r = ref 0 in fun () -> incr r; !r

• Or pure: state threaded through Seq thunks

Rust Approach

• Struct with state fields

• impl Iterator with next(&mut self)

• State mutation is explicit and checked by borrow rules

Comparison Table

Feature	OCaml	Rust
State	Closure ref / pure threading	Struct fields
Mutation	`ref` / `incr`	`&mut self`
Infinite	Lazy Seq	Iterator (never returns None)

Exercises

Implement a Prime iterator that produces prime numbers infinitely using a sieve stored in a HashSet.

Add DoubleEndedIterator for Counter (since it has a known step, going backwards is well-defined).

Write a Zip3<A, B, C> iterator that zips three iterators into triples.

Implement a RunLength<I: Iterator> adapter that wraps any iterator and emits (count, value) pairs.

In OCaml, implement a lazy Sieve of Eratosthenes using Seq that streams primes without storing the full sieve.

Open Source Repos

functional-rust

View the source for this example on GitHub — OCaml and Rust side by side in the repo.

Rust