feat(codec): supported union with custom id

.changeset/tender-apes-speak.md

@@ -0,0 +1,5 @@
+---
+"@ckb-lumos/codec": minor
+---
+
+feat: supported custom union id in union

packages/codec/src/molecule/layout.ts

@@ -131,11 +131,9 @@ export function struct<T extends Record<string, FixedBytesCodec>>(
       }, Uint8Array.from([]));
     },
     unpack(buf) {
-      const result = {} as PartialNullable<
-        {
-          [key in keyof T]: UnpackResult<T[key]>;
-        }
-      >;
+      const result = {} as PartialNullable<{
+        [key in keyof T]: UnpackResult<T[key]>;
+      }>;
       let offset = 0;
 
       fields.forEach((field) => {
@@ -296,11 +294,9 @@ export function table<T extends Record<string, BytesCodec>>(
         );
       }
       if (totalSize <= 4 || fields.length === 0) {
-        return {} as PartialNullable<
-          {
-            [key in keyof T]: UnpackResult<T[key]>;
-          }
-        >;
+        return {} as PartialNullable<{
+          [key in keyof T]: UnpackResult<T[key]>;
+        }>;
       } else {
         const offsets = fields.map((_, index) =>
           Uint32LE.unpack(buf.slice(4 + index * 4, 8 + index * 4))
@@ -315,11 +311,9 @@ export function table<T extends Record<string, BytesCodec>>(
           const itemBuf = buf.slice(start, end);
           Object.assign(obj, { [field]: itemCodec.unpack(itemBuf) });
         }
-        return obj as PartialNullable<
-          {
-            [key in keyof T]: UnpackResult<T[key]>;
-          }
-        >;
+        return obj as PartialNullable<{
+          [key in keyof T]: UnpackResult<T[key]>;
+        }>;
       }
     },
   });
@@ -328,19 +322,28 @@ export function table<T extends Record<string, BytesCodec>>(
 /**
  * Union is a dynamic-size type.
  * Serializing a union has two steps:
- * - Serialize a item type id in bytes as a 32 bit unsigned integer in little-endian. The item type id is the index of the inner items, and it's starting at 0.
+ * - Serialize an item type id in bytes as a 32 bit unsigned integer in little-endian. The item type id is the index of the inner items, and it's starting at 0.
  * - Serialize the inner item.
  * @param itemCodec the union item record
- * @param fields the list of itemCodec's keys. It's also provide an order for pack/unpack.
+ * @param fields the union item keys, can be an array or an object with custom id
+ * @example
+ * // without custom id
+ * union({ cafe: Uint8, bee: Uint8 }, ['cafe', 'bee'])
+ * // with custom id
+ * union({ cafe: Uint8, bee: Uint8 }, { cafe: 0xcafe, bee: 0xbee })
  */
 export function union<T extends Record<string, BytesCodec>>(
   itemCodec: T,
-  fields: (keyof T)[]
+  fields: (keyof T)[] | Record<keyof T, number>
 ): UnionCodec<T> {
+  checkShape(itemCodec, Array.isArray(fields) ? fields : Object.keys(fields));
+
   return createBytesCodec({
     pack(obj) {
+      const availableFields: (keyof T)[] = Object.keys(itemCodec);
+
       const type = obj.type;
-      const typeName = `Union(${fields.join(" | ")})`;
+      const typeName = `Union(${availableFields.join(" | ")})`;
 
       /* c8 ignore next */
       if (typeof type !== "string") {
@@ -350,20 +353,38 @@ export function union<T extends Record<string, BytesCodec>>(
         );
       }
 
-      const fieldIndex = fields.indexOf(type);
-      if (fieldIndex === -1) {
+      const fieldId = Array.isArray(fields)
+        ? fields.indexOf(type)
+        : fields[type];
+
+      if (fieldId < 0) {
         throw new CodecBaseParseError(
           `Unknown union type: ${String(obj.type)}`,
           typeName
         );
       }
-      const packedFieldIndex = Uint32LE.pack(fieldIndex);
+      const packedFieldIndex = Uint32LE.pack(fieldId);
       const packedBody = itemCodec[type].pack(obj.value);
       return concat(packedFieldIndex, packedBody);
     },
     unpack(buf) {
-      const typeIndex = Uint32LE.unpack(buf.slice(0, 4));
-      const type = fields[typeIndex];
+      const fieldId = Uint32LE.unpack(buf.slice(0, 4));
+
+      const type: keyof T | undefined = (() => {
+        if (Array.isArray(fields)) {
+          return fields[fieldId];
+        }
+
+        const entry = Object.entries(fields).find(([, id]) => id === fieldId);
+        return entry?.[0];
+      })();
+
+      if (!type) {
+        throw new Error(
+          `Unknown union field id: ${fieldId}, only ${fields} are allowed`
+        );
+      }
+
       return { type, value: itemCodec[type].unpack(buf.slice(4)) };
     },
   });

packages/codec/tests/molecule.test.ts

@@ -12,7 +12,7 @@ import {
 import { Bytes, createFixedHexBytesCodec } from "../src/blockchain";
 import { bytify } from "../src/bytes";
 import test, { ExecutionContext } from "ava";
-import { Uint16, Uint16BE, Uint32, Uint8 } from "../src/number";
+import { Uint16, Uint16BE, Uint32, Uint32LE, Uint8 } from "../src/number";
 import { byteOf } from "../src/molecule";
 import { CodecExecuteError } from "../src/error";
 
@@ -193,6 +193,37 @@ test("test layout-union", (t) => {
   t.throws(() => codec.pack({ type: "unknown", value: [] }));
 });
 
+test("test union with custom id", (t) => {
+  const codec = union(
+    { key1: Uint8, key2: Uint32LE },
+    { key1: 0xaa, key2: 0xbb }
+  );
+
+  // prettier-ignore
+  const case1 = bytify([
+    0xaa, 0x00, 0x00, 0x00, // key1
+    0x11, // value
+  ]);
+
+  t.deepEqual(codec.unpack(case1), { type: "key1", value: 0x11 });
+  t.deepEqual(codec.pack({ type: "key1", value: 0x11 }), case1);
+
+  // prettier-ignore
+  const case2 = bytify([
+    0xbb, 0x00, 0x00, 0x00, // key2
+    0x00, 0x00, 0x00, 0x11, // value u32le
+  ])
+
+  t.deepEqual(codec.unpack(case2), { type: "key2", value: 0x11_00_00_00 });
+  t.deepEqual(codec.pack({ type: "key2", value: 0x11_00_00_00 }), case2);
+
+  // @ts-expect-error
+  t.throws(() => codec.pack({ type: "unknown", value: 0x11 }));
+
+  // @ts-expect-error
+  t.throws(() => union({ key1: Uint8, key2: Uint32LE }, { unknown: 0x1 }));
+});
+
 test("test byteOf", (t) => {
   t.deepEqual(byteOf(Uint8).pack(1), bytify([1]));
   t.throws(() => byteOf(Uint16).pack(1));
@@ -316,7 +347,7 @@ test("nested type", (t) => {
     ["byteField", "arrayField", "structField", "fixedVec", "dynVec", "option"]
   );
 
-  const validInput: Parameters<typeof codec["pack"]>[0] = {
+  const validInput: Parameters<(typeof codec)["pack"]>[0] = {
     byteField: 0x1,
     arrayField: [0x2, 0x3, 0x4],
     structField: { f1: 0x5, f2: 0x6 },