Merge pull request #34 from mlabs-haskell/rmgaray/fix-tagged-variants

rmgaray · web-flow · commit 88d055d79335 · 2025-01-31T10:59:32.000-03:00
fix: union types with multiple tagged variants not being decoded correctly
diff --git a/conway-cddl/codegen/generators/union.ts b/conway-cddl/codegen/generators/union.ts
@@ -4,6 +4,7 @@ import { SchemaTable } from "..";
 export type Variant = {
   tag: number; // number assigned to the variant in the FooKind enum
   peek_type: string | string[]; // decode this variant if the CBOR type tag equals any of these values
+  valid_tags?: number[];
   name: string; // used in Class.new_foo()
   type: string; // used to do if (reader.getTag() == tag) type.deserialize()
   kind_name?: string; // name of the variant in the FooKind enum
@@ -13,6 +14,15 @@ export type GenUnionOptions = {
   variants: Variant[];
 } & CodeGeneratorBaseOptions;
 
+// We say tagged to refer to variants that are encoded with a CBOR tag
+type TaggedVariant = {
+  tag: number; // number assigned to the variant in the FooKind enum
+  valid_tags: number[];
+  name: string; // used in Class.new_foo()
+  type: string; // used to do if (reader.getTag() == tag) type.deserialize()
+  kind_name?: string; // name of the variant in the FooKind enum
+}
+
 export class GenUnion extends CodeGeneratorBase {
   variants: Variant[];
 
@@ -87,33 +97,103 @@ export class GenUnion extends CodeGeneratorBase {
   }
 
   generateDeserialize(reader: string, path: string): string {
+    const constructUntagged = (v: Variant) => {
+      let out = "";
+      if (Array.isArray(v.peek_type)) {
+        for (let t of v.peek_type) {
+          out += `case "${t}":\n`;
+        }
+      } else {
+        out += `case "${v.peek_type}":\n`;
+      }
+      return out +
+        `
+        variant = {
+          kind: ${this.name}Kind.${v.kind_name ?? v.type},
+          value: ${this.typeUtils.readType(reader, v.type, `[...${path}, '${v.type}(${v.name})']`)}
+        };
+        break;
+        `
+    }
+
+    const constructTagged = (v: TaggedVariant) => {
+      if(v.valid_tags.length == 0) {
+        throw new Error("Expected a non-empty 'valid_tags' field because multiple tagged variants exist. These are needed to disambiguate.")
+      } else {
+        return `if ([${v.valid_tags.toString()}].includes(tagNumber)) {
+                    variant = {
+                      kind: ${this.name}Kind.${v.kind_name ?? v.type},
+                      value: ${this.typeUtils.readType(reader, v.type, `[...${path}, '${v.type}(${v.name})']`)}
+                    };
+                    break;
+                }`
+      }
+    }
+
+    // split variants into tagged and untagged types
+    let [taggedVariants, untaggedVariants] = this.variants.reduce((acc, v) => {
+      let [tagged, untagged] = acc;
+      if (typeof v.peek_type == "string") {
+        if (v.peek_type == "tagged") {
+          let tagged_v: TaggedVariant = {
+            tag: v.tag,
+            valid_tags: v.valid_tags ? v.valid_tags : [],
+            name: v.name,
+            type: v.type,
+            kind_name: v.kind_name
+          }
+          tagged.push(tagged_v);
+        } else {
+          untagged.push(v);
+        }
+      } else if (v.peek_type.includes("tagged")) {
+          let untagged_v: Variant = structuredClone(v)
+          untagged_v.peek_type = v.peek_type.filter((t) => t != "tagged");
+
+          let tagged_v: TaggedVariant = {
+            tag: v.tag,
+            valid_tags: v.valid_tags ? v.valid_tags : [],
+            name: v.name,
+            type: v.type,
+            kind_name: v.kind_name
+          }
+
+          untagged.push(untagged_v);
+          tagged.push(tagged_v);
+      } else {
+          untagged.push(v);
+      }
+
+      return [tagged, untagged]
+    }, [[], []] as [TaggedVariant[], Variant[]]);
+
     return `
       let tag = ${reader}.peekType(${path});
       let variant: ${this.name}Variant;
 
       switch(tag) {
-        ${this.variants
-          .map((x) => {
-            let out = "";
-            if (Array.isArray(x.peek_type)) {
-              for (let t of x.peek_type) {
-                out += `case "${t}":\n`;
-              }
-            } else {
-              out += `case "${x.peek_type}":\n`;
-            }
-            return (
-              out +
-              `
-              variant = {
-                kind: ${this.name}Kind.${x.kind_name ?? x.type},
-                value: ${this.typeUtils.readType(reader, x.type, `[...${path}, '${x.type}(${x.name})']`)}
-              };
-              break;
-              `
-            );
-          })
+        ${untaggedVariants
+          .map(constructUntagged)
           .join("\n")}
+        ${taggedVariants.length > 0
+            ? (taggedVariants.length == 1
+                ? `case "tagged":
+                      variant = {
+                        kind: ${this.name}Kind.${taggedVariants[0].kind_name ?? taggedVariants[0].type},
+                        value: ${this.typeUtils.readType(reader, taggedVariants[0].type, `[...${path}, '${taggedVariants[0].type}(${taggedVariants[0].name})']`)}
+                      };
+                      break;
+                  `
+                : `case "tagged":
+                      const tagNumber = ${reader}.peekTagNumber(${path});
+                      ${constructTagged(taggedVariants[0])}      
+                      ${taggedVariants.slice(1).map((v) => "else " + constructTagged(v)).join("\n")}
+                      else {
+                        throw new Error("Unexpected tag number " + tagNumber + " (at " + ${path}.join("/") + ")")
+                      }
+                  `)
+            : ''
+          }
         default:
           throw new Error("Unexpected subtype for ${this.name}: " + tag + "(at " + ${path}.join("/") + ")");
       }
diff --git a/conway-cddl/codegen/types.ts b/conway-cddl/codegen/types.ts
@@ -120,6 +120,7 @@ export const Schema = Type.Intersect([
         Type.Object({
           tag: Type.Number(),
           peek_type: Type.Union([Type.String(), Type.Array(Type.String())]),
+          valid_tags: Type.Optional(Type.Array(Type.Number())),
           name: Type.String(),
           type: Type.String(),
           kind_name: Type.Optional(Type.String()),
diff --git a/conway-cddl/yaml/custom/plutus_data.yaml b/conway-cddl/yaml/custom/plutus_data.yaml
@@ -3,6 +3,7 @@ PlutusData:
   variants:
     - tag: 0
       peek_type: "tagged"
+      valid_tags: [102, 121, 122, 123, 124, 125, 126, 127]
       name: constr_plutus_data
       type: ConstrPlutusData
     - tag: 1
@@ -15,6 +16,7 @@ PlutusData:
       type: PlutusList
     - tag: 3
       peek_type: ["uint", "nint", "tagged"]
+      valid_tags: [2, 3]
       name: integer
       type: CSLBigInt
     - tag: 4
diff --git a/src/generated.ts b/src/generated.ts
@@ -10026,16 +10026,6 @@ export class PlutusData {
     let variant: PlutusDataVariant;
 
     switch (tag) {
-      case "tagged":
-        variant = {
-          kind: PlutusDataKind.ConstrPlutusData,
-          value: ConstrPlutusData.deserialize(reader, [
-            ...path,
-            "ConstrPlutusData(constr_plutus_data)",
-          ]),
-        };
-        break;
-
       case "map":
         variant = {
           kind: PlutusDataKind.PlutusMap,
@@ -10052,7 +10042,6 @@ export class PlutusData {
 
       case "uint":
       case "nint":
-      case "tagged":
         variant = {
           kind: PlutusDataKind.CSLBigInt,
           value: CSLBigInt.deserialize(reader, [...path, "CSLBigInt(integer)"]),
@@ -10066,6 +10055,36 @@ export class PlutusData {
         };
         break;
 
+      case "tagged":
+        const tagNumber = reader.peekTagNumber(path);
+        if ([102, 121, 122, 123, 124, 125, 126, 127].includes(tagNumber)) {
+          variant = {
+            kind: PlutusDataKind.ConstrPlutusData,
+            value: ConstrPlutusData.deserialize(reader, [
+              ...path,
+              "ConstrPlutusData(constr_plutus_data)",
+            ]),
+          };
+          break;
+        } else if ([2, 3].includes(tagNumber)) {
+          variant = {
+            kind: PlutusDataKind.CSLBigInt,
+            value: CSLBigInt.deserialize(reader, [
+              ...path,
+              "CSLBigInt(integer)",
+            ]),
+          };
+          break;
+        } else {
+          throw new Error(
+            "Unexpected tag number " +
+              tagNumber +
+              " (at " +
+              path.join("/") +
+              ")",
+          );
+        }
+
       default:
         throw new Error(
           "Unexpected subtype for PlutusData: " +
diff --git a/src/lib/cbor/reader.ts b/src/lib/cbor/reader.ts
@@ -58,6 +58,33 @@ export class CBORReader {
     throw err
   }
 
+  peekTagNumber(path: string[]): number {
+    let tag = this.buffer[0];
+
+    let len = tag & 0b11111;
+
+    // the value of the length field must be between 0x00 and 0x1b
+    if (!(len >= 0x00 && len <= 0x1b)) {
+      let err = new CBORInvalidTag(tag);
+      err.message += ` (at ${path.join("/")})`;
+      throw err;
+    }
+
+    let slicedBuffer = this.buffer.slice(1);
+
+    // if the length field is less than 0x18, then that itself is the value
+    // (optimization for small values)
+    if (len < 0x18) {
+      return Number(BigInt(len));
+    }
+
+    // Else the length is 2^(length - 0x18)
+    let nBytes = Math.pow(2, len - 0x18);
+
+    let x = Number(bigintFromBytes(nBytes, slicedBuffer));
+    return x;
+  }
+
   isBreak(): boolean {
     return this.buffer[0] == 0xff;
   }
