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

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

import type {
    HandlerMap,
    Matchable,
    Member,
    PatternKey,
    PatternReturns,
    RedundantFallback,
    UnaryFn,
    UniverseGate,
} 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; a broad tag is then
// rejected by the universe gate.
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` narrowed to the tag(s) `V`. Distributes over `T`, so a
// duplicated tag maps to a union of members rather than silently dropping one.
// A member whose `K` cannot take `V` drops out; the rest have `K` narrowed to
// `Extract<T[K], V>`. `T[K] extends V` keeps an exact member (a discriminated
// union's declared interface) untouched, so only a property that is itself a
// union — or optional — goes through the mapped form. `SetRequired` makes `K`
// required unless `V` admits `undefined`: a defined tag yields `{ type: "x" }`,
// while the `undefined` tag keeps `{ type?: never }` (absence) or
// `{ type?: undefined }` when the property admits an explicit `undefined`.
// Keyed by `PatternKey`, so `boolean`/`null`/`undefined` tags can key a mapped
// type; `Member` inverts the projection against the tag set to recover the tag.
type Narrowed<T extends object, K extends keyof T, V> = T extends object
    ? [Extract<T[K], V>] extends [never]
        ? never
        : T[K] extends V
          ? T
          : undefined extends V
            ? { [P in keyof T]: P extends K ? Extract<T[P], V> : T[P] }
            : SetRequired<
                  { [P in keyof T]: P extends K ? Extract<T[P], V> : T[P] },
                  K
              >
    : never;

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

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

// The tag values whose `PatternKey` is handled.
type HandledTags<T extends object, K extends keyof T, Handled> = Member<
    Tags<T, K>,
    keyof Handled
>;

// The fallback is a *second argument*, not a property of the handler map, so
// its parameter can be the remainder the map left uncovered: the members
// narrowed to the tags the map did not handle. See development/library.md.
type Fallback<T extends object, K extends keyof T, Handled, R> = UnaryFn<
    Narrowed<T, K, Exclude<Tags<T, K>, HandledTags<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> & UniverseGate<Tags<T, K>>): UnaryFn<T, R>;
    <
        R,
        Handled extends Exact<Partial<Handlers<T, K, R>>, Handled> &
            MustBePartial<T, K, Handled>,
    >(
        handlers: Handled &
            Partial<Handlers<T, K, R>> &
            UniverseGate<Tags<T, K>>,
        fallback: Fallback<T, K, Handled, R> & UniverseGate<Tags<T, K>>,
    ): 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 & UniverseGate<Tags<T, K>>,
    ): UnaryFn<T, PatternReturns<P>>;
    <
        R,
        Handled extends Exact<Partial<Handlers<T, K, unknown>>, Handled> &
            MustBePartial<T, K, Handled>,
    >(
        handlers: Handled & UniverseGate<Tags<T, K>>,
        fallback: Fallback<T, K, Handled, R> & UniverseGate<Tags<T, K>>,
    ): UnaryFn<T, PatternReturns<Handled> | R>;
}

// The key-taking step of the curried factory. Naming it lets the factory return
// `dispatch` directly, the tacit twin of the primitive-union factory's bare
// `=> dispatch`.
type TaggedUnionMatcherFactory<T extends object> = <K extends Discriminated<T>>(
    k: K,
) => TaggedUnionMatcherStrict<T, K>;

type TaggedUnionMatcherWideningFactory<T extends object> = <
    K extends Discriminated<T>,
>(
    k: K,
) => TaggedUnionMatcherWidening<T, K>;

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,
>(): TaggedUnionMatcherFactory<T> => dispatch;

export const getTaggedUnionMatcherW = <
    T extends object,
>(): TaggedUnionMatcherWideningFactory<T> => dispatch;