How can I narrow a union type based on a key, value and a shape provided to a function?

Question

I would like to discriminate a union type based on args that provided to a function, but for some reason I can't use a generic type for a shape of data. It brokes my narrowing. What do you think how can I achieve this?

export type DiscriminateUnionType<Map, Tag extends keyof Map, TagValue extends Map[Tag]> = Map extends Record<
  Tag,
  TagValue
>
  ? Map
  : never;

function inStateOfType<Map extends { [index in Tag]: TagValue }, Tag extends keyof Map, TagValue extends Map[Tag], DiscriminatedState extends DiscriminateUnionType<Map, Tag, TagValue>>(tag: Tag, value: TagValue, state: Map): DiscriminatedState | undefined {
  return state[tag] === value ? state as DiscriminatedState : undefined
}

type State = { type: 'loading', a: string } | { type: 'loaded', b: string } | { type: 'someOtherState', c: string }

export function main(state: State) {
  const loadedState = inStateOfType('type', 'loading', state)

  if (loadedState) {
    loadedState.b // Property 'b' does not exist on type 'State'. Property 'b' does not exist on type '{ type: "loading"; a: string; }'
  }
}

function inStateOfType<Map extends { type: string }, Tag extends 'type', TagValue extends Map[Tag], DiscriminatedState extends DiscriminateUnionType<Map, Tag, TagValue>>(state: Map, value: TagValue): DiscriminatedState | undefined {
  return state['type'] === value ? state as DiscriminatedState : undefined
}

function main(state: State) {
  // { type: "loaded"; b: string }, everything is fine, narrowing works
  // but in this case, inStateOfType function is not generic
  const loadedState = inStateOfType(state, 'loaded')

  if (loadedState) {
    loadedState.b 
  }
}

In order to investigate this I created an executable snippet with a code, so you can debug it on TS playground

Answer 1

In what follows I am going to change the names of your type parameters to be more in line with TypeScript conventions (single uppercase characters); Map will become M, Tag will become K (as it is a key of M), TagValue will become V, index will become I, and DiscriminatedState will become S. So now we have:

function inStateOfType<
  M extends { [I in K]: V },
  K extends keyof M,
  V extends M[K], 
  S extends Extract<M, Record<K, V>>
>(tag: K, value: V, state: M): S | undefined {
  return state[tag] === value ? state as S : undefined
}

And note that { [I in K]: V } is equivalent to Record<K, V> using the Record<K, V> utility type and that

type DiscriminateUnionType<M, K extends keyof M, V extends M[K]> =
  M extends Record<K, V> ? M : never;

can be dispensed with in favor of the built-in Extract<T, U> utility type as Extract<M, Record<K, V>>, so now we have:

function inStateOfType<
  M extends Record<K, V>,
  K extends keyof M,
  V extends M[K], S extends Extract<M, Record<K, V>>
>(tag: K, value: V, state: M): S | undefined {
  return state[tag] === value ? state as S : undefined
}

We're almost done cleaning this up to the point where we can answer. One more thing; the S type parameter is superfluous. There is no good inference site for it (no parameter is of type S or a function of S) so the compiler will just fall back to having S be exactly Extract<M, Record<K, V>>, meaning it's just a synonym for it.

And if you're going to write return xxx ? yyy as S : undefined then you don't need to annotate the return type at all, since it will be inferred as S | undefined.

So you could write the following and have everything work (or fail to work) the same:

function inStateOfType<
  M extends Record<K, V>,
  K extends keyof M,
  V extends M[K]
>(tag: K, value: V, state: M) {
  return state[tag] === value ?
    state as Extract<M, Record<K, V>> :
    undefined
}

---

So why doesn't that work? The big problem here is that M is supposed to be the full discriminated union type, so you can't constrain it to Record<K, V>, since V is just one of the various possible values for the key K. If you constrain M to Record<K, V>, then the compiler will not let you pass in a value for state unless it already knows that its tag property is the same type as value. Or, as in your case, the compiler will widen V so that it is the full set of possibilities for tag. Oops.

So if we can't constrain M to Record<K, V>, what should we constrain it to? It needs a key at K, but the value type there should only be constrained to be a viable discriminant property. Something like

type DiscriminantValues = string | number | boolean | null | undefined;

Let's try it:

function inStateOfGenericType<
  M extends Record<K, DiscriminantValues>,
  K extends keyof M,
  V extends M[K]
>(tag: K, value: V, state: M) {
  return state[tag] === value ?
    state as Extract<M, Record<K, V>> :
    undefined
}

function main(state: State) {
  const loadedState = inStateOfGenericType('type', 'loaded', state)

  if (loadedState) {
    loadedState.b // okay
  }
}

And that does it!

---

Do note that in TypeScript it is a little more conventional to rewrite this as a user defined type guard function where inStateOfType() returns a boolean that can be used to decide whether the compiler may narrow state to Record<K, V> or not:

function inStateOfGenericType<
  M extends Record<K, DiscriminantValues>,
  K extends keyof M,
  V extends M[K] 
>(tag: K, value: V, state: M):
  state is Extract<M, Record<K, V>> {
  return state[tag] === value
}

function main(state: State) {
  if (inStateOfGenericType('type', 'loaded', state)) {
    state.b // okay
  }
}

Playground link to code