union codecs, error messages

2 files changed, 113 insertions, 6 deletions

diff --git a/src/util/codec-test.ts b/src/util/codec-test.ts
index d7edd545f..0d1ab5603 100644
--- a/src/util/codec-test.ts
+++ b/src/util/codec-test.ts
@@ -19,14 +19,34 @@
  */
 
 import test from "ava";
-import { stringCodec, objectCodec } from "./codec";
+import {
+  stringCodec,
+  objectCodec,
+  unionCodec,
+  Codec,
+  stringConstCodec,
+} from "./codec";
 
 interface MyObj {
   foo: string;
 }
 
-test("basic codec", (t) => {
-  const myObjCodec = objectCodec<MyObj>().property("foo", stringCodec).build("MyObj");
+interface AltOne {
+  type: "one";
+  foo: string;
+}
+
+interface AltTwo {
+  type: "two";
+  bar: string;
+}
+
+type MyUnion = AltOne | AltTwo;
+
+test("basic codec", t => {
+  const myObjCodec = objectCodec<MyObj>()
+    .property("foo", stringCodec)
+    .build<MyObj>("MyObj");
   const res = myObjCodec.decode({ foo: "hello" });
   t.assert(res.foo === "hello");
 
@@ -34,3 +54,24 @@ test("basic codec", (t) => {
     const res2 = myObjCodec.decode({ foo: 123 });
   });
 });
+
+test("union", t => {
+  const altOneCodec: Codec<AltOne> = objectCodec<AltOne>()
+    .property("type", stringConstCodec("one"))
+    .property("foo", stringCodec)
+    .build("AltOne");
+  const altTwoCodec: Codec<AltTwo> = objectCodec<AltTwo>()
+    .property("type", stringConstCodec("two"))
+    .property("bar", stringCodec)
+    .build("AltTwo");
+  const myUnionCodec: Codec<MyUnion> = unionCodec<MyUnion, "type">("type")
+    .alternative("one", altOneCodec)
+    .alternative("two", altTwoCodec)
+    .build<MyUnion>("MyUnion");
+
+  const res = myUnionCodec.decode({ type: "one", foo: "bla" });
+  t.is(res.type, "one");
+  if (res.type == "one") {
+    t.is(res.foo, "bla");
+  }
+});
diff --git a/src/util/codec.ts b/src/util/codec.ts
index 690486b7d..78516183c 100644
--- a/src/util/codec.ts
+++ b/src/util/codec.ts
@@ -69,6 +69,11 @@ interface Prop {
   codec: Codec<any>;
 }
 
+interface Alternative {
+  tagValue: any;
+  codec: Codec<any>;
+}
+
 class ObjectCodecBuilder<T, TC> {
   private propList: Prop[] = [];
 
@@ -89,10 +94,10 @@ class ObjectCodecBuilder<T, TC> {
    * @param objectDisplayName name of the object that this codec operates on,
    *   used in error messages.
    */
-  build(objectDisplayName: string): Codec<TC> {
+  build<R extends (TC & T)>(objectDisplayName: string): Codec<R> {
     const propList = this.propList;
     return {
-      decode(x: any, c?: Context): TC {
+      decode(x: any, c?: Context): R {
         if (!c) {
           c = {
             path: [`(${objectDisplayName})`],
@@ -107,12 +112,53 @@ class ObjectCodecBuilder<T, TC> {
           );
           obj[prop.name] = propVal;
         }
-        return obj as TC;
+        return obj as R;
       },
     };
   }
 }
 
+class UnionCodecBuilder<T, D extends keyof T, TC> {
+  private alternatives = new Map<any, Alternative>();
+
+  constructor(private discriminator: D) {}
+
+  /**
+   * Define a property for the object.
+   */
+  alternative<V>(
+    tagValue: T[D],
+    codec: Codec<V>,
+  ): UnionCodecBuilder<T, D, TC | V> {
+    this.alternatives.set(tagValue, { codec, tagValue });
+    return this as any;
+  }
+
+  /**
+   * Return the built codec.
+   *
+   * @param objectDisplayName name of the object that this codec operates on,
+   *   used in error messages.
+   */
+  build<R extends TC>(objectDisplayName: string): Codec<R> {
+    const alternatives = this.alternatives;
+    const discriminator = this.discriminator;
+    return {
+      decode(x: any, c?: Context): R {
+        const d = x[discriminator];
+        if (d === undefined) {
+          throw new DecodingError(`expected tag for ${objectDisplayName} at ${renderContext(c)}.${discriminator}`);
+        }
+        const alt = alternatives.get(d);
+        if (!alt) {
+          throw new DecodingError(`unknown tag for ${objectDisplayName} ${d} at ${renderContext(c)}.${discriminator}`);
+        }
+        return alt.codec.decode(x);
+      }
+    };
+  }
+}
+
 /**
  * Return a codec for a value that must be a string.
  */
@@ -126,6 +172,20 @@ export const stringCodec: Codec<string> = {
 };
 
 /**
+ * Return a codec for a value that must be a string.
+ */
+export function stringConstCodec<V extends string>(s: V): Codec<V> {
+  return {
+    decode(x: any, c?: Context): V {
+      if (x === s) {
+        return x;
+      }
+      throw new DecodingError(`expected string constant "${s}" at ${renderContext(c)}`);
+    }
+  }
+};
+
+/**
  * Return a codec for a value that must be a number.
  */
 export const numberCodec: Codec<number> = {
@@ -179,3 +239,9 @@ export function mapCodec<T>(innerCodec: Codec<T>): Codec<{ [x: string]: T }> {
 export function objectCodec<T>(): ObjectCodecBuilder<T, {}> {
   return new ObjectCodecBuilder<T, {}>();
 }
+
+export function unionCodec<T, D extends keyof T>(
+  discriminator: D,
+): UnionCodecBuilder<T, D, never> {
+  return new UnionCodecBuilder<T, D, never>(discriminator);
+}