Tagged Unions in TypeScript

In my last post we looked at tagged unions and how they work at a low level in languages like C or Rust. Today we try to get as close as possible to the Rust version using TypeScript.

Here the Rust code again as a reminder, and spoiler alert, the TypeScript version won’t be as pretty.

pub enum Event {
    KeyPress(i32),
    MouseClick(i32, i32),
    MouseMove(i32, i32),
}

pub fn process_event(evt: &Event) -> i32 {
    match *evt {
        Event::KeyPress(key_code) => key_code * 1234,
        Event::MouseClick(x_pos, y_pos) => x_pos + y_pos * 457,
        Event::MouseMove(x_delta, y_delta) => x_delta + y_delta * 987,
    }
}

Mapping the example above to TypeScript directly won’t really work. TypeScript also has an enum type but they don’t allow you to associate data with the enum key. You can however use Union Types to combine multiple different types into one. An example might look like this:

type MyType = string | number

Easy enough. We do still need a way to check at runtime what type we are talking about (same as in Rust or the C version from the previous post). We probably want to handle strings and numbers differently. Something like this would work for the example above.

const values: MyType[] = ["Hi", 100]

for (let i = 0; i < values.length; i++) {
    const v = values[i]
    if (typeof v === "string") {
        console.log("handle string:", v)
    } else {
        console.log("handle number:", v)
    }
}

This approach can however get complicated quickly. Imagine using more complicated object types instead of strings and numbers. We could still determine the type of the object by checking the existence of certain fields but that seems kind of error prone. Especially if some objects have fields with the same name and so on. Instead we can just use the same approach Rust does automatically and what we implemented manually using C. A tag. Something like this.

enum EventTag {
    KeyPress,
    MouseClick,
    MouseMove,
}

type KeyPressEvent = {
    tag: EventTag.KeyPress
    keyCode: number
}

type MouseClickEvent = {
    tag: EventTag.MouseClick
    xPos: number
    yPos: number
}

type MouseMoveEvent = {
    tag: EventTag.MouseMove
    xDelta: number
    yDelta: number
}

type UserEvent = KeyPressEvent | MouseClickEvent | MouseMoveEvent

// the math in here is just an example of doing some work
function processEvent(evt: UserEvent) {
    switch (evt.tag) {
        case EventTag.KeyPress:
            return evt.keyCode * 1234
        case EventTag.MouseClick:
            return evt.xPos + evt.yPos * 457
        case EventTag.MouseMove:
            return evt.xDelta + evt.yDelta * 987
        default:
            missingCase(evt)
    }
}

function missingCase(_value: never): never {
    throw new Error("Not all cases where handled!")
}

Two interesting things you might have noticed. We are not using EventTag as the type for the tags, instead we are using one of the enum values as the type. This is the key to the tagged union approach in TypeScript. This allows us to get good support from the TypeScript compiler. Also noteworthy is the missingCase function. Calling this function in this way in the default case allows the compiler to detect if we forgot to handle any case (e.g. if we add a new type to UserEvent). See never for more infos on the never type.

OK, but how does this “good support” from the compiler look like in practice. Just like the Rust version, the compiler knowns the exact type of evt inside each case block. This means we can access the evt fields in a type-safe way. We also can’t create events that don’t exist (at least not without type casting).

And this is how you might use the example types defined above.

function main() {
    const events: Event[] = [
        { tag: EventTag.KeyPress, keyCode: 100 },
        { tag: EventTag.MouseClick, xPos: 10, yPos: 20 },
        { tag: EventTag.MouseMove, xDelta: 3, yDelta: 1 },
    ]

    for (let i = 0; i < events.length; i++) {
        let result = processEvent(events[i])
        console.log("Event Result:", result)
    }
}

main()

Practical example

Just in case it is helpful, another example based on a real world project. Imagine an API which returns you a list of configurations for UI components that you want to render in your web app. This could be the API from your CMS or some other dynamic data.

The following is a simple react example on how you might structure your code to render a dynamic list of UI components in an almost type-safe way. Almost because we still have any in TypeScript. If the API gives us an incorrect component structure and we don’t validate the data, the TypeScript compiler cannot help us.

enum UITag {
    Markdown,
    ImageSlideshow,
    // ...
}

type Markdown = {
    tag: UITag.Markdown
    markdown: string
}

type ImageSlideshow = {
    tag: UITag.ImageSlideshow
    images: string[]
}

type ComponentData = Markdown | ImageSlideshow // | ...

const UIComponent = ({ data }: { data: ComponentData }) => {
    switch (data.tag) {
        case UITag.Markdown:
            return <MyMarkdownComponent markdown={data.markdown} />
        case UITag.ImageSlideshow:
            return <MySlideshowComponent images={data.images} />
        default:
            missingCase(data)
    }
}

const UIComponentList = ({ components }: { components: ComponentData[] }) => {
    return (
        <div>
            {components.map((component, idx) => (
                <UIComponent key={idx} data={component} />
            ))}
        </div>
    )
}

Not too much to say here. Same as the previous example just using UI components in React. You can imagine how we could fetch a list of UI component configurations from our CMS and then pass that data to the UIComponentList to render all the components.

Closing thoughts

I do realise that comparing TypeScript with something like Rust and C doesn’t make too much sense as those languages are quite different. But that wasn’t really the point. I just wanted to give a practical introduction into tagged unions and how they can be used in some popular languages. I do think that tagged unions in general are a very useful tool in your programmers toolbox.

In TypeScript we don’t really have control of things like memory layout or fine grained memory management. But even ignoring that for a second I do still think that from a code maintainability perspective tagged unions are pretty useful even in languages like TypeScript. They usually result in a structure which is quite easy to read and modify.

I hope this post was useful to you and maybe you learned something.