minifsm

statemachine.test.ts

import { describe, it } from "node:test"
import assert from "node:assert"
import { StateMachine } from "./statemachine"

describe("StateMachine", () => {
  describe("Stoplight", () => {
    interface State {
      name: "Green" | "Yellow" | "Red"
      count?: number
    }

    type Action = { type: "TIMER" }

    const createStoplight = (initialState: State) =>
      new StateMachine<State, Action>(initialState, {
        Green: {
          TIMER: () => ({ name: "Yellow" }),
        },
        Yellow: {
          TIMER: () => ({ name: "Red" }),
        },
        Red: {
          TIMER: () => ({ name: "Green" }),
        },
      })

    it("transitions from Green to Yellow", () => {
      const fsm = createStoplight({ name: "Green" })

      assert.strictEqual(fsm.state.name, "Green")
      fsm.transition({ type: "TIMER" })
      assert.strictEqual(fsm.state.name, "Yellow")
    })

    it("transitions from Yellow to Red", () => {
      const fsm = createStoplight({ name: "Yellow" })

      assert.strictEqual(fsm.state.name, "Yellow")
      fsm.transition({ type: "TIMER" })
      assert.strictEqual(fsm.state.name, "Red")
    })

    it("transitions from Red to Green", () => {
      const fsm = createStoplight({ name: "Red" })

      assert.strictEqual(fsm.state.name, "Red")
      fsm.transition({ type: "TIMER" })
      assert.strictEqual(fsm.state.name, "Green")
    })

    it("does a whole cycle", () => {
      const fsm = createStoplight({ name: "Green" })

      assert.strictEqual(fsm.state.name, "Green")
      fsm.transition({ type: "TIMER" })
      assert.strictEqual(fsm.state.name, "Yellow")
      fsm.transition({ type: "TIMER" })
      assert.strictEqual(fsm.state.name, "Red")
      fsm.transition({ type: "TIMER" })
      assert.strictEqual(fsm.state.name, "Green")
    })
  })

  describe("State", () => {
    interface State {
      name: "on" | "off"
      count?: number
    }

    type Action = { type: "TOGGLE" }

    const createState = (initialState: State) =>
      new StateMachine<State, Action>(initialState, {
        on: {
          TOGGLE: (_, { count }) => ({
            name: "off",
            count: count ? count + 1 : 1,
          }),
        },
        off: {
          TOGGLE: () => ({ name: "on" }),
        },
      })

    it("increments count when on is toggled", () => {
      const fsm = createState({ name: "on", count: 0 })

      assert.strictEqual(fsm.state.name, "on")
      assert.strictEqual(fsm.state.count, 0)
      fsm.transition({ type: "TOGGLE" })
      assert.strictEqual(fsm.state.name, "off")
      assert.strictEqual(fsm.state.count, 1)
      fsm.transition({ type: "TOGGLE" })
      assert.strictEqual(fsm.state.name, "on")
      assert.strictEqual(fsm.state.count, 1)
      fsm.transition({ type: "TOGGLE" })
      assert.strictEqual(fsm.state.name, "off")
      assert.strictEqual(fsm.state.count, 2)
    })
  })

  describe("Transitions", () => {
    interface State {
      name: "A" | "B" | "C"
    }

    type Action = { type: "TO_A" } | { type: "TO_B" } | { type: "TO_C" }

    const createTransitions = (initialState: State) =>
      new StateMachine<State, Action>(initialState, {
        A: {
          TO_B: () => ({ name: "B" }),
          TO_C: () => ({ name: "C" }),
        },
        B: {
          TO_C: () => ({ name: "C" }),
        },
        C: {
          TO_B: () => ({ name: "B" }),
        },
      })

    it("transitions from A to B", () => {
      const fsm = createTransitions({ name: "A" })

      assert.strictEqual(fsm.state.name, "A")
      fsm.transition({ type: "TO_B" })
      assert.strictEqual(fsm.state.name, "B")
    })

    it("transitions from A to C", () => {
      const fsm = createTransitions({ name: "A" })

      assert.strictEqual(fsm.state.name, "A")
      fsm.transition({ type: "TO_C" })
      assert.strictEqual(fsm.state.name, "C")
    })

    it("cannot transition from B to A", () => {
      const fsm = createTransitions({ name: "B" })
      assert.throws(
        () => {
          assert.strictEqual(fsm.state.name, "B")
          fsm.transition({ type: "TO_A" })
        },
        {
          message: "Invalid transition: B -> TO_A",
        }
      )
    })

    it("transitions from B to C", () => {
      const fsm = createTransitions({ name: "B" })

      assert.strictEqual(fsm.state.name, "B")
      fsm.transition({ type: "TO_C" })
      assert.strictEqual(fsm.state.name, "C")
    })

    it("transitions from C to B", () => {
      const fsm = createTransitions({ name: "C" })

      assert.strictEqual(fsm.state.name, "C")
      fsm.transition({ type: "TO_B" })
      assert.strictEqual(fsm.state.name, "B")
    })

    it("canot transition from C to A", () => {
      const fsm = createTransitions({ name: "C" })
      assert.throws(
        () => {
          assert.strictEqual(fsm.state.name, "C")
          fsm.transition({ type: "TO_A" })
        },
        {
          message: "Invalid transition: C -> TO_A",
        }
      )
    })
  })

  describe(".reducer", () => {
    interface State {
      name: "X" | "Y"
      count?: number
    }

    type Action = { type: "INCREMENT" } | { type: "DECREMENT" }

    const createCounter = (initialState: State) =>
      new StateMachine<State, Action>(initialState, {
        X: {
          INCREMENT: (_, { count }) => ({ name: "X", count: (count || 0) + 1 }),
          DECREMENT: (_, { count }) => ({ name: "X", count: (count || 0) - 1 }),
        },
        Y: {
          INCREMENT: () => ({ name: "Y" }),
          DECREMENT: () => ({ name: "Y" }),
        },
      })

    it("increments count in state X", () => {
      const { reducer, state } = createCounter({ name: "X", count: 0 })

      const newState = reducer(state, { type: "INCREMENT" })
      assert.strictEqual(newState.name, "X")
      assert.strictEqual(newState.count, 1)

      const newState2 = reducer(newState, { type: "INCREMENT" })
      assert.strictEqual(newState2.name, "X")
      assert.strictEqual(newState2.count, 2)

      const newState3 = reducer(newState2, { type: "DECREMENT" })
      assert.strictEqual(newState3.name, "X")
      assert.strictEqual(newState3.count, 1)
    })
  })
})

statemachine.ts

export type TransitionGraph<
  T extends { name: string },
  U extends { type: string }
> = {
  [view in T["name"]]: {
    [action in U["type"]]?: (action: U, state: T) => T
  }
}

/**
 * Extremely simple and featureless finite state machine
 */
export class StateMachine<
  T extends { name: string },
  U extends { type: string }
> {
  /**
   * @constructor
   */
  constructor(public state: T, private graph: TransitionGraph<T, U>) {}

  /**
   * Transition to a new state
   * @throws {Error} if the transition is invalid
   */
  transition(action: U): void {
    if (this.can(action.type)) {
      const nextStateProducer =
        this.graph[this.state.name as T["name"]][action.type]

      if (nextStateProducer) {
        this.state = {
          ...this.state,
          ...nextStateProducer(action, this.state),
        }
      }
    } else {
      throw new Error(
        `Invalid transition: ${this.state.name} -> ${action.type}`
      )
    }
  }

  /**
   * Check if an action is valid in the given state
   */
  can(action: U["type"]): action is keyof TransitionGraph<T, U> {
    return (
      typeof this.graph[this.state.name as T["name"]][action] === "function"
    )
  }

  /**
   * Reducer function to be used with useReducer
   * @throws {Error} if the transition is invalid
   */
  reducer = (_: T, action: U): T => {
    if (this.can(action.type)) {
      this.transition(action)
    }
    return this.state
  }
}