A Pure Implementation of the ST Monad
Tags: functional
Published: 2022-07-10 10:58:06

---
title: A Pure Implementation of the ST Monad
published: true
description: A ST monad with no external dependencies in less than 200 lines of code
tags: functional
cover_image: https://dev-to-uploads.s3.amazonaws.com/uploads/articles/4qkerdz8f0niiwcsjvjb.jpeg
# Use a ratio of 100:42 for best results.
---

I've been part of a couple discussions recently about the nature of `ST` and, more specifically, its relationship to purity and certain monadic laws. Folks including myself have been asking:

- Is it possible to implement `ST` in a "pure" way? By "pure" here, I mean a hand-wavy notion of referentially transparent, no dependencies on external libs or low-level hacks, etc.
- Does `ST` act like a transformer? Can it have an instance of `MonadTrans`?

The answer is yes and yes! In this article, I'll share a no-frills pure implementation of `ST` in one file with no external dependencies. I'll also provide a strategy to give it a valid instance of `MonadTrans`.

## ST - the basics

`ST` is a monad that allows you to write code that feels like JavaScript circa 1995. Everything is a variable, all variables can be modified, and the assignment of a variable to another variable captures the original variable's value at the point of assignment.

```javascript
var foo = 1; // assignment 
var bar = foo; // read + assignment
foo = 3; // write
console.log(bar); // 1
bar = 5; // write
console.log(bar) // 5
```

In Haskell and PureScript, the corresponding code would be:

```haskell
foo <- new 1 -- assignment 
bar <- read foo >>= new -- read + assignment
write 3 foo -- write
log bar -- 1
write 5 bar -- write
log bar -- 5
```

`ST` is nice for several reasons:

- for those who have never programmed in an ML-family language like Haskell or PureScript, `ST` is a great way to translate imperative code into functional code with minimal modifications.
- many algorithms are faster and easier to express using `ST`.
- some types, like those representing event streams, need to create _ad hoc_ buffers of values for situations like debouncing and underflow, and `ST` is a useful way to implement these features (this is how, for example, [`purescript-hyrule`](https://github.com/mikesol/purescript-hyrule) works).

`ST` is almost always implemented by writing its logic in a non-functional language that more closely resembles imperative code. For example, in `PureScript`, almost all of `ST` is implemented in `JavaScript`, and the compiler even has special rules to rewrite `do` blocs of `ST` code as imperative code. This leads to substantial performance gains when working with `ST`.

For understandable reasons, there's a common misconception that `ST` is a sort of imperative "back door" to functional languages. Even though `ST` often has an imperative implementation, it is possible to conceive of it in an entirely "pure" way. That's what this article is about!

## How is this done?

The first thing to do is define a `ST` type.

```haskell
newtype ST r i o a = ST ((Nat /\ i) -> (Nat /\ o /\ a))
```

Let's unpack what the type is saying:

- `r` is a phantom type that will "lock" `ST` computations into a context.
- `i` represents all of the types for which we have made _ad hoc_ values _before_ this computation. For example, in the JS and Haskell/PS examples above, we are using integers, so this would contain `Int`.
- `o` represents the types that are used _after_ this computation. For example, if the computation creates a new reference to a `Boolean` and we have no other references to `Booleans` in our computation, it will add `Boolean` to the cache of types.
- `a` is the actual value in the context, like in any garden-variety monad.

The constructor itself looks a lot like the `State` monad from Haskell and PureScript's MTL libraries. Here's a simplified definition of `State`:

```haskell
type State = s -> Tuple a s
```

The difference between `State` and `ST` is that, instead of an open `s` type, we restrict the variable in `ST` to a cache that is managed automatically by type classes. The `ST` cache will contain all of the references with their current value.

Then, we need some type classes to manage two operations:

- creating a new reference (we'll call this `STCons`)
- modifying a reference (we'll call this `STModify`, and we'll use it to define both `read` and `write` operations)

### STCons

The definition of `STCons` is as follows:

```haskell
class STCons (x :: Type) (i :: Type) (o :: Type) | x i -> o where
  stCons :: Nat -> x -> i -> o
```

It is very similar to the `cons` one would see for any associative list. There is a `Nat` for the index, an `x` type we are storing, the input list `i` and the output list `o`. At both the term level (the definition of `stCons` and the type level (the functional dependency `x i -> o`), we define a function that inserts `x` in `i` to produce `o`. When we call `new`, this is how will update our cached values.

### STModify

The definition of `STModify` is as follows:

```haskell
class STModify (x :: Type) (y :: Type) where
  stModify :: (x -> x) -> Nat -> x -> y -> x /\ y
```

We can only legally modify a value in the cache if its type has already been added via a call to `stCons`. So, instead of tracking and `i` and `o` type, we have a single type `y` representing the cache with this value already present. If when solving the class constraint the compiler finds that the type `x` is not present in `y`, the program will not compile. This is how the correctness of the cache is enforced at the type level.

`stModify` can be used to implement both `read` and `write`. In the case of `read`, we use the identity function as our `x -> x`. In the case of the write function, we stash the modified `y` in our `ST` and return `Unit`.

## Show me the code!

Here's the whole thing. Everything before `-- test` is the implementation, and everything afterwards is a test that throws a lot of stuff at the type (several `new`-s, `read` and `write` in different orders, etc).

{% gist https://gist.github.com/mikesol/43ecacfd5b2f2b62d464bfaf0339606e %}

Because the file is self-contained with no side effects, the test can't actually execute. So here's a bit of impure code we can use to test that our test code in the gist actually works:

```haskell
module Main where

import Prelude

import Effect (Effect)
import Effect.Console (log)
import ST (myColorVehicle)
import Unsafe.Coerce (unsafeCoerce)

main :: Effect Unit
main = do
  log $ unsafeCoerce myColorVehicle
  pure unit
```

When we run this, we see in the console:

```bash
Tuple { value0: Red {}, value1: Car {} }
```

Sure enough, following the control flow of `myColorVehicle`, this is exactly what we'd expect!

## What about MonadTrans

An implementation of `MonadTrans` would require a new type `STT`:

```haskell
newtype STT r i o m a = ST ((Nat /\ i) -> m (Nat /\ o /\ a))
```

This type has the same topology as `StateT` and an instance of `MonadTrans` can be defined in a similar manner, allowing it to be used as a transformer in an MTL stack.