“Pattern matching” with Typescript

lib.test.ts

// Inspired by: https://medium.com/@fillopeter/pattern-matching-with-typescript-done-right-94049ddd671c

import { match, None } from "./lib";

const A = Symbol("A");
interface A {
  readonly _tag: typeof A;
  num: number;
}

const B = Symbol("B");
interface B {
  readonly _tag: typeof B;
  str: string;
}

type Set = A | B;

describe("match", () => {
  () => {
    // it typechecks

    const m = match<Set>()({
      [A]: (v) => 1, // typeof v === A
      [B]: (v) => "lkj", // typeof v === B
      [None]: (v) => true, // typeof v === unknown
    })({ _tag: A, num: 1 });

    type test = typeof m; // string | number | boolean

    const m2 = match<Set>()({
      [A]: (v) => 1, // typeof v === A
      [B]: (v) => "lkj", // typeof v === B
      [None]: (v) => true as unknown, // typeof v === unknown
    })({ _tag: A, num: 1 });
    type test2 = typeof m2; // unknown, because returntype of None handler

    const m3 = match<Set>()({
      [A]: (v) => 1, // typeof v === A
      [B]: (v): any => "lkj", // typeof v === B
      [None]: (v) => true, // typeof v === unknown
    })({ _tag: A, num: 1 });
    type test3 = typeof m3; // any, because returntype of B handler

    // @ts-expect-error missing None branch
    const mError1 = match<Set>()({
      [A]: (v) => 1, // typeof v === A
      [B]: (v) => "lkj", // typeof v === B
    })({ _tag: A, num: 1 });

    // @ts-expect-error missing A branch
    const mError2 = match<Set>()({
      [B]: (v) => "lkj", // typeof v === B
      [None]: (v) => true, // typeof v === unknown
    })({ _tag: A, num: 1 });

    const mError3 = match<Set>()({
      // @ts-expect-error A handler's v mistyped
      [A]: (v: string) => 1, // typeof v === A
      [B]: (v) => "lkj", // typeof v === B
      [None]: (v) => true, // typeof v === unknown
    })({ _tag: A, num: 1 });
  };

  it("matches", () => {
    expect(
      match<Set>()({
        [A]: (v) => v.num * 2,
        [B]: (v) => "Hello " + v.str,
        [None]: (v) => "something broke",
      })({ _tag: A, num: 2 })
    ).toBe(4);

    expect(
      match<Set>()({
        [A]: (v) => v.num * 2,
        [B]: (v) => "Hello " + v.str,
        [None]: (v) => "something broke",
      })({ _tag: B, str: "world" })
    ).toBe("Hello world");

    // invalid: no tag
    expect(
      match<Set>()({
        [A]: (v) => v.num * 2,
        [B]: (v) => "Hello " + v.str,
        [None]: (v) => "something broke",
      })({ num: 2 } as A)
    ).toBe("something broke");

    // invalid: wrong tag
    expect(
        match<Set>()({
          [A]: (v) => v.num * 2,
          [B]: (v) => "Hello " + v.str,
          [None]: (v) => "something broke",
        })({ _tag: Symbol("A"), num: 2 } as A)
      ).toBe("something broke");
  });
});

lib.ts

export interface Tagged {
  readonly _tag: symbol;
}

type TagMap<T extends Tagged> = {
  [K in T["_tag"]]: T extends { _tag: K } ? T : never;
};

export const None = Symbol("None");

export type Pattern<S extends Tagged> = {
  [K in keyof TagMap<S>]: (v: TagMap<S>[K]) => any;
} & { [None]: (v: unknown) => any };

export function match<S extends Tagged>() {
  return <
    P extends Pattern<S>,
    R extends {
      [K in keyof P]: ReturnType<P[K]>;
    },
    R2 extends R[keyof R]
  >(
    p: P
  ) => <A extends S>(v: A): R2 => {
    if("_tag" in v && v._tag in p) {
      return p[v._tag](v);
    }
    return p[None](v);
  };
}