Implement discriminated union schema

ambar-core/src/helpers/object.ts

@@ -0,0 +1,43 @@
+export { filterKeys, filterMap, mapValues };
+
+function filterKeys<O extends object, K extends keyof O>(
+  obj: O,
+  pred: <KK extends K>(key: KK, value: O[KK]) => boolean
+): Partial<O> {
+  const result: Partial<O> = {};
+  const keys = Object.keys(obj) as K[];
+  for (const key of keys) {
+    if (pred(key, obj[key])) {
+      result[key] = obj[key];
+    }
+  }
+  return result;
+}
+
+// Filter keys and transform values at the same time
+function filterMap<O extends object, R>(
+  obj: O,
+  fn: <K extends keyof O>(key: K, value: O[K]) => R | undefined
+): Partial<{ [K in keyof O]: R }> {
+  const result = {} as Partial<{ [K in keyof O]: R }>;
+  const keys = Object.keys(obj) as Array<keyof O>;
+  for (const key of keys) {
+    const mapped = fn(key, obj[key]);
+    if (mapped !== undefined) {
+      result[key] = mapped;
+    }
+  }
+  return result;
+}
+
+function mapValues<T extends object, R>(
+  obj: T,
+  fn: <K extends keyof T>(key: K, value: T[K]) => R
+): { [K in keyof T]: R } {
+  const out = {} as { [K in keyof T]: R };
+  const keys = Object.keys(obj) as Array<keyof T>;
+  for (const k of keys) {
+    out[k] = fn(k, obj[k]);
+  }
+  return out;
+}

ambar-core/src/json/schema.ts

@@ -6,6 +6,7 @@ export {
   type Infer,
   type SchemaDef,
   type SchemaOptional,
+  type Variant,
   object,
   pair,
   triple,
@@ -24,6 +25,8 @@ export {
   stringLiteral,
   stringEnum,
   oneOf,
+  discriminatedUnion,
+  variant,
   from,
   decode,
   encode,
@@ -40,6 +43,7 @@ import { Encoder, EncoderDef } from './encoder';
 import { Json } from './types';
 import * as E from './encoder';
 import { Maybe, Nullable } from '../maybe';
+import { filterMap, mapValues } from '../helpers/object';
 
 // Infer the type from a schema definition
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -181,3 +185,74 @@ const oneOf = <V>(f: (v: V) => Schema<V>, ss: Array<Schema<V>>): Schema<V> =>
     D.oneOf(ss.map((s) => s.decoder)),
     E.oneOf((v) => f(v).encoder)
   );
+
+const discriminatedUnion = <const Variants extends readonly Variant<any>[]>(
+  vars: Variants
+): Schema<Infer<Variants[number]['schema']>> => {
+  type Ty = Infer<Variants[number]['schema']>;
+  const d: D.Decoder<Ty> = D.oneOf(vars.map((v) => v.schema.decoder));
+  const e: E.Encoder<Ty> = E.oneOf<Ty>((v) => {
+    const found = vars.find((variant) => matches(variant.pattern, v));
+    if (found == undefined) {
+      throw new Error(`Invalid discriminant in union type: '${v}'`);
+    }
+
+    return found.schema.encoder;
+  });
+
+  return new Schema<Ty>(d, e);
+};
+
+// Check whether a value matches a pattern.
+const matches = (pattern: Record<string, string>, val: unknown): boolean => {
+  if (typeof val !== 'object' || val === null) {
+    return false;
+  }
+
+  for (const key in pattern) {
+    if (!(key in val)) {
+      return false;
+    }
+    if (pattern[key] != undefined && pattern[key] !== (val as any)[key]) {
+      return false;
+    }
+  }
+
+  return true;
+};
+
+// One option in a sum type.
+// Includes the pattern that differentiates it from the other options in the type.
+class Variant<T> {
+  constructor(
+    readonly pattern: Record<string, string>,
+    readonly schema: Schema<T>
+  ) {}
+}
+
+type VariantDef<T extends string, A> = {
+  [D in keyof A]: Schema<A[D]> | SchemaOptional<A[D]> | (A[D] & T);
+};
+
+const variant = <const T extends string, const A>(
+  def: VariantDef<T, A>
+): Variant<A> => {
+  const pattern = filterMap(def, (_, v): string | undefined =>
+    typeof v == 'string' ? v : undefined
+  ) as Record<string, string>;
+
+  if (Object.keys(pattern).length == 0) {
+    throw new Error(
+      'Invalid variant definition. No discriminant identified. Discriminant must be provided as a string'
+    );
+  }
+
+  const schemaDef: SchemaDef<A> = mapValues(def, (_, value) =>
+    // @ts-expect-error hard to prove the types, but this is correct.
+    typeof value == 'string' ? stringLiteral<T>(value) : value
+  );
+
+  const schema: Schema<A> = object(schemaDef);
+
+  return new Variant(pattern, schema);
+};

ambar-core/tests/suites/json.ts

@@ -6,6 +6,7 @@ import * as fc from 'fast-check';
 import { Schema } from 'json/schema';
 import { Json } from 'json/types';
 import { Maybe, Nothing, Just, Nullable } from 'maybe';
+import * as assert from 'node:assert';
 
 const tests = group('json', [
   group('schema', [
@@ -106,6 +107,65 @@ const tests = group('json', [
         expect.deep_equals(encoded, { one: 'wat' });
       }),
     ]),
+    group('discriminatedUnion', [
+      test('roundtrips first variant', () =>
+        roundtrip(
+          s.discriminatedUnion([
+            s.variant({ ty: 'one', a: s.string }),
+            s.variant({ ty: 'two', b: s.number }),
+          ]),
+          { ty: 'one', a: 'wat' }
+        )),
+      test('roundtrips second variant', () =>
+        roundtrip(
+          s.discriminatedUnion([
+            s.variant({ ty: 'one', a: s.string }),
+            s.variant({ ty: 'two', b: s.number }),
+          ]),
+          { ty: 'two', b: 2 }
+        )),
+      test('fails on payload mismatch', () =>
+        expect.throws(
+          () =>
+            roundtrip(
+              s.discriminatedUnion([
+                s.variant({ ty: 'one', a: s.string }),
+                s.variant({ ty: 'two', b: s.number }),
+              ]),
+              { ty: 'one', a: 2 } as any
+            ),
+          (err) => assert.match(err.message, /Unable to decode/)
+        )),
+      test('fails on unknown variant', () =>
+        expect.throws(
+          () =>
+            roundtrip(
+              s.discriminatedUnion([
+                s.variant({ ty: 'one', a: s.string }),
+                s.variant({ ty: 'two', b: s.number }),
+              ]),
+              { c: 2 } as any
+            ),
+          (err) =>
+            assert.match(err.message, /Invalid discriminant in union type/)
+        )),
+      test('fails if discriminant is not specified', () =>
+        expect.throws(
+          () =>
+            roundtrip(
+              s.discriminatedUnion([
+                s.variant({ ty: s.stringLiteral('one'), a: s.string }),
+                s.variant({ ty: 'two', b: s.number }),
+              ]),
+              { ty: 'two', b: 2 }
+            ),
+          (err) =>
+            assert.match(
+              err.message,
+              /Discriminant must be provided as a string/
+            )
+        )),
+    ]),
   ]),
 ]);