Provable qualified unions

packages/common/src/union/union.ts

@@ -0,0 +1,45 @@
+import { Provable, Struct, Field, InferProvable, Unconstrained } from "o1js";
+
+import { padArray } from "../utils";
+
+export function createQualifiedUnion<
+  T extends (Provable<any> & { name: string })[],
+>(provables: T) {
+  const maxLength = Math.max(
+    ...provables.map((provable) => provable.sizeInFields())
+  );
+
+  const typeMap = Object.fromEntries(
+    provables.map(({ name }, index) => [name, index])
+  );
+
+  class ProvableUnion extends Struct({
+    array: Provable.Array(Field, maxLength),
+    type: Field,
+  }) {
+    public static from<Type extends T[number]>(
+      provable: Type,
+      value: InferProvable<Type>
+    ) {
+      const fields = provable.toFields(value);
+      const fullFields = padArray(fields, maxLength, () => Field(0));
+      return new ProvableUnion({
+        array: fullFields,
+        type: Field(typeMap[provable.name]),
+      });
+    }
+
+    public into<Type extends T[number]>(provable: Type): InferProvable<Type> {
+      const size = provable.sizeInFields();
+      const fields = this.array.slice(0, size);
+
+      this.array.slice(size).forEach((field) => field.assertEquals(0));
+
+      this.type.assertEquals(typeMap[provable.name]);
+
+      return provable.fromFields(fields, []);
+    }
+  }
+
+  return ProvableUnion;
+}

packages/common/test/union/union.test.ts

@@ -0,0 +1,16 @@
+import { Field, UInt64 } from "o1js";
+
+import { createQualifiedUnion } from "../../src/union/union";
+
+describe("union", () => {
+  it("should serialize correctly", () => {
+    const provable = createQualifiedUnion([Field, UInt64]);
+    const p = provable.from(UInt64, UInt64.from(1));
+    // const p2 = provable.from(UInt32, UInt32.from(1));
+
+    const uint = p.into(UInt64);
+    const x = uint.add(UInt64.from(2)).toString();
+
+    expect(x).toBe("3");
+  });
+});