feat(molecule): make union support the custom id

ckb-js · Dec 13, 2023 · be1d723 · be1d723
1 parent 04f5146
commit be1d723
Show file tree

Hide file tree

Showing 9 changed files with 153 additions and 32 deletions.
diff --git a/.changeset/fluffy-rabbits-tell.md b/.changeset/fluffy-rabbits-tell.md
@@ -0,0 +1,5 @@
+---
+"@ckb-lumos/molecule": minor
+---
+
+feat: supported union with custom id
diff --git a/commitlint.config.js b/commitlint.config.js
@@ -22,6 +22,7 @@ const scopeEnumValues = [
   "utils",
   "runner",
   "e2e-test",
+  "molecule",
 ];
 const Configuration = {
   extends: ["@commitlint/config-conventional"],

diff --git a/packages/molecule/README.md b/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);
+```
diff --git a/packages/molecule/src/codec.ts b/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": {

diff --git a/packages/molecule/src/grammar/mol.js b/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]),
           };
         },
       },

diff --git a/packages/molecule/src/grammar/mol.ne b/packages/molecule/src/grammar/mol.ne
@@ -117,20 +117,20 @@ top_level_statement
 
 array_definition
     -> "array" __ identifier _ lbracket _ identifier _ semicolon _ number _ rbracket _ semicolon _ comment_opt
-        {% 
+        {%
             function(data) {
                 return {
                     type: "array",
                     name: data[2].value,
                     item:  data[6].value,
-                    item_count: data[10].value 
+                    item_count: data[10].value
                 };
             }
         %}
 
 vector_definition
      -> "vector" __ identifier _ labracket _ identifier _ rabracket _ semicolon _ comment_opt
-        {% 
+        {%
             function(data) {
                 return {
                     type: "vector",
@@ -142,7 +142,7 @@ vector_definition
 
 option_definition
      -> "option" __ identifier _ lparan _ identifier _ rparan _ semicolon _ comment_opt
-        {% 
+        {%
             function(data) {
                 return {
                     type: "option",
@@ -152,21 +152,35 @@ option_definition
             }
         %}
 
+union_item_decl
+     -> identifier _ ":" _ number
+     {%
+        function (data) {
+            return [data[0].value, Number(data[4].value)]
+        }
+     %}
+     | identifier
+     {%
+        function (data) {
+            return data[0].value
+        }
+    %}
+
 union_definition
-     -> "union" __ identifier _ lbrace _ (multi_line_ws_char _ identifier _ comma _ comment_opt _ multi_line_ws_char):+  _ rbrace
-        {% 
+     -> "union" __ identifier _ lbrace _ (multi_line_ws_char _ union_item_decl _ comma _ comment_opt _ multi_line_ws_char):+  _ rbrace
+        {%
             function(data) {
                 return {
                     type: "union",
                     name: data[2].value,
-                    items:  data[6].map(d => d[2].value),
+                    items:  data[6].map(d => d[2]),
                 };
             }
         %}
 
 struct_definition
      -> "struct" __ identifier _ block_definition
-        {% 
+        {%
             function(data) {
                 return {
                     type: "struct",
@@ -178,7 +192,7 @@ struct_definition
 
 table_definition
      -> "table" __ identifier _ block_definition
-        {% 
+        {%
             function(data) {
                 return {
                     type: "table",
@@ -217,4 +231,4 @@ multi_line_ws_char
 
 __ -> %ws:+
 
-_ -> %ws:*
+_ -> %ws:*
diff --git a/packages/molecule/src/nearley.ts b/packages/molecule/src/nearley.ts
@@ -13,6 +13,7 @@ import {
   ParseOptions,
 } from "./type";
 import { nonNull, toMolTypeMap } from "./utils";
+import { Uint32 } from "@ckb-lumos/codec/lib/number";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const grammar = require("./grammar/mol.js");
@@ -92,10 +93,16 @@ export const checkDependencies = (results: MolType[]): void => {
       }
       case "union": {
         const unionDeps = (molItem as Union).items;
-        unionDeps.forEach((dep: string) => {
-          if (dep !== byte) {
+        unionDeps.forEach((dep) => {
+          if (typeof dep === "string" && dep !== byte) {
             nonNull(map[dep]);
           }
+          if (Array.isArray(dep)) {
+            const [key, id] = dep;
+            // check if the id is a valid uint32
+            Uint32.pack(id);
+            nonNull(map[key]);
+          }
         });
         break;
       }

diff --git a/packages/molecule/src/type.ts b/packages/molecule/src/type.ts
@@ -27,7 +27,7 @@ export type Option = {
 export type Union = {
   type: "union";
   name: string;
-  items: string[];
+  items: (string | [string, number])[];
 };
 
 export type Struct = {

diff --git a/packages/molecule/tests/grammar.test.ts b/packages/molecule/tests/grammar.test.ts
@@ -45,6 +45,51 @@ test("should parse sample with refs", (t) => {
   );
 });
 
+test("union with custom id", (t) => {
+  const parser = createParser();
+  const withCustomId = parser.parse(`
+    array Uint8 [byte; 1];
+    array Uint16 [byte; 2];
+    array Uint32 [byte; 4];
+    
+    union JSNumber {
+      Uint8: 8,
+      Uint16: 16,
+      Uint32: 32,
+    }
+  `);
+
+  t.deepEqual(
+    withCustomId.JSNumber.pack({ type: "Uint8", value: 1 }),
+    // prettier-ignore
+    Uint8Array.from([
+      0x08, 0x00, 0x00, 0x00, // id should be 8
+      0x01
+    ])
+  );
+
+  const withoutCustomId = parser.parse(`
+    array Uint8 [byte; 1];
+    array Uint16 [byte; 2];
+    array Uint32 [byte; 4];
+    
+    union JSNumber {
+      Uint8,
+      Uint16,
+      Uint32,
+    }
+  `);
+
+  t.deepEqual(
+    withoutCustomId.JSNumber.pack({ type: "Uint8", value: 1 }),
+    // prettier-ignore
+    Uint8Array.from([
+      0x00, 0x00, 0x00, 0x00, // id should be 0
+      0x01
+    ])
+  );
+});
+
 test("should parse blockchain.mol", (t) => {
   const parser = createParser();
   // https://github.com/nervosnetwork/ckb/blob/5a7efe7a0b720de79ff3761dc6e8424b8d5b22ea/util/types/schemas/blockchain.mol