Allow union type instantiation to be faster by specifying a discriminator property

Gadget side, we instantiate a lot of unions that have many elements. Both MQT (or MST) have to take an incoming snapshot and figure out which element of the union is the right one to instantiate and return for the union. Both of them have a pretty naive algorithm for doing this: take each element and check the snapshot against it, returning the first union element which does. This means a LOT of `type.is` checks, where each type in the union is compared against the snapshot. These type.is checks basically do a full instantiation under the hood -- there's no fancy fast path for `.is` for a lot of types, there's just running the runtime code that would boot an instance, and seeing if it fails.

This works in the general case but has gotten really slow in Gadget, specifically for the FieldConfig enum. It has a lot of elements, and each element's type is complicated-ish such that it takes a little bit of time to check with `.is`.

Our snapshots for this type have a really easy to get at `"type": "StringConfig"` property on them though that allows us to discriminate between the union elements super easily. We added this a long time ago for typescript discriminated unions at type type, but, I think we can use this at runtime too for a much faster path for instantiation. MST already has support for a custom union type "dispatcher", that returns the right type to use for instantiation given a snapshot. This PR adds support for that, and then adds a macro of sorts on top that does type dispatching using a discriminating property. We expect all the types in the union to have a value for this property we can inspect at the time the union is defined, and then we build a map from the value of the discriminator to the type. When a snapshot comes in, we look up the value of the snapshot's discriminator property within the map.

Author: airhorns

spec/simple.spec.ts

@@ -1,4 +1,4 @@
-import { getSnapshot, types } from "../src";
+import { types } from "../src";
 
 describe("boolean", () => {
   test("can create a read-only instance", () => {
@@ -64,45 +64,6 @@ describe("literal", () => {
   });
 });
 
-describe("union", () => {
-  const unionType = types.union(types.literal("value 1"), types.literal("value 2"));
-
-  test("can create a read-only instance", () => {
-    expect(unionType.createReadOnly("value 1")).toEqual("value 1");
-    expect(unionType.createReadOnly("value 2")).toEqual("value 2");
-    expect(() => unionType.createReadOnly("value 3" as any)).toThrow();
-  });
-
-  test("can create an eager, read-only instance", () => {
-    const unionType = types.lazyUnion(types.literal("value 1"), types.literal("value 2"));
-    expect(unionType.createReadOnly("value 1")).toEqual("value 1");
-    expect(unionType.createReadOnly("value 2")).toEqual("value 2");
-    expect(() => unionType.createReadOnly("value 3" as any)).toThrow();
-  });
-
-  test("can be verified with is", () => {
-    expect(unionType.is("value 1")).toEqual(true);
-    expect(unionType.is("value 2")).toEqual(true);
-    expect(unionType.is("value 3")).toEqual(false);
-    expect(unionType.is(null)).toEqual(false);
-    expect(unionType.is(true)).toEqual(false);
-    expect(unionType.is({})).toEqual(false);
-  });
-
-  test("can create a union of model types", () => {
-    const modelTypeA = types.model({ x: types.string });
-    const modelTypeB = types.model({ y: types.number });
-    const unionType = types.union(modelTypeA, modelTypeB);
-    const unionInstance = unionType.createReadOnly({ x: "test" });
-
-    expect(getSnapshot(unionInstance)).toEqual(
-      expect.objectContaining({
-        x: "test",
-      })
-    );
-  });
-});
-
 describe("refinement", () => {
   const smallStringsType = types.refinement(types.string, (v: string) => v.length <= 5);
 

spec/union.spec.ts

@@ -0,0 +1,158 @@
+import { getSnapshot, types } from "../src";
+import { create } from "./helpers";
+
+describe("union", () => {
+  const unionType = types.union(types.literal("value 1"), types.literal("value 2"));
+
+  test("can be verified with is", () => {
+    expect(unionType.is("value 1")).toEqual(true);
+    expect(unionType.is("value 2")).toEqual(true);
+    expect(unionType.is("value 3")).toEqual(false);
+    expect(unionType.is(null)).toEqual(false);
+    expect(unionType.is(true)).toEqual(false);
+    expect(unionType.is({})).toEqual(false);
+  });
+
+  describe.each([
+    ["readonly", true],
+    ["observable", false],
+  ])("%s nodes", (_, readonly) => {
+    test("can create an instance", () => {
+      expect(create(unionType, "value 1", readonly)).toEqual("value 1");
+      expect(create(unionType, "value 2", readonly)).toEqual("value 2");
+      expect(() => create(unionType, "value 3" as any, readonly)).toThrow();
+    });
+
+    test("can create an eager, read-only instance", () => {
+      const unionType = types.lazyUnion(types.literal("value 1"), types.literal("value 2"));
+      expect(create(unionType, "value 1", readonly)).toEqual("value 1");
+      expect(create(unionType, "value 2", readonly)).toEqual("value 2");
+      expect(() => create(unionType, "value 3" as any, readonly)).toThrow();
+    });
+
+    test("can create a union of model types", () => {
+      const modelTypeA = types.model({ x: types.string });
+      const modelTypeB = types.model({ y: types.number });
+      const unionType = types.union(modelTypeA, modelTypeB);
+      const unionInstance = create(unionType, { x: "test" }, readonly);
+
+      expect(getSnapshot(unionInstance)).toEqual(
+        expect.objectContaining({
+          x: "test",
+        })
+      );
+    });
+
+    test("nested unions", () => {
+      const one = types.union(types.literal("A"), types.literal("B"));
+      const two = types.union(types.literal("C"), types.literal("D"));
+      const Union = types.union(one, two);
+
+      const aInstance = create(Union, "A", readonly);
+      expect(aInstance).toEqual("A");
+
+      const dInstance = create(Union, "D", readonly);
+      expect(dInstance).toEqual("D");
+    });
+
+    describe("with an explicit dispatcher option", () => {
+      test("can use a dispatcher", () => {
+        const Apple = types.model({ color: types.string });
+        const Banana = types.model({ ripeness: types.number });
+        const Union = types.union(
+          {
+            dispatcher: (snapshot: any) => {
+              if (snapshot.color) {
+                return Apple;
+              } else {
+                return Banana;
+              }
+            },
+          },
+          Apple,
+          Banana
+        );
+
+        const appleInstance = create(Union, { color: "red" }, readonly);
+        expect(Apple.is(appleInstance)).toBeTruthy();
+        expect(Banana.is(appleInstance)).toBeFalsy();
+
+        const bananaInstance = create(Union, { ripeness: 2 }, readonly);
+        expect(Apple.is(bananaInstance)).toBeFalsy();
+        expect(Banana.is(bananaInstance)).toBeTruthy();
+      });
+    });
+
+    describe("with an explicit discriminator property", () => {
+      test("can use a literal discriminator", () => {
+        const Apple = types.model({ type: types.literal("apple"), color: types.string });
+        const Banana = types.model({ type: types.literal("banana"), ripeness: types.number });
+        const Union = types.union({ discriminator: "type" }, Apple, Banana);
+
+        const appleInstance = create(Union, { type: "apple", color: "red" }, readonly);
+        expect(Apple.is(appleInstance)).toBeTruthy();
+        expect(Banana.is(appleInstance)).toBeFalsy();
+
+        const bananaInstance = create(Union, { type: "banana", ripeness: 2 }, readonly);
+        expect(Apple.is(bananaInstance)).toBeFalsy();
+        expect(Banana.is(bananaInstance)).toBeTruthy();
+      });
+
+      test("can use an optional literal discriminator", () => {
+        const Apple = types.model({ type: types.optional(types.literal("apple"), "apple"), color: types.string });
+        const Banana = types.model({ type: types.optional(types.literal("banana"), "banana"), ripeness: types.number });
+        const Union = types.union({ discriminator: "type" }, Apple, Banana);
+
+        const appleInstance = create(Union, { type: "apple", color: "red" }, readonly);
+        expect(Apple.is(appleInstance)).toBeTruthy();
+        expect(Banana.is(appleInstance)).toBeFalsy();
+
+        const bananaInstance = create(Union, { type: "banana", ripeness: 2 }, readonly);
+        expect(Apple.is(bananaInstance)).toBeFalsy();
+        expect(Banana.is(bananaInstance)).toBeTruthy();
+      });
+
+      test("does not get an error when passing a snapshot that doesn't have a value for the discriminator property", () => {
+        const Apple = types.model({ type: types.optional(types.literal("apple"), "apple"), color: types.string });
+        const Banana = types.model({ type: types.optional(types.literal("banana"), "banana"), ripeness: types.number });
+        const Union = types.union({ discriminator: "type" }, Apple, Banana);
+
+        const apple = create(Union, { color: "red" }, readonly);
+        expect(Apple.is(apple)).toBeTruthy();
+      });
+
+      test("can use an literal discriminators within nested unions", () => {
+        const Apple = types.model({ type: types.optional(types.literal("apple"), "apple"), color: types.string });
+        const Banana = types.model({ type: types.optional(types.literal("banana"), "banana"), ripeness: types.number });
+        const InnerUnion = types.union({ discriminator: "type" }, Apple, Banana);
+        const Pear = types.model({ type: types.optional(types.literal("pear"), "pear"), species: types.string });
+        const Union = types.union({ discriminator: "type" }, InnerUnion, Pear);
+
+        const appleInstance = create(Union, { type: "apple", color: "red" }, readonly);
+        expect(Apple.is(appleInstance)).toBeTruthy();
+        expect(Banana.is(appleInstance)).toBeFalsy();
+        expect(Pear.is(appleInstance)).toBeFalsy();
+
+        const bananaInstance = create(Union, { type: "banana", ripeness: 2 }, readonly);
+        expect(Apple.is(bananaInstance)).toBeFalsy();
+        expect(Banana.is(bananaInstance)).toBeTruthy();
+        expect(Pear.is(bananaInstance)).toBeFalsy();
+
+        const pearInstance = create(Union, { type: "pear", species: "anjou" }, readonly);
+        expect(Apple.is(pearInstance)).toBeFalsy();
+        expect(Banana.is(pearInstance)).toBeFalsy();
+        expect(Pear.is(pearInstance)).toBeTruthy();
+      });
+    });
+  });
+
+  test("gets an error when passing a snapshot that has an unrecognized value for the discriminator property", () => {
+    const Apple = types.model({ type: types.optional(types.literal("apple"), "apple"), color: types.string });
+    const Banana = types.model({ type: types.optional(types.literal("banana"), "banana"), ripeness: types.number });
+    const Union = types.union({ discriminator: "type" }, Apple, Banana);
+
+    expect(() => create(Union, { type: "pear" } as any, true)).toThrowErrorMatchingInlineSnapshot(
+      `"Discriminator property value \`pear\` for property \`type\` on incoming snapshot didn't correspond to a type. Options: apple, banana. Snapshot was \`{"type":"pear"}\`"`
+    );
+  });
+});

src/class-model.ts

@@ -471,3 +471,7 @@ function getPropertyDescriptor(obj: any, property: string) {
   }
   return null;
 }
+
+export const isClassModel = (type: IAnyType): type is IClassModelType<any, any, any> => {
+  return (type as any).isMQTClassModel;
+};

src/errors.ts

@@ -3,3 +3,6 @@ export class CantRunActionError extends Error {}
 
 /** Thrown when an invalid registration is passed to the class model register function */
 export class RegistrationError extends Error {}
+
+/** Thrown when a type in a union can't be used for discrimination because the value of the descriminator property can't be determined at runtime */
+export class InvalidDiscriminatorError extends Error {}

src/optional.ts

@@ -13,7 +13,7 @@ import type {
 
 export type DefaultFuncOrValue<T extends IAnyType> = T["InputType"] | T["OutputType"] | (() => CreateTypes<T>);
 
-class OptionalType<T extends IAnyType, OptionalValues extends [ValidOptionalValue, ...ValidOptionalValue[]]> extends BaseType<
+export class OptionalType<T extends IAnyType, OptionalValues extends [ValidOptionalValue, ...ValidOptionalValue[]]> extends BaseType<
   T["InputType"] | OptionalValues[number],
   T["OutputType"],
   InstanceWithoutSTNTypeForType<T>

src/simple.ts

@@ -80,7 +80,7 @@ export class NullType extends BaseType<null, null, null> {
   }
 }
 
-class LiteralType<T extends Primitives> extends SimpleType<T> {
+export class LiteralType<T extends Primitives> extends SimpleType<T> {
   constructor(readonly value: T) {
     super(typeof value, types.literal<T>(value));
   }