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

Answer 2

You try to sort of do a reverse discrimination, which really isn't possible. This is because your TagValue extends Map[Tag], where Tag = type which unfortunately will always produce a union 'loaded' | 'loading' | 'someOtherValue'. When you extend Map[Tag] it is a sort of narrowing in of itself, so TS does not narrow it further, since 'loaded' will fulfill the union it'll pass there, but then goes on to pass on the whole union to the DiscriminatedState.

So when it tries to use TagValue in DiscriminateUnionType<>, it'll pass in the entire union and not really narrow it at all.

Instead of having TagValue extends Map[Tag] you should instead allow any string, then conditionally check it later. That way your narrowing occurs outside of the generic, where TS might infer and narrow it incorrectly.

function inStateOfGenericTypeFix<
  Map extends Record<string, any>, 
  Tag extends keyof Map, 
  TagValue extends string = Map[Tag], 
> (tag: Tag, value: TagValue, state: Map): 
  TagValue extends Map[Tag] 
  ? DiscriminateUnionType<Map, Tag, typeof value> 
  : unknown
  | undefined 
{
  return state[tag] === value ? state as any : undefined
}

  const loadedState = inStateOfGenericTypeFix('type', 'loaded', state)
  //  ^?

  if (loadedState) {
    loadedState.b //No more error!
  }

This unfortunately means you can pass in any string to the value, but if the string isn't a valid key, it will return unknown, basically letting you know the intended use was wrong. You can also configure this to return undefined or something else as you please.

Here is a lot of examples/the code on Playground

This has been my personal method of doing this sort of type inferencing, and if anyone else has a better method I'd love to hear it.