tiny-pattern-ts/src/tagged-union.ts

import type { Exact, UnknownRecord } from "type-fest";

import type {
    HandlerMap,
    Matchable,
    PatternKey,
    PatternParam,
    PatternReturns,
    RedundantFallback,
    UnaryFn,
} from "./matcher-shared.ts";

// A tagged union is discriminated by one property whose values are the tags (a
// `Matchable`). `string` and `number` tags key a handler map directly;
// `boolean`, `null` and `undefined` are admitted too but are not property keys,
// so they go through the `PatternKey` projection. `symbol` has no literal syntax
// to write a handler under, and `bigint` is not a property key.

// The discriminant values of `T` under `K`. `Extract` keeps the finite literal
// tags and leaves a widened `string`/`number` as itself, so an open universe
// keeps an open fallback.
type Tags<T extends object, K extends keyof T> = Extract<T[K], Matchable>;

// The keys of `T` that can act as a discriminant. `getTaggedUnionMatcher<T>()`
// accepts only these, so the factory rejects a key whose values are not tags.
type Discriminated<T extends object> = {
    [K in keyof T]: T[K] extends Matchable ? K : never;
}[keyof T];

// The member(s) of `T` tagged `V`. `Extract` distributes over the union, so a
// duplicated tag maps to a union of members rather than silently dropping one.
// Keyed by `PatternKey`, so `boolean`/`null`/`undefined` tags can key a mapped
// type; `PatternParam` inverts the projection to recover the member.
type MapTaggedUnion<T extends object, K extends keyof T> = {
    [P in PatternKey<Tags<T, K>>]: Extract<T, Record<K, PatternParam<P>>>;
};

type Handlers<T extends object, K extends keyof T, R> = {
    [P in PatternKey<Tags<T, K>>]: UnaryFn<MapTaggedUnion<T, K>[P], R>;
};

// The members `Handled` covers. Mapping over the projected keys keeps every
// index within `MapTaggedUnion`'s keys, and the conditional drops a stray key
// outside `T` so it cannot widen the remainder. The remainder is
// `Exclude<T, …>`, mirroring the primitive-union matcher's
// `Exclude<T, PatternParam<keyof Handled>>`.
type HandledMembers<T extends object, K extends keyof T, Handled> = {
    [P in PatternKey<Tags<T, K>>]: P extends keyof Handled
        ? MapTaggedUnion<T, K>[P]
        : never;
}[PatternKey<Tags<T, K>>];

// The fallback is a *second argument*, not a property of the handler map, so
// its parameter can be the remainder the map left open. See development/library.md.
type Fallback<T extends object, K extends keyof T, Handled, R> = UnaryFn<
    Exclude<T, HandledMembers<T, K, Handled>>,
    R
>;

// A fallback is redundant once the handler map covers every tag of `T`. Folded
// into `Handled`'s own (self-referential) constraint so it is checked *after*
// inference; see the primitive-union matcher for why a conditional in the fallback's
// parameter is evaluated too early.
type MustBePartial<T extends object, K extends keyof T, Handled> =
    PatternKey<Tags<T, K>> extends keyof Handled ? RedundantFallback : unknown;

// TypeScript does not apply the excess-property check to a generic constraint,
// so `Exact` restores it for the generic forms.

// Strict returns: one common `R`. Overload order is load-bearing:
//   #1 Handlers (first) -> the exhaustive form and the autocomplete popup
//   #2 Fallback (last)  -> accepts a partial handler map plus a fallback
interface TaggedUnionMatcherStrict<T extends object, K extends keyof T> {
    <R>(handlers: Handlers<T, K, R>): UnaryFn<T, R>;
    <
        R,
        Handled extends Exact<Partial<Handlers<T, K, R>>, Handled> &
            MustBePartial<T, K, Handled>,
    >(
        handlers: Handled & Partial<Handlers<T, K, R>>,
        fallback: Fallback<T, K, Handled, R>,
    ): UnaryFn<T, R>;
}

// Widened returns: the union of every handler's return type. `P` is inferred
// from the whole handler map, whose closed constraint supplies the
// contextual/autocomplete type.
interface TaggedUnionMatcherWidening<T extends object, K extends keyof T> {
    <P extends Exact<Handlers<T, K, unknown>, P>>(
        handlers: P,
    ): UnaryFn<T, PatternReturns<P>>;
    <
        R,
        Handled extends Exact<Partial<Handlers<T, K, unknown>>, Handled> &
            MustBePartial<T, K, Handled>,
    >(
        handlers: Handled,
        fallback: Fallback<T, K, Handled, R>,
    ): UnaryFn<T, PatternReturns<Handled> | R>;
}

const dispatch =
    (k: PropertyKey) =>
    (handlers: HandlerMap, fallback?: UnaryFn<never, unknown>) =>
    (shape: object): unknown => {
        // `object` carries no index signature, so the read needs the assertion;
        // the factory admits only keys whose values are tags, and the map keys
        // them by `PatternKey`, so the result is narrowed to the map's key space.
        // oxlint-disable-next-line typescript/no-unsafe-type-assertion
        const tag = (shape as UnknownRecord)[k] as string | number;
        return (
            handlers[tag] ??
            fallback ??
            (() => {
                throw new Error(`Unhandled tag: ${String(tag)}`);
            })
        )(
            // oxlint-disable-next-line typescript/no-unsafe-type-assertion
            shape as never,
        );
    };

export const getTaggedUnionMatcher =
    <T extends object>() =>
    <K extends Discriminated<T>>(k: K): TaggedUnionMatcherStrict<T, K> =>
        dispatch(k);

export const getTaggedUnionMatcherW =
    <T extends object>() =>
    <K extends Discriminated<T>>(k: K): TaggedUnionMatcherWidening<T, K> =>
        dispatch(k);