How do I make my function specialize on a union tag?

Question

I have union types Item and ItemSlot and a mapped type that maps Items to Slots.

type WeaponSlot = "melee" | "ranged";

type ArmorSlot = "helmet" | "gloves" | "boots";

type Weapon = {
  type: "weapon";
  slot: WeaponSlot;
  damage: number;
};

type Armor = {
  type: "armor";
  slot: ArmorSlot;
  armorValue: number;
};

type Item = Weapon | Armor;

type ItemSlot = WeaponSlot | ArmorSlot;

// type ItemSlots = {
//     melee: Weapon;
//     ranged: Weapon;
//     helmet: Armor;
//     gloves: Armor;
//     boots: Armor;
// }
type ItemSlots = {
  [O in Item as O["slot"]]: O;
};

There is a clear relation between items and item slots, and this is reflected in the type mapping. But I couldn't figure out how to make use of this relationship so I can make generic functions that can get an Item and map it to the correct ItemSlot in a type safe way.

These work

// manual way
function putItem(item: Item, slots: ItemSlots) {
  if (item.slot === "boots") {
    slots["boots"] = item;
  } else if (item.slot === "melee") {
    slots["melee"] = item;
  }
}


// can do generic when slot is passed from outside
function putItem2<S extends ItemSlot>(
  item: ItemSlots[S],
  slots: ItemSlots,
  slot: S,
) {
  slots[slot] = item;
}

But I don't want the manual option. And slot may not get passed from outside like examples below

// a function that takes an item and puts it in its corresponding slot
function putItem3(item: Item, slots: ItemSlots) {
  slots[item.slot] = item;
}

// or a function that takes an item list and slot as an input and
// filters items to get an item suitable for the slot
function putItem4(items: Item[], slot: ItemSlot, slots: ItemSlots) {
  // just return the first item that suits to slot
  const item = items.filter((item): item is Something => item.slot === slot)[0];

  slots[item.slot] = item;
}

Somehow I need to hint that the item has a slot that would point to one of the correct slots. Do I need to make Item generic on Slot? How would that work together with union types? Or is there a better solution maybe that doesn't require the use of more generics?

This doesn't seem like a unique problem. I would also appreciate if you can link some examples in the wild where union types are handled in a generic way.

Answer 1

You're running into an issue I've called "correlated union types", as discussed in microsoft/TypeScript#30581. Sometimes people write code where an unchanging value of a union type is used multiple times, and which is clearly safe if you keep in mind that the value's type must be the same union member in all occurrences. For example, if x is of type number[] | string[], then x.push(x[0]) is safe (assuming x is non-empty) because x is either a number[] both times or it's a string[] both times. But the compiler does not analyze the code this way because it's not tracking the identity of the value, just the type. It treats it the same as if you wrote x1.push(x2[0]) where both x1 and x2 are of type number[] | string[]. And that would clearly be unsafe.

That's the problem you're having with your putItem() code. slots[item.slot] and item must be of the same type, but that's only discoverable if you tract the identity of item and not its type.

---

The recommended approach to correlated unions is a refactoring described at microsoft/TypeScript#47109, where you move away from unions and toward generics that are constrained to unions, and specifically involving a "base" interface and generic indexes into that interface and into mapped types over that interface. The goal is to make it so that the two correlated expressions are seen as being of a visibly-equivalent generic type.

With your example as shown, it could look like this. First the base interface:

type ItemMap = { [I in Item as I["type"]]: I }
/* type ItemMap = {
    weapon: Weapon;
    armor: Armor;
} */

And then the ItemSlots type becomes a generic type of the following form:

type ItemSlots<K extends keyof ItemMap> = Record<ItemMap[K]["slot"], ItemMap[K]>;

You could keep around your original ItemSlots is you want, but the point of this is that ItemSlots<"weapon"> is seen is the Weapon-appropriate piece of ItemSlots, and ItemSlotes<"armor"> is the Armor-appropriate piece.

Then putItem() can be made generic in the following way:

function putItem3<K extends keyof ItemMap>(item: ItemMap[K], slots: ItemSlots<K>) {
  const slot: ItemMap[K]['slot'] = item.slot
  slots[slot] = item;
}

Note that there's a bit of a workaround to microsoft/TypeScript#33181 since when you index into a generic (like ItemMap[K]) with a specific key (like "slot") the compiler widens things to a non-generic. So I needed to annotate that slot is of type ItemMap[K]['slot'], since that is the key of ItemSlots<K>.

Your filter() version is similarly generic:

function putItem4<K extends keyof ItemMap>(
  items: Item[], slot: ItemMap[K]["slot"], slots: ItemSlots<K>
) {
  const item = items.filter((item): item is ItemMap[K] => item.slot === slot)[0];
  slots[slot] = item;
}

---

So that works and is at least vaguely type safe, especially compared to just using type assertions. It would be nice if the compiler could always "see" the safety by distributing its analysis over unions, or tracking identities, but for now microsoft/TypeScript#47109 is the best approach we have for this sort of issue.

Playground link to code