Iterators and Finite State Machines

I have a soft spot for iterators.

An iterator is an abstraction in programming that represents a sequence of values that can be iterated (looped) over one at a time. It provides a way to access elements in a collection, data structure, or sequence without exposing the underlying details of how the elements are stored or generated. Iterators are widely used in programming languages to simplify the process of iterating through data.

ChatGPT (validated by this human)

Iterators are extremely versatile, and a crucial component of virtually every programming language. A common use of iterators is creating some sort of data structure, such as a vector, in memory and iterating over the collection. There are a lot of scenarios where this is perfectly suitable, and even the best way to do it. However, this is a sub-optimal approach in situations where the data is consumed exactly once (or a handful of times), and never manipulated further. Managing memory on the heap is expensive for a variety of reasons, such as operating system calls, pointer management, and cache misses. What if, instead of generating a sequence into a collection and then iterating over the collection, we just iterated over the sequence while we generate it?

A Simple Problem

Let us look at a simple example: a method that iterates through the set of unique integer partitions that contain N or fewer elements. For example, if N=4, the unique integer partitions of 6 would be:

/*
[6]
[5, 1]
[4, 2]
[4, 1, 1]
[3, 3]
[3, 2, 1]
[3, 1, 1, 1]
[2, 2, 2]
[2, 2, 1, 1]
*/

Note that the elements contained in partitions are descending, and the partitions themselves are listed in a descending order. We can think of sorting the elements within a partition as a sort of canonical form, where we can ensure that no partition is included more than once. Listing the partitions in this descending order also ensures that each partition is included at least once. All of our implementations will follow these ordering patterns.

The integer partition function appears frequently in discrete mathematics problems. As an extension, we could envision variants of this function with other properties such as partitions whose members are unique, or generating all orderings of each partition rather than simply the canonical form.

Recursive Implementation

This method is fairly easy to implement as a recursive function.

fn get_partition_set<const N: usize>(x: u32) -> Vec<[u32; N]> {
    fn inner<const N: usize>(
        output: &mut Vec<[u32; N]>,
        stub: [u32; N],
        index: usize,
        remainder: u32,
    ) {
        if remainder == 0 {
            output.push(item);
        } else if index < N {
            let max_value = match index {
                0 => remainder,
                _ => stub[index - 1].min(remainder),
            };
            let mut new_stub = stub;
            for i in (1..=max_value).rev() {
                new_stub[index] = i;
                inner(output, new_stub, index + 1, remainder - i);
            }
        }
    }

    let mut output = vec![];
    let mut stub = [0; N];
    inner::<N>(&mut output, &mut stub, 0, x);
    output
}

And when we test it, we get the expected result.

#[test]
fn test_get_partition_set_function() {
    let set = get_partition_set::<4>(6);
    assert_eq!(set.len(), 9);
    for partition in &set {
        println!("{partition:?}");
    }
}

/*
[6, 0, 0, 0]
[5, 1, 0, 0]
[4, 2, 0, 0]
[4, 1, 1, 0]
[3, 3, 0, 0]
[3, 2, 1, 0]
[3, 1, 1, 1]
[2, 2, 2, 0]
[2, 2, 1, 1]
*/

Let’s look a little more closely at what’s going on in this simple implementation.

Each call of the recursive function locks in the values of the first elements of stub, and we can think of all following elements of stub as a smaller version of the same problem. This is the role of the function signature, and it’s apparent through the stub, index, and remainder inputs. stub is responsible for keeping track of the locked-in portion of the solution, index keeps track of how many elements are already locked in (and whether we need to stop recursing), and remainder keeps track of how much of the original integer is left to partition.

fn inner<const N: usize>(
    output: &mut Vec<[u32; N]>,
    stub: [u32; N],
    index: usize,
    remainder: u32,
)

Each instance of the recursive function can fork into an arbitrary number of instances of itself. This implementation treats the state of the function not as a loop or finite state machine, but rather as something more complicated and powerful – a stack.

Each call of inner pushes a new frame to the call stack of the program. Each frame contains information about the current stage of the recursive problem. Whenever a frame is finished, that frame is popped from the stack, and execution resumes from the point in the prior frame where the next frame was instantiated. To re-implement this function as an iterator, we need to emulate this process in a struct.

Iterator Implementation

The end goal for creating an iterator in Rust is making it conform to this interface:

impl<const N: usize> Iterator for IterPartitionSet<N> {
    type Item = [u8; N];
    fn next(&mut self) -> Option<Self::Item>;
}

Essentially, we do two things every time we call next:

1. We return Some(_) until the iterator runs out and we return None instead
2. We advance the state of the iterator so that the next invocation of next returns the next iterator value

Implementing iterators can be challenging because it’s frequently difficult or impossible to transform the original code into a form that can be stored in an (ideally) fixed-size struct and have next logic expressible as a hierarchy of if, while, and other control blocks. Finite state machines (FSM) are a good solution to this problem. The FSM is a general concept that shows up frequently in computer science and computer engineering applications and can effectively and intuitively, if not always elegantly, tackle complex control flows.

We will take a couple of steps to transform this function to an iterator:

1. Replace moves and copies associated with the signature of inner with mutable references
2. Rearrange and label the control flow
3. Create an FSM

Mutable References

This is a natural first step because all iterator logic ultimately needs to be stored in a struct which ideally has constant size. We create these mutable “global” variables in the stack of the get_partition_set and pass them by reference to inner. These variables are stub, index, and remainder. We ignore output as it serves as a storage buffer for the result of the calculation, and the logic of our method does not depend on its contents.

The transformation looks something like this:

fn get_partition_set<const N: usize>(x: u32) -> Vec<[u32; N]> {
    fn inner<const N: usize>(
        output: &mut Vec<[u32; N]>,
        stub: &mut [u32; N],
        index: &mut usize,
        remainder: &mut u32,
    ) {
        if *remainder == 0 {
            output.push(*stub);
        } else if *index < N {
            match *index {
                0 => stub[*index] = *remainder,
                _ => stub[*index] = stub[*index - 1].min(*remainder),
            }
            *remainder -= stub[*index];
            while stub[*index] > 0 {
                *index += 1;
                inner(output, stub, index, remainder);
                *index -= 1;
                stub[*index] -= 1;
                *remainder += 1;
            }
        }
    }

    let mut output = vec![];
    let mut stub = [0; N];
    let mut index = 0;
    let mut remainder = x;
    inner::<N>(&mut output, &mut stub, &mut index, &mut remainder);
    output
}

There really isn’t much to it. The only thing that requires further explanation is the transformation from a for loop to a while loop. Our intent is to vary stub[index] between 1 and the greatest value that stub[index] can contain. The for loop involves i, which is a redundant variable. The while loop achieves our intent by storing the progress of the loop with state we need to maintain anyway. Coincidentally, it also helps us with our next transformation.

Quick aside: It is possible to infer index from stub by finding the first index in stub that contains 0. It is also possible to infer remainder by subtracting the sum of stub from x. However, for the purposes of this exercise, and making a more efficient iterator, we explicitly maintain index and remainder.

Another quick aside: It would be cleaner to make inner a mutable lambda function that does not take any arguments and simply mutates stub, index, and remainder within its closure, but Rust does not allow self-referencing lambdas, and neither does Rust allow mutable closures for fn. Another choice would be placing stub, index, and remainder in a struct and writing a mutable inner method, but I arbitrarily did it the other way instead.

Control Flow

This next part is inspired by compilers. Sometime after the abstract syntax tree (AST) is parsed, but before an intermediate representation (IR) is created, the compiler produces produces a control flow graph (CFG). The nodes of a CFG are called “basic blocks”, which are collections of sequentially-executed statements, or basically lines of code. These basic blocks are connected by edges derived from a nested hierarchy of control statements, such as if or while, that have some sort of conditional branching behavior. I won’t dive into examples here, but there are lots of good sources on it.

Essentially, after some minor rearranging, we label the basic blocks in our code and reason about how they connect to each other.

fn get_partition_set<const N: usize>(x: u8) -> Vec<[u8; N]> {
    fn inner<const N: usize>(
        output: &mut Vec<[u8; N]>,
        stub: &mut [u8; N],
        index: &mut usize,
        remainder: &mut u8,
    ) {
        // START
        match *index {
            0 => stub[*index] = *remainder,
            _ => stub[*index] = stub[*index - 1].min(*remainder),
        }
        *remainder -= stub[*index];
        while stub[*index] > 0 {
            // LOOP
            if *remainder == 0 {
                // YIELD
                output.push(*stub);
            } else {
                // DESCEND
                *index += 1;
                if *index < N {
                    inner(output, stub, index, remainder);
                }
                // JOIN
                *index -= 1;
            }
            // RESUME
            stub[*index] -= 1;
            *remainder += 1;
        }
    }

    let mut output = vec![];
    let mut stub = [0; N];
    let mut index = 0;
    let mut remainder = x;
    inner::<N>(&mut output, &mut stub, &mut index, &mut remainder);
    // DONE
    output
}

The only major change from the previous code is that we have inverted the nesting of the if and while blocks. Similar analysis is possible without this inversion, but I thought it was somewhat easier this way. The inner call in the original code was basically inlined. If stub is a valid partition, then it is appended to output; otherwise, we recurse into inner anyway and bypass this check.

There were two basic patterns applied to create labels:

• Wherever there is a branching statement (if and while in this case), we prepend a label to the statement to which the branch can jump. and we prepend a label to the statement to which the branch can continue
• Wherever there is a function call, we prepend a label to the first statement of the function, and we append a label to the function call itself

Then, for each basic block, we determine its successor basic blocks and construct the CFG:

/*
START -> LOOP | JOIN
LOOP -> YIELD | DESCEND
YIELD -> RESUME
DESCEND -> START | JOIN
JOIN -> RESUME
RESUME -> LOOP | JOIN | DONE
DONE -> {}
*/

The labels chosen here are somewhat simplified from what the patterns would naively yield. We did this by eliminating any labels that have exactly one incoming edge and one outgoing edge. These labels are superfluous because the basic blocks they connect can be joined into a single basic block without affecting the control flow.

RESUME is the most conceptually difficult because it finishes at the end of the while loop as well as the end of the function. If the while loop terminates, then inner returns. Since inner is called from two locations, it is possible for inner to return to two locations, which are JOIN and DONE. If the while loop does not terminate, then the next label is LOOP.

Another exercise is reasoning about the preconditions that exist at these labels. This is not strictly necessary, but can be a useful debugging technique.

fn get_partition_set<const N: usize>(x: u8) -> Vec<[u8; N]> {
    fn inner<const N: usize>(
        output: &mut Vec<[u8; N]>,
        stub: &mut [u8; N],
        index: &mut usize,
        remainder: &mut u8,
    ) {
        // START
        //  - index < N
        //  - index == 0 || stub[index - 1] > 0
        //  - stub[index] == 0
        //  - remainder > 0
        match *index {
            0 => stub[*index] = *remainder,
            _ => stub[*index] = stub[*index - 1].min(*remainder),
        }
        *remainder -= stub[*index];
        while stub[*index] > 0 {
            // LOOP
            //  - index < N
            //  - stub[index] > 0
            if *remainder == 0 {
                // YIELD
                //  - index < N
                //  - stub[index] > 0
                //  - remainder == 0
                output.push(*stub);
            } else {
                // DESCEND
                //  - index < N
                //  - stub[index] > 0
                //  - remainder > 0
                *index += 1;
                if *index < N {
                    inner(output, stub, index, remainder);
                }
                // JOIN
                //  - index > 0
                //  - stub[index] == 0
                //  - remainder > 0
                *index -= 1;
            }
            // RESUME
            //  - index < N
            //  - stub[index] > 0
            stub[*index] -= 1;
            *remainder += 1;
        }
    }

    let mut output = vec![];
    let mut stub = [0; N];
    let mut index = 0;
    let mut remainder = x;
    inner::<N>(&mut output, &mut stub, &mut index, &mut remainder);
    // DONE
    //  - index == 0
    //  - stub[index] == 0
    //  - remainder > 0
    output
}

Finite State Machine

We are almost done! This last transformation takes the following steps:

1. Create a IterPartitionSet enum whose variants are the states we just identified
2. Create the iterator struct which contains stub, index, remainder, and a state
3. Write a private method for IterPartitionSet that modifies the iterator based on its current state and finds the next state for the iterator
4. Implement the next method for the Iterator trait

Step number 1 is easy.

#[derive(PartialEq, Eq)]
enum IterPartitionSetState {
    Start,
    Loop,
    Yield,
    Descend,
    Join,
    Resume,
    Done,
}

Step number 2 is also easy. It is a natural extension of the function signature we already had.

pub struct IterPartitionSet<const N: usize> {
    stub: [u8; N],
    index: usize,
    remainder: u8,
    state: IterPartitionSetState,
}

Step number 3 is a little trickier, but also pretty straightforward. All we have to do is take the basic blocks we identified in the previous section and transcribe their semantics into our state enum and our struct.

impl<const N: usize> IterPartitionSet<N> {
    fn advance(&mut self) {
        self.state = match self.state {
            IterPartitionSetState::Start => {
                let init_value = match self.index {
                    0 => self.remainder,
                    _ => self.stub[self.index - 1].min(self.remainder),
                };
                self.stub[self.index] = init_value;
                self.remainder -= init_value;
                IterPartitionSetState::Loop
            }
            IterPartitionSetState::Loop => {
                match self.remainder {
                    0 => IterPartitionSetState::Yield,
                    _ => IterPartitionSetState::Descend,
                }
            }
            IterPartitionSetState::Yield => {
                IterPartitionSetState::Resume
            }
            IterPartitionSetState::Descend => {
                self.index += 1;
                if self.index < N {
                    IterPartitionSetState::Start
                } else {
                    IterPartitionSetState::Join
                }
            }
            IterPartitionSetState::Join => {
                self.index -= 1;
                IterPartitionSetState::Resume
            }
            IterPartitionSetState::Resume => {
                self.stub[self.index] -= 1;
                self.remainder += 1;
                match (self.index, self.stub[self.index]) {
                    (0, 0) => IterPartitionSetState::Done,
                    (_, 0) => IterPartitionSetState::Join,
                    _ => IterPartitionSetState::Loop,
                }
            }
            IterPartitionSetState::Done => {
                IterPartitionSetState::Done
            }
        }
    }
}

The sticky wicket is figuring out how to initialize the struct. Here, we assume/know that the iterator will contain at least one valid integer partition, and we initialize it in the Yield state as though it has just found this first partition.

impl<const N: usize> IterPartitionSet<N> {
    pub fn new(x: u8) -> Self {
        let mut output = Self {
            stub: [0; N],
            index: 0,
            remainder: 0,
            state: IterPartitionSetState::Yield,
        };
        output.stub[0] = x;
        output
    }
}

Now that we have all of these pieces in place, we just need to implement the Iterator trait interface using the advance method. This is probably a bit clunkier than it needs to be, but it gets the job done, and I find it easy to understand.

impl<const N: usize> Iterator for IterPartitionSet<N> {
    type Item = [u8; N];

    fn next(&mut self) -> Option<Self::Item> {
        if self.remainder == 0 {
            // The iterator starts in a valid output state
            let output = Some(self.stub);

            // Advance the state to the next valid output (or the "done" state)
            self.advance();
            while self.state != IterPartitionSetState::Yield
                && self.state != IterPartitionSetState::Done
            {
                self.advance();
            }

            output
        } else {
            // No more outputs
            None
        }
    }
}

So now, we would use the iterator as follows:

fn do_something() {
    for partition in IterPartitionSet::<4>::new(6) {
        println!("{partition}")
    }
}

Afterword

Limitations of FSMs

The technique shown in this post is hopefully illustrative and inspiring. Among the classes of automata, FSMs are relatively limited. They can, definitionally, only represent a finite number of states, and consequently can only model algorithms with a finite number of states. More powerful alternatives, such as pushdown automatons, should be considered where FSMs are insufficient.

Iterators vs. Generators

This article was intended to demystify Rust iterators and make a case for low-memory alternatives to containers. However, the problem we solved with an iterator is technically an example of a generator. Iterators are more general than generators, and they have some subtle differences in typical usage. Their implementations are very similar.

Unstable Rust has a draft for a Generator trait, which has striking resemblance to the iterator we implemented in this article.

pub trait Generator<R = ()> {
    type Yield;
    type Return;
    fn resume(self: Pin<&mut Self>, resume: R)
          -> GeneratorState<Self::Yield, Self::Return>;
}

This trait interface translates to our state machine via Yield and Return (which we call Done), but our generator would have additional states to represent the more complicated call structure. The differences between the Iterator and Generator trait interfaces echo the initialization and completion states for our iterator. I leave it as an exercise for the reader to repackage IterPartitionSet as a generator.