Lazy state machines

I’m not sure if this is a new thing or not, and I’m too lazy to look it up as it’s 3 AM right now, so here it is.

I’ve been thinking about state machines lately, and how Clojure’s multimethods are a cool way to implement a state machine. Multimethods are somewhat like pattern matching, in a way, and you can use them to do some declarative programming in Clojure, like this:

(defmulti factorial identity)

(defmethod factorial 0 [_] 1)

(defmethod factorial 1 [_] 1)

(defmethod factorial :default [n]
  (*' n (factorial (dec n))))

(factorial 5) ;; => 120

So, here, we’re declaring the factorial multimethod that dispatches on numbers. Let’s ignore negative numbers for now - pretend we’re ancient Greeks, and there is no such thing as a negative number.

So, when factorial sees a value of 0, it simply returns 1. The same thing happens if it sees 1. Otherwise, we’ll multiply the number with the result of calling factorial with the said number decreased by 1.

In reality, there’s and if hiding in here, as we might have just written factorial the conventional way:

(defn factorial [n]
  (cond (= n 0) 1
        (= n 1) 1
        :else   (*' n (factorial (dec n)))))

(factorial 5) ;; => 120

It’s the same thing, except we’ve avoided concretion of our function, allowing it to be extended with arbitrary dispatch values from outside, without modifying existing code.

I say it’s declarative because we’re declaring how a function should behave, given a specific input. I say it’s like pattern matching because we kinda match a pattern with the dispatching function.

E.g. we can write a pattern matching function that can do stuff based on a pattern of arguments given:

(defmulti foo (fn [& args] args))

(defmethod foo [0 0] [& _] "draw")
(defmethod foo [1 0] [& _] "win")
(defmethod foo [0 1] [& _] "loose")
(defmethod foo [1 1] [& _] "snake eyes")

(foo 1 1) ;; => "snake eyes"

It’s not quite pattern matching, as it doesn’t do binding, and there’s no match-all/I-don’t-care clause, but it can be used like that, and we certainly do.

Now, the factorial example is also a kind of a state machine. When it sees the number 1 or 0 its action is simple - return 1 and halt. If it sees anything else, it changes the state by decreasing a number, and proceeds. After the state reaches 1 or 0, the stack of actions is reduced, and we get the final result.

However, there’s a problem. Such a state machine can’t run for an arbitrary amount of time - it will consume a lot of the stack, and eventually fail. How can we avoid that?

We know, there’s a way of doing recursion in Clojure, using the recur special. It isn’t like a proper recursion though, rather, it’s a jump to the beginning of the function which also rebinds its arguments. It’s a handy thing, but it can’t handle mutual recursion.

The other way of doing recursion in Clojure is the trampoline function. Instead of calling the function you want to recurse into, you return it as a callback, and trampoline executes it instead. Because of that the stack is constantly incremented and decremented, but we can’t have non-tail mutual recursion.

Yet, there’s a way of having non-tail mutual recursion that doesn’t blow up the stack - lazy sequences!

Here’s a great example, a lazy sequence of Fibonacci numbers:

(defn fib-seq [a b]
  (cons a (lazy-seq (fib-seq b (+' a b)))))

(defn fib [n]
  (first (drop (inc n) (fib-seq 0 1))))

(fib 9) ;; => 55

(map fib (range 10)) ;; => (1 1 2 3 5 8 13 21 34 55)

(take 10 (fib-seq 0 1)) ;; => (0 1 1 2 3 5 8 13 21 34)

As you can see, fib-seq creates a lazy sequence by constantly calling itself. Its call is delayed until we consume the sequence, so it acts mostly like a trampoline, except we can retain all previous results.

Thinking of this, I came up with this kind of state machine:

(defmulti machine first)

(defmethod machine :a [[state & states]]
  (cons "A" (lazy-seq (machine (cons :c states)))))

(defmethod machine :b [[state & states]]
  (cons "B" (lazy-seq (machine states))))

(defmethod machine :c [[state & states]]
  (cons "C" (lazy-seq (machine (cons :b states)))))

(defmethod machine nil [[state & states]]
  ["done"])

This machine is simple - when it sees the state :a it will produce the "A" as a result, and change to the state :c. In state :c it produces the result "C" and changes to the state :b. Finally, in the state :b the machine produces the result "B" and doesn’t change its state. Then it halts. So, if we call (machine [:a]) we’ll get the following list ("A" "C" "B" "done").

One important observation would be that, unlike the fibbonaci, this state machine accepts a sequence of states, and modifies it to change state. This means that we can provide more tasks in a sequence, like:

lazy-state-machines
> (machine [:a :a])
("A" "C" "B" "A" "C" "B" "done")
lazy-state-machines
> (machine [:a :c])
("A" "C" "B" "C" "B" "done")
lazy-state-machines
> (machine [:b :b])
("B" "B" "done")
lazy-state-machines
> (machine [:b :b :a :c])
("B" "B" "A" "C" "B" "C" "B" "done")

And, since this state machine is lazy, we don’t have to consume all of its results, as it might never actually halt, say there’s one more state :d:

(defmethod machine :d [[state & states]]
  (cons "D" (lazy-seq (machine (cons :d states)))))

No other state triggers this one, so the machine would always halt if we only gave it states :a, :b, and :c. However, if we give it the state :d manually, we would get this:

lazy-state-machines
> (take 10 (machine [:a :d]))
("A" "C" "B" "D" "D" "D" "D" "D" "D" "D")
lazy-state-machines
> (first (drop 1000000000 (machine [:a :d])))
"D"

Unless I used take we would’ve stuck on the :d state forever. However, it doesn’t concern us much, as we lazily take the results out of the state machine, and even if we were to process an extreme amount of steps, the stack would not overflow. The memory footprint would also not be that big unless we retained the head of the produced list somewhere.

Now, onto mutual recursion. Here are two more state machines:

(defmulti dog first)
(defmulti kid first)

(defmethod dog :play [[state & states]]
  (cons "dog: runs back with a ball"
        (lazy-seq (kid (cons :throw states)))))

(defmethod kid :throw [[state & states]]
  (cons "kid: throws the ball"
        (lazy-seq (dog (cons :ball states)))))

(defmethod dog :ball [[state & states]]
  (cons "dog: runs for the ball"
        (lazy-seq (dog (cons :play states)))))

(defmethod kid nil [[state & states]]
  (cons "kid: does nothing" (lazy-seq (dog states))))

(defmethod dog nil [[state & states]]
  ["dog: stops playing"])

Hopefully it’s not too complicated.

There are two communicating state machines the dog and the kid. The kid can :throw the ball, and the dog can :play which means that it brings the ball back to the kid asking to throw again. As can be seen, there’s a mutual recursion going on, and there’s no state that will make these machines halt. The kid and the dog are doomed to play til the end of time, and the StackOverflow won’t save them. Thankfully, they’re lazy, so they won’t unless we check on them.

Now, you may be wondering, apart from being able to specify many states upfront, what is the reason for having states as a sequence? Most of the time we’re just consing a new state, and discarding the previous one, so we could just use a singular state instead.

That’s right, but lazy sequences also give us one more advantage - such state machines can influence its future!

Let’s introduce a random chance for the dog to get tired:

(defmethod dog :ball [[state & states]]
  (cons "dog: runs for the ball"
        (lazy-seq (if (> (rand) 0.75)
                    (dog (lazy-cat [:tired] states [:rest]))
                    (dog (cons :play states))))))

(defmethod dog :rest [[state & states]]
  (cons "dog: rests" (lazy-seq (kid states))))

(defmethod dog :tired [states]
  (cons "dog: walks back with a ball"
        (lazy-seq (kid states))))

(defmethod kid :tired [[state & states]]
  (cons "kid: sees dog is tired"
        (lazy-seq (dog states))))

As you can see, we’ve changed the :ball state for the dog. If the random number is higher than 0.75, the dog changes its state to :tired and appends a task for later to :rest.

Then, in the :tired state, the dog doesn’t want to do anything yet, so it looks at the kid with the same :tired state. The kid notices that the dog is tired, and stops playing, allowing dog to check if it wants to do anything else. Earlier the dog appended the :rest task, so it decides to rest, and let’s kid check if it has any more tasks to do.

Since no more states are present, both stop playing:

lazy-state-machines
> (run! println (kid [:throw]))
kid: throws the ball
dog: runs for the ball
dog: runs back with a ball
kid: throws the ball
dog: runs for the ball
dog: walks back with a ball
kid: sees dog is tired
dog: rests
kid: does nothing
dog: stops playing

If the kid is extremely unlucky, this machine can still loop forever though.

We can, however, force the dog to play, by giving the kid two throwing tasks. This would require implementing a :default method for the dog, so if it doesn’t know what to do with the state, it asks the kid to check if it can. The same applies to the kid:

(defmethod kid :default [states]
  (lazy-seq (dog states)))

(defmethod dog :default [states]
  (lazy-seq (kid states)))

Now, we can make the kid more ignorant:

lazy-state-machines
> (run! println (kid [:throw :throw]))
kid: throws the ball
dog: runs for the ball
dog: walks back with a ball
kid: sees dog is tired
kid: throws the ball
dog: runs for the ball
dog: runs back with a ball
kid: throws the ball
dog: runs for the ball
dog: runs back with a ball
kid: throws the ball
dog: runs for the ball
dog: walks back with a ball
kid: sees dog is tired
dog: rests
dog: rests
kid: does nothing
dog: stops playing

The dog needs double rest now though.

Anyway, that’s it! I hope you’ve found this entertaining, and it would be cool to hear from you if you have ever used this kind of lazy state machine in the past, or if you can think of other cool applications for them. Thank you for reading!