Modeling State Machines

A state machine is a simple way to describe a program that moves between fixed states.

For example, a network connection may be:

disconnected
connecting
connected
closing

At any moment, it is in exactly one state. It is not both connected and closing. It is one state, then another state, then another.

This is exactly the kind of problem where Zig enums and tagged unions are useful.

Start with an Enum

If each state only needs a name, use an enum.

const ConnectionState = enum {
    disconnected,
    connecting,
    connected,
    closing,
};

Now a value of type ConnectionState must be one of those four states.

var state: ConnectionState = .disconnected;

You can change it later:

state = .connecting;
state = .connected;

This is already safer than using strings or integers.

var state = "connected";

A string can contain a typo.

var state = "conected";

The compiler cannot know that this is wrong.

With an enum, a typo is rejected:

var state: ConnectionState = .conected; // error

Use `switch` to Handle States

A state machine usually has behavior for each state.

fn canSend(state: ConnectionState) bool {
    return switch (state) {
        .disconnected => false,
        .connecting => false,
        .connected => true,
        .closing => false,
    };
}

This function answers one question: can we send data in this state?

Only the connected state returns true.

The switch is exhaustive. It handles every possible ConnectionState.

That gives you a useful safety property. If you add a new state later, Zig can force you to update the switch.

State Transitions

A state transition is a move from one state to another.

For example:

disconnected -> connecting
connecting -> connected
connected -> closing
closing -> disconnected

You can write that as a function:

const ConnectionState = enum {
    disconnected,
    connecting,
    connected,
    closing,
};

fn next(state: ConnectionState) ConnectionState {
    return switch (state) {
        .disconnected => .connecting,
        .connecting => .connected,
        .connected => .closing,
        .closing => .disconnected,
    };
}

Use it like this:

var state: ConnectionState = .disconnected;

state = next(state); // connecting
state = next(state); // connected
state = next(state); // closing

The state changes are explicit. There is no hidden logic.

Some Transitions Should Be Rejected

Real state machines often have rules.

For example, a connection may allow these transitions:

disconnected -> connecting
connecting -> connected
connecting -> disconnected
connected -> closing
closing -> disconnected

But this should not be allowed:

disconnected -> connected

You can encode that rule with an error union.

const TransitionError = error{
    InvalidTransition,
};

fn transition(
    current: ConnectionState,
    next_state: ConnectionState,
) TransitionError!ConnectionState {
    return switch (current) {
        .disconnected => switch (next_state) {
            .connecting => next_state,
            else => error.InvalidTransition,
        },

        .connecting => switch (next_state) {
            .connected, .disconnected => next_state,
            else => error.InvalidTransition,
        },

        .connected => switch (next_state) {
            .closing => next_state,
            else => error.InvalidTransition,
        },

        .closing => switch (next_state) {
            .disconnected => next_state,
            else => error.InvalidTransition,
        },
    };
}

The return type says the function may fail:

TransitionError!ConnectionState

That means: either return the new state, or return error.InvalidTransition.

Using the Transition Function

const std = @import("std");

const ConnectionState = enum {
    disconnected,
    connecting,
    connected,
    closing,
};

const TransitionError = error{
    InvalidTransition,
};

fn transition(
    current: ConnectionState,
    next_state: ConnectionState,
) TransitionError!ConnectionState {
    return switch (current) {
        .disconnected => switch (next_state) {
            .connecting => next_state,
            else => error.InvalidTransition,
        },

        .connecting => switch (next_state) {
            .connected, .disconnected => next_state,
            else => error.InvalidTransition,
        },

        .connected => switch (next_state) {
            .closing => next_state,
            else => error.InvalidTransition,
        },

        .closing => switch (next_state) {
            .disconnected => next_state,
            else => error.InvalidTransition,
        },
    };
}

pub fn main() void {
    var state: ConnectionState = .disconnected;

    state = transition(state, .connecting) catch {
        std.debug.print("invalid transition\n", .{});
        return;
    };

    state = transition(state, .connected) catch {
        std.debug.print("invalid transition\n", .{});
        return;
    };

    std.debug.print("state = {}\n", .{state});
}

Output:

state = main.ConnectionState.connected

The exact printed name may include the type path, but the important value is connected.

Keep State and Data Together

Sometimes a state needs extra data.

For example:

disconnected
connecting, with attempt number
connected, with peer address
closing

An enum alone cannot store that extra data.

Use a tagged union.

const Connection = union(enum) {
    disconnected,
    connecting: u8,
    connected: []const u8,
    closing,
};

This says:

disconnected carries no data
connecting carries a u8 attempt number
connected carries a peer address
closing carries no data

Now the state and its data stay together.

Handling State with Payloads

const std = @import("std");

const Connection = union(enum) {
    disconnected,
    connecting: u8,
    connected: []const u8,
    closing,

    fn print(self: Connection) void {
        switch (self) {
            .disconnected => {
                std.debug.print("disconnected\n", .{});
            },
            .connecting => |attempt| {
                std.debug.print("connecting, attempt {}\n", .{attempt});
            },
            .connected => |peer| {
                std.debug.print("connected to {s}\n", .{peer});
            },
            .closing => {
                std.debug.print("closing\n", .{});
            },
        }
    }
};

pub fn main() void {
    var conn = Connection.disconnected;

    conn = Connection{ .connecting = 1 };
    conn.print();

    conn = Connection{ .connected = "127.0.0.1:9000" };
    conn.print();

    conn = Connection.closing;
    conn.print();
}

Output:

connecting, attempt 1
connected to 127.0.0.1:9000
closing

This model is stronger than keeping state and data in separate fields.

Avoid Loose Optional Fields

A weak model might look like this:

const Connection = struct {
    state: ConnectionState,
    attempt: ?u8 = null,
    peer: ?[]const u8 = null,
};

This can represent invalid combinations:

const conn = Connection{
    .state = .disconnected,
    .attempt = 3,
    .peer = "127.0.0.1:9000",
};

A disconnected connection should not have a peer address.

A tagged union avoids that:

const Connection = union(enum) {
    disconnected,
    connecting: u8,
    connected: []const u8,
    closing,
};

Now each state carries only the data that belongs to that state.

Model Events Separately

Many state machines respond to events.

For a connection, events might be:

start
connected
failed
close
closed

Use another enum or tagged union for events.

const Event = enum {
    start,
    connected,
    failed,
    close,
    closed,
};

Then write a function that combines current state and event:

fn apply(state: ConnectionState, event: Event) ?ConnectionState {
    return switch (state) {
        .disconnected => switch (event) {
            .start => .connecting,
            else => null,
        },

        .connecting => switch (event) {
            .connected => .connected,
            .failed => .disconnected,
            else => null,
        },

        .connected => switch (event) {
            .close => .closing,
            else => null,
        },

        .closing => switch (event) {
            .closed => .disconnected,
            else => null,
        },
    };
}

The return type is:

?ConnectionState

That means: either a new state, or null if the event is not valid in the current state.

Complete Event-Driven Example

const std = @import("std");

const ConnectionState = enum {
    disconnected,
    connecting,
    connected,
    closing,
};

const Event = enum {
    start,
    connected,
    failed,
    close,
    closed,
};

fn apply(state: ConnectionState, event: Event) ?ConnectionState {
    return switch (state) {
        .disconnected => switch (event) {
            .start => .connecting,
            else => null,
        },

        .connecting => switch (event) {
            .connected => .connected,
            .failed => .disconnected,
            else => null,
        },

        .connected => switch (event) {
            .close => .closing,
            else => null,
        },

        .closing => switch (event) {
            .closed => .disconnected,
            else => null,
        },
    };
}

pub fn main() void {
    var state: ConnectionState = .disconnected;

    const events = [_]Event{
        .start,
        .connected,
        .close,
        .closed,
    };

    for (events) |event| {
        state = apply(state, event) orelse {
            std.debug.print("invalid event: {}\n", .{event});
            return;
        };

        std.debug.print("state = {}\n", .{state});
    }
}

This walks through a valid sequence:

disconnected -> connecting -> connected -> closing -> disconnected

Put the Logic Near the Type

You can place the transition logic inside a struct or enum namespace.

const ConnectionState = enum {
    disconnected,
    connecting,
    connected,
    closing,

    fn canSend(self: ConnectionState) bool {
        return switch (self) {
            .disconnected => false,
            .connecting => false,
            .connected => true,
            .closing => false,
        };
    }
};

Then call:

const ok = state.canSend();

This keeps state-related behavior close to the state definition.

For larger machines, you may use a struct:

const ConnectionMachine = struct {
    state: ConnectionState,

    pub fn init() ConnectionMachine {
        return .{
            .state = .disconnected,
        };
    }

    pub fn apply(self: *ConnectionMachine, event: Event) bool {
        const next_state = applyState(self.state, event) orelse return false;
        self.state = next_state;
        return true;
    }
};

The struct stores the current state. Its methods change that state safely.

Prefer Explicit States

Do not hide important states inside booleans.

This is weak:

const Connection = struct {
    is_connected: bool,
    is_connecting: bool,
};

It can represent invalid combinations:

is_connected = true
is_connecting = true

A connection should not usually be both connected and connecting.

An enum is clearer:

const ConnectionState = enum {
    disconnected,
    connecting,
    connected,
    closing,
};

One value, one state.

Prefer Explicit Events

Do not pass vague strings as events:

handleEvent("connected");

Use an enum:

handleEvent(.connected);

This lets the compiler check spelling and valid cases.

If an event has data, use a tagged union:

const Event = union(enum) {
    start,
    connected: []const u8,
    failed: ErrorCode,
    close,
};

Now event data is attached to the correct event.

The Main Idea

A state machine becomes easier to write when states are explicit.

Use enums when states are just names:

const State = enum {
    idle,
    running,
    stopped,
};

Use tagged unions when states carry data:

const State = union(enum) {
    idle,
    running: JobId,
    stopped,
};

Use switch to handle every state. Put transition rules in one place. Let the type system prevent invalid combinations.