feat(codec): supported union with custom id (#586)

.changeset/fluffy-rabbits-tell.md

@@ -0,0 +1,6 @@
+---
+"@ckb-lumos/molecule": minor
+"@ckb-lumos/codec": minor
+---
+
+feat: supported union works with custom id

.changeset/pre.json

@@ -0,0 +1,28 @@
+{
+  "mode": "pre",
+  "tag": "next",
+  "initialVersions": {
+    "@ckb-lumos/base": "0.21.1",
+    "@ckb-lumos/bi": "0.21.1",
+    "@ckb-lumos/ckb-indexer": "0.21.1",
+    "@ckb-lumos/codec": "0.21.1",
+    "@ckb-lumos/common-scripts": "0.21.1",
+    "@ckb-lumos/config-manager": "0.21.1",
+    "@ckb-lumos/debugger": "0.21.1",
+    "@ckb-lumos/e2e-test": "0.21.1",
+    "@ckb-lumos/experiment-tx-assembler": "0.21.1",
+    "@ckb-lumos/hd": "0.21.1",
+    "@ckb-lumos/hd-cache": "0.21.1",
+    "@ckb-lumos/helpers": "0.21.1",
+    "@ckb-lumos/light-client": "0.21.1",
+    "@ckb-lumos/lumos": "0.21.1",
+    "@ckb-lumos/molecule": "0.21.1",
+    "@ckb-lumos/rpc": "0.21.1",
+    "@ckb-lumos/runner": "0.21.1",
+    "@ckb-lumos/testkit": "0.21.1",
+    "@ckb-lumos/toolkit": "0.21.1",
+    "@ckb-lumos/transaction-manager": "0.21.1",
+    "@ckb-lumos/utils": "0.21.1"
+  },
+  "changesets": []
+}

commitlint.config.js

@@ -22,6 +22,7 @@ const scopeEnumValues = [
   "utils",
   "runner",
   "e2e-test",
+  "molecule",
 ];
 const Configuration = {
   extends: ["@commitlint/config-conventional"],

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,36 @@ 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));
+
+  // check duplicated id
+  if (!Array.isArray(fields)) {
+    const ids = Object.values(fields);
+    if (ids.length !== new Set(ids).size) {
+      throw new Error(`Duplicated id in union: ${ids.join(", ")}`);
+    }
+  }
+
   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 +361,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,46 @@ 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 }));
+  // prettier-ignore
+  t.throws(() => codec.unpack([
+    0x00, 0x00, 0x00, 0x00, // unknown key
+    0x11,
+  ]));
+});
+
+test("test union with duplicated custom id", (t) => {
+  t.throws(() => union({ key1: Uint8, key2: Uint32LE }, { key1: 0, key2: 0 }));
+});
+
 test("test byteOf", (t) => {
   t.deepEqual(byteOf(Uint8).pack(1), bytify([1]));
   t.throws(() => byteOf(Uint16).pack(1));
@@ -316,7 +356,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 },

packages/molecule/README.md

@@ -0,0 +1,14 @@
+# @ckb-lumos/molecule
+
+A molecule parser written in JavaScript that helps developers to parse molecule into a codec map.
+
+```js
+const { createParser } = require("@ckb-lumos/molecule");
+
+const parser = createParser();
+const codecMap = parser.parse(`
+  array Uint8 [byte; 1];
+`);
+
+codecMap.Uint8.pack(1);
+```

packages/molecule/src/codec.ts

@@ -1,8 +1,8 @@
 import {
-  FixedBytesCodec,
-  createBytesCodec,
-  BytesLike,
   BytesCodec,
+  BytesLike,
+  createBytesCodec,
+  FixedBytesCodec,
 } from "@ckb-lumos/codec/lib/base";
 import {
   array,
@@ -46,7 +46,7 @@ export const toCodec = (
   }
   const molType: MolType = molTypeMap[key];
   nonNull(molType);
-  let codec = null;
+  let codec: BytesCodec | null = null;
   switch (molType.type) {
     case "array": {
       if (molType.name.startsWith("Uint")) {
@@ -104,21 +104,42 @@ export const toCodec = (
       break;
     }
     case "union": {
-      const unionCodecs: Record<string, BytesCodec> = {};
-      molType.items.forEach((itemMolTypeName) => {
-        if (itemMolTypeName === byte) {
-          unionCodecs[itemMolTypeName] = createFixedHexBytesCodec(1);
+      // Tuple of [UnionFieldName, UnionFieldId, UnionTypeCodec]
+      const unionCodecs: [string, number, BytesCodec][] = [];
+
+      molType.items.forEach((unionTypeItem, index) => {
+        if (unionTypeItem === byte) {
+          unionCodecs.push([unionTypeItem, index, createFixedHexBytesCodec(1)]);
         } else {
-          const itemMolType = toCodec(
-            itemMolTypeName,
-            molTypeMap,
-            result,
-            refs
-          );
-          unionCodecs[itemMolTypeName] = itemMolType;
+          if (typeof unionTypeItem === "string") {
+            const itemMolType = toCodec(
+              unionTypeItem,
+              molTypeMap,
+              result,
+              refs
+            );
+            unionCodecs.push([unionTypeItem, index, itemMolType]);
+          } else if (Array.isArray(unionTypeItem)) {
+            const [key, fieldId] = unionTypeItem;
+
+            const itemMolType = toCodec(key, molTypeMap, result, refs);
+            unionCodecs.push([key, fieldId, itemMolType]);
+          }
         }
       });
-      codec = union(unionCodecs, Object.keys(unionCodecs));
+
+      const unionFieldsCodecs: Record<string, BytesCodec> = unionCodecs.reduce(
+        (codecMap, [fieldName, _fieldId, fieldCodec]) =>
+          Object.assign(codecMap, { [fieldName]: fieldCodec }),
+        {}
+      );
+      const unionFieldIds: Record<string, number> = unionCodecs.reduce(
+        (idMap, [fieldName, fieldId, _fieldCodec]) =>
+          Object.assign(idMap, { [fieldName]: fieldId }),
+        {}
+      );
+
+      codec = union(unionFieldsCodecs, unionFieldIds);
       break;
     }
     case "table": {

packages/molecule/src/grammar/mol.js

@@ -228,12 +228,26 @@
           };
         },
       },
+      {
+        name: "union_item_decl",
+        symbols: ["identifier", "_", { literal: ":" }, "_", "number"],
+        postprocess: function (data) {
+          return [data[0].value, Number(data[4].value)];
+        },
+      },
+      {
+        name: "union_item_decl",
+        symbols: ["identifier"],
+        postprocess: function (data) {
+          return data[0].value;
+        },
+      },
       {
         name: "union_definition$ebnf$1$subexpression$1",
         symbols: [
           "multi_line_ws_char",
           "_",
-          "identifier",
+          "union_item_decl",
           "_",
           "comma",
           "_",
@@ -251,7 +265,7 @@
         symbols: [
           "multi_line_ws_char",
           "_",
-          "identifier",
+          "union_item_decl",
           "_",
           "comma",
           "_",
@@ -287,7 +301,7 @@
           return {
             type: "union",
             name: data[2].value,
-            items: data[6].map((d) => d[2].value),
+            items: data[6].map((d) => d[2]),
           };
         },
       },