discriminate.ts

/**
 * Exhaustive visitation for discriminated unions.
 * Let's say you have something like
 *
 *     type T = { type: "number", value: number } | { type: "string", value: string };
 *     const t: T = ...;
 *
 * Then you can say
 *
 *     discriminate(t, "type", {
 *         number: ({ value }) => ..., // `value` is a `number` here!
 *         string: ({ value }) => ..., // and a `string` here!
 *     })
 */
export function discriminate<
    Union extends Discriminated<Discriminator>,
    Discriminator extends Discriminators<Union>,
    Arms extends DiscriminateArms<Union, Discriminator>,
>(
    value: Union,
    discriminator: Discriminator,
    arms: Restrict<Arms, DiscriminateArms<Union, Discriminator>>,
): DiscriminateResult<Union, Discriminator, Arms>;
export function discriminate<
    Union extends Discriminated<Discriminator>,
    Discriminator extends Discriminators<Union>,
    Arms extends Partial<DiscriminateArms<Union, Discriminator>>,
    Fallback,
>(
    value: Union,
    discriminator: Discriminator,
    arms: Restrict<Arms, Partial<DiscriminateArms<Union, Discriminator>>>,
    fallback: (value: Union) => Fallback,
): DiscriminateResult<Union, Discriminator, Arms> | Fallback;
export function discriminate<
    Union extends Discriminated<Discriminator>,
    Discriminator extends Discriminators<Union>,
    Arms extends Partial<DiscriminateArms<Union, Discriminator>>,
    Fallback,
>(
    value: Union,
    discriminator: Discriminator,
    arms: Restrict<Arms, Partial<DiscriminateArms<Union, Discriminator>>>,
    fallback?: (value: Union) => Fallback,
): DiscriminateResult<Union, Discriminator, Arms> | Fallback {
    // Implementing this so that TS accepts it is hard to impossible.
    // The code is simple enough, though.
    // See also `match` above.
    type Result = DiscriminateResult<Union, Discriminator, Arms>;
    const arm = arms[value[discriminator]] as (
        <Variant extends Union>(value: Variant) => Result
    ) | undefined;
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    return (arm ?? fallback)!(value);
}

type Discriminated<Discriminator extends string | number> = {
    [Discriminant in Discriminator]: string | number;
};

type Discriminators<Union> = keyof Union & (string | number);

type DiscriminateArms<
    Union extends Discriminated<Discriminator>,
    Discriminator extends Discriminators<Union>,
    Out = unknown,
> = {
    [Discriminant in Union[Discriminator]]: (
        variant: Extract<Union, Record<Discriminator, Discriminant>>,
    ) => Out;
};

type Restrict<T, U> = Exclude<keyof T, keyof U> extends never ? T : U;

type DiscriminateResult<
    Union extends Discriminated<Discriminator>,
    Discriminator extends Discriminators<Union>,
    Arms extends Partial<DiscriminateArms<Union, Discriminator>>,
> = Arms extends Partial<DiscriminateArms<Union, Discriminator, infer Out>>
    ? Out
    : never;