Using discriminate unions with Map objects

Question

Consider this working typescript example that uses discriminated unions:

type Alpha = { name: "alpha"; alpha: 1}
type Beta = { name: "beta"; beta: 2}
type Letters = Alpha | Beta

const calculateAlpha = (cfg: Alpha): number => {
    return cfg.alpha
}
const calculateBeta = (cfg: Beta): number => {
    return cfg.beta
}

const CalculatorFactory1 = (cfg: Letters) => {
  switch (cfg.name) {
    case "alpha":
      return () => calculateAlpha(cfg)
    case "beta":
      return () => calculateBeta(cfg)
    default:
      throw new Error()
  }
}

I want to refactor this to make it more configurable by moving the switch statement into a map object. Something like this:

type LetterMap = Map<Letters["name"], (c: Letters) => number>
const letterMap: LetterMap = new Map([
    ["alpha", calculateAlpha],
    ["beta", calculateBeta]
])

const CalculatorFactory2 = (cfg: Letters) => {
    const func = letterMap.get(cfg.name)
    return func(cfg)        
}

However I can't figure out how to type letterMap, in this example I get the following error:

Property 'alpha' is missing in type 'Beta' but required in type 'Alpha'.

Typescript Playground Example

Answer 1

First, TypeScript doesn't by default recognize that a Map could be strongly typed so that each key corresponds to a different value type. A Map<K, V> is like Record<K, V> in that it treats all values as the same type. There are ways to try to encode different value types for Map in the type system (see Typescript: How can I make entries in an ES6 Map based on an object key/value type) but it's much, much easier to just use a plain object, since object types naturally represent such key-value type pairs. So let's make letterMap as a plain object like

const letterMap = {
    alpha: cfg => cfg.alpha,
    beta: cfg => cfg.beta
}

and work on how to use it the way you'd like.

---

TypeScript doesn't really support the concept of "correlated unions". If you have a value cfg of a discriminated union type Letters, and look up cfg.name in your letterMap map you will get a union of functions. And then if you try to call that union of functions by passing it cfg as an argument, the compiler will be unhappy. It doesn't understand that the union of functions and the union of arguments are correlated to each other such that the function and argument are appropriate for each other. Instead it treats them as independent unions, and therefore complains that you might be calling calculateAlpha with an Beta argument, or calculateBeta with an Alpha argument. See microsoft/TypeScript#30581 for more information.

The recommended approach to deal with correlation unions, as described in microsoft/TypeScript#47109, is to refactor to replace the union with generics. The technique involves encoding the underlying relationship you need the compiler to track as a base interface, and then writing operations on this type as generic indexes into this interface or mapped types over that interface.

In your case, the base interface looks like this:

interface LetterMapping {
    alpha: { alpha: 1 },
    beta: { beta: 2 }
}

And then we can redefine your Letters type as a mapped type over it:

type Letter<K extends keyof LetterMapping = keyof LetterMapping> =
    { [P in K]: { name: P } & LetterMapping[P] }[K]

This is a distributive object type, so Letter<"alpha"> and Letter<"beta"> correspond to your original Alpha and Beta types:

type Alpha = Letter<"alpha">;
type Beta = Letter<"beta">;

and Letter<"alpha" | "beta"> corresponds to your original Letters discriminated union type (and we can just write Letter for that, since "alpha" | "beta" is the default type argument).

Now letterMap's type needs to be explicitly annotated as another mapped type over LetterMapping:

const letterMap: { [K in keyof LetterMapping]:
    (cfg: Letter<K>) => number } = {
    alpha: cfg => cfg.alpha,
    beta: cfg => cfg.beta
}

And now we can write calculatorFactory as a generic function like this:

const calculatorFactory = <K extends keyof LetterMapping>(
    cfg: Letter<K>
) => letterMap[cfg.name](cfg); // okay

Here the compiler sees cfg.name as having the non-union generic type K. And because letterMap's annotated type is a mapped type over the base interface, when we look up letterMap[cfg.name] we get the non-union generic type (cfg: Letter<K>) => number. And this function expects an argument of type Letter<K>, which is exactly the type of cfg. So letterMap[cfg.name](cfg) compiles with no error, and is seen to be of type number, as desired.

And callers will also see reasonable behavior:

calculatorFactory({ name: "alpha", alpha: 1 });
/* <"alpha">(cfg: { name: "alpha"; } & { alpha: 1; }) => number */
calculatorFactory({ name: "beta", beta: 2 });
/* <"beta">(cfg: { name: "beta"; } & { beta: 2; }) => number */

Here the compiler is able to infer that K is "alpha" for the first call and "beta" for the second. They both return number, so in some sense it doesn't matter, but you could imagine changing the functions so that the return type was also looked up in a base interface:

interface LetterReturn {
    alpha: boolean,
    beta: string
}

const letterMap: { [K in keyof LetterMapping]:
    (cfg: Letter<K>) => LetterReturn[K] } = {
    alpha: cfg => Math.random() < (cfg.alpha / 2),
    beta: cfg => cfg.beta.toFixed(2)
}

const a = calculatorFactory({ name: "alpha", alpha: 1 });
/* const a: boolean */
const b = calculatorFactory({ name: "beta", beta: 2 });
/* const b: string */

which is nice.

Playground link to code