IdlCoder

import { Layout } from "buffer-layout";
import * as borsh from "@coral-xyz/borsh";
import {
  IdlField,
  IdlTypeDef,
  IdlType,
  IdlGenericArg,
  Idl,
  handleDefinedFields,
  IdlArrayLen,
} from "../../idl.js";
import { IdlError } from "../../error.js";

type PartialField = { name?: string } & Pick<IdlField, "type">;

export class IdlCoder {
  public static fieldLayout(
    field: PartialField,
    types: IdlTypeDef[] = [],
    genericArgs?: IdlGenericArg[] | null
  ): Layout {
    const fieldName = field.name;
    switch (field.type) {
      case "bool": {
        return borsh.bool(fieldName);
      }
      case "u8": {
        return borsh.u8(fieldName);
      }
      case "i8": {
        return borsh.i8(fieldName);
      }
      case "u16": {
        return borsh.u16(fieldName);
      }
      case "i16": {
        return borsh.i16(fieldName);
      }
      case "u32": {
        return borsh.u32(fieldName);
      }
      case "i32": {
        return borsh.i32(fieldName);
      }
      case "f32": {
        return borsh.f32(fieldName);
      }
      case "u64": {
        return borsh.u64(fieldName);
      }
      case "i64": {
        return borsh.i64(fieldName);
      }
      case "f64": {
        return borsh.f64(fieldName);
      }
      case "u128": {
        return borsh.u128(fieldName);
      }
      case "i128": {
        return borsh.i128(fieldName);
      }
      case "u256": {
        return borsh.u256(fieldName);
      }
      case "i256": {
        return borsh.i256(fieldName);
      }
      case "bytes": {
        return borsh.vecU8(fieldName);
      }
      case "string": {
        return borsh.str(fieldName);
      }
      case "pubkey": {
        return borsh.publicKey(fieldName);
      }
      default: {
        if ("option" in field.type) {
          return borsh.option(
            IdlCoder.fieldLayout(
              { type: field.type.option },
              types,
              genericArgs
            ),
            fieldName
          );
        }
        if ("vec" in field.type) {
          return borsh.vec(
            IdlCoder.fieldLayout({ type: field.type.vec }, types, genericArgs),
            fieldName
          );
        }
        if ("array" in field.type) {
          let [type, len] = field.type.array;
          len = IdlCoder.resolveArrayLen(len, genericArgs);

          return borsh.array(
            IdlCoder.fieldLayout({ type }, types, genericArgs),
            len,
            fieldName
          );
        }
        if ("defined" in field.type) {
          if (!types) {
            throw new IdlError("User defined types not provided");
          }

          const definedName = field.type.defined.name;
          const typeDef = types.find((t) => t.name === definedName);
          if (!typeDef) {
            throw new IdlError(`Type not found: ${field.name}`);
          }
          console.log("typeDef", typeDef);
          console.log("fieldName", fieldName);
          return IdlCoder.typeDefLayout({
            typeDef,
            types,
            genericArgs: genericArgs ?? field.type.defined.generics,
            name: fieldName,
          });
        }
        if ("generic" in field.type) {
          const genericArg = genericArgs?.at(0);
          if (genericArg?.kind !== "type") {
            throw new IdlError(`Invalid generic field: ${field.name}`);
          }

          return IdlCoder.fieldLayout(
            { ...field, type: genericArg.type },
            types
          );
        }

        throw new IdlError(
          `Not yet implemented: ${JSON.stringify(field.type)}`
        );
      }
    }
  }

  /**
   * Get the type layout of the given defined type(struct or enum).
   */
  public static typeDefLayout({
    typeDef,
    types,
    name,
    genericArgs,
  }: {
    typeDef: IdlTypeDef;
    types: IdlTypeDef[];
    genericArgs?: IdlGenericArg[] | null;
    name?: string;
  }): Layout {
    switch (typeDef.type.kind) {
      case "struct": {
        const fieldLayouts = handleDefinedFields(
          typeDef.type.fields,
          () => [],
          (fields) =>
            fields.map((f) => {
              const genArgs = genericArgs
                ? IdlCoder.resolveGenericArgs({
                  type: f.type,
                  typeDef,
                  genericArgs,
                })
                : genericArgs;
              return IdlCoder.fieldLayout(f, types, genArgs);
            }),
          (fields) =>
            fields.map((f, i) => {
              const genArgs = genericArgs
                ? IdlCoder.resolveGenericArgs({
                  type: f,
                  typeDef,
                  genericArgs,
                })
                : genericArgs;
              return IdlCoder.fieldLayout(
                { name: i.toString(), type: f },
                types,
                genArgs
              );
            })
        );

        return borsh.struct(fieldLayouts, name);
      }

      case "enum": {
        const variants = typeDef.type.variants.map((variant) => {
          const fieldLayouts = handleDefinedFields(
            variant.fields,
            () => [],
            (fields) =>
              fields.map((f) => {
                const genArgs = genericArgs
                  ? IdlCoder.resolveGenericArgs({
                    type: f.type,
                    typeDef,
                    genericArgs,
                  })
                  : genericArgs;
                return IdlCoder.fieldLayout(f, types, genArgs);
              }),
            (fields) =>
              fields.map((f, i) => {
                const genArgs = genericArgs
                  ? IdlCoder.resolveGenericArgs({
                    type: f,
                    typeDef,
                    genericArgs,
                  })
                  : genericArgs;
                return IdlCoder.fieldLayout(
                  { name: i.toString(), type: f },
                  types,
                  genArgs
                );
              })
          );

          return borsh.struct(fieldLayouts, variant.name);
        });

        if (name !== undefined) {
          // Buffer-layout lib requires the name to be null (on construction)
          // when used as a field.
          return borsh.rustEnum(variants).replicate(name);
        }

        return borsh.rustEnum(variants, name);
      }

      case "type": {
        return IdlCoder.fieldLayout({ type: typeDef.type.alias, name }, types);
      }
    }
  }

  /**
   * Get the type of the size in bytes. Returns `1` for variable length types.
   */
  public static typeSize(
    ty: IdlType,
    idl: Idl,
    genericArgs?: IdlGenericArg[] | null
  ): number {
    switch (ty) {
      case "bool":
        return 1;
      case "u8":
        return 1;
      case "i8":
        return 1;
      case "i16":
        return 2;
      case "u16":
        return 2;
      case "u32":
        return 4;
      case "i32":
        return 4;
      case "f32":
        return 4;
      case "u64":
        return 8;
      case "i64":
        return 8;
      case "f64":
        return 8;
      case "u128":
        return 16;
      case "i128":
        return 16;
      case "u256":
        return 32;
      case "i256":
        return 32;
      case "bytes":
        return 1;
      case "string":
        return 1;
      case "pubkey":
        return 32;
      default:
        if ("option" in ty) {
          return 1 + IdlCoder.typeSize(ty.option, idl, genericArgs);
        }
        if ("coption" in ty) {
          return 4 + IdlCoder.typeSize(ty.coption, idl, genericArgs);
        }
        if ("vec" in ty) {
          return 1;
        }
        if ("array" in ty) {
          let [type, len] = ty.array;
          len = IdlCoder.resolveArrayLen(len, genericArgs);
          return IdlCoder.typeSize(type, idl, genericArgs) * len;
        }
        if ("defined" in ty) {
          const typeDef = idl.types?.find((t) => t.name === ty.defined.name);
          if (!typeDef) {
            throw new IdlError(`Type not found: ${JSON.stringify(ty)}`);
          }

          const typeSize = (type: IdlType) => {
            const genArgs = genericArgs ?? ty.defined.generics;
            const args = genArgs
              ? IdlCoder.resolveGenericArgs({
                type,
                typeDef,
                genericArgs: genArgs,
              })
              : genArgs;

            return IdlCoder.typeSize(type, idl, args);
          };

          switch (typeDef.type.kind) {
            case "struct": {
              return handleDefinedFields(
                typeDef.type.fields,
                () => [0],
                (fields) => fields.map((f) => typeSize(f.type)),
                (fields) => fields.map((f) => typeSize(f))
              ).reduce((acc, size) => acc + size, 0);
            }

            case "enum": {
              const variantSizes = typeDef.type.variants.map((variant) => {
                return handleDefinedFields(
                  variant.fields,
                  () => [0],
                  (fields) => fields.map((f) => typeSize(f.type)),
                  (fields) => fields.map((f) => typeSize(f))
                ).reduce((acc, size) => acc + size, 0);
              });

              return Math.max(...variantSizes) + 1;
            }

            case "type": {
              return IdlCoder.typeSize(typeDef.type.alias, idl, genericArgs);
            }
          }
        }
        if ("generic" in ty) {
          const genericArg = genericArgs?.at(0);
          if (genericArg?.kind !== "type") {
            throw new IdlError(`Invalid generic: ${ty.generic}`);
          }

          return IdlCoder.typeSize(genericArg.type, idl, genericArgs);
        }

        throw new Error(`Invalid type ${JSON.stringify(ty)}`);
    }
  }

  /**
   * Resolve the generic array length or return the constant-sized array length.
   */
  private static resolveArrayLen(
    len: IdlArrayLen,
    genericArgs?: IdlGenericArg[] | null
  ): number {
    if (typeof len === "number") return len;

    if (genericArgs) {
      const genericLen = genericArgs.find((g) => g.kind === "const");
      if (genericLen?.kind === "const") {
        len = +genericLen.value;
      }
    }

    if (typeof len !== "number") {
      throw new IdlError("Generic array length did not resolve");
    }

    return len;
  }

  /**
   * Recursively resolve generic arguments i.e. replace all generics with the
   * actual type that they hold based on the initial `genericArgs` given.
   */
  private static resolveGenericArgs({
    type,
    typeDef,
    genericArgs,
    isDefined,
  }: {
    type: IdlType;
    typeDef: IdlTypeDef;
    genericArgs: IdlGenericArg[];
    isDefined?: boolean;
  }): IdlGenericArg[] | null {
    if (typeof type !== "object") return null;

    for (const index in typeDef.generics) {
      const defGeneric = typeDef.generics[index];

      if ("generic" in type && defGeneric.name === type.generic) {
        return [genericArgs[index]];
      }

      if ("option" in type) {
        const args = IdlCoder.resolveGenericArgs({
          type: type.option,
          typeDef,
          genericArgs,
          isDefined,
        });
        if (!args || !isDefined) return args;

        if (args[0].kind === "type") {
          return [
            {
              kind: "type",
              type: { option: args[0].type },
            },
          ];
        }
      }

      if ("vec" in type) {
        const args = IdlCoder.resolveGenericArgs({
          type: type.vec,
          typeDef,
          genericArgs,
          isDefined,
        });
        if (!args || !isDefined) return args;

        if (args[0].kind === "type") {
          return [
            {
              kind: "type",
              type: { vec: args[0].type },
            },
          ];
        }
      }

      if ("array" in type) {
        const [elTy, len] = type.array;
        const isGenericLen = typeof len === "object";

        const args = IdlCoder.resolveGenericArgs({
          type: elTy,
          typeDef,
          genericArgs,
          isDefined,
        });
        if (args) {
          // Has generic type, also check for generic length
          for (const i in typeDef.generics.slice(+index)) {
            const curIndex = +index + +i;
            if (
              isGenericLen &&
              typeDef.generics[curIndex].name === len.generic
            ) {
              args.push(genericArgs[curIndex]);
            }
          }

          if (!isDefined) return args;

          if (args[0].kind === "type" && args[1].kind === "const") {
            return [
              {
                kind: "type",
                type: { array: [args[0].type, +args[1].value] },
              },
            ];
          }
        }

        // Only generic len
        if (isGenericLen && defGeneric.name === len.generic) {
          const arg = genericArgs[index];
          if (!isDefined) return [arg];

          return [
            {
              kind: "type",
              type: { array: [elTy, +arg.value] },
            },
          ];
        }

        // Non-generic
        return null;
      }

      if ("defined" in type) {
        if (!type.defined.generics) return null;

        return type.defined.generics
          .flatMap((g) => {
            switch (g.kind) {
              case "type":
                return IdlCoder.resolveGenericArgs({
                  type: g.type,
                  typeDef,
                  genericArgs,
                  isDefined: true,
                });
              case "const":
                return [g];
            }
          })
          .filter((g) => g !== null) as IdlGenericArg[];
      }
    }

    return null;
  }
}

这段代码实现了一个名为 IdlCoder 的类，用于处理基于 IDL（Interface Definition Language） 的数据结构编码和解码。IDL 通常用于定义数据结构的接口，特别是在区块链和智能合约领域（例如 Solana 的 Anchor 框架）。以下是对代码的逐步分析：

---

1. 核心功能概述

IdlCoder 提供以下主要功能：

• 字段布局生成：
  • 根据字段类型 (IdlField.type) 动态生成字段的内存布局。
  • 使用 borsh 库（一个二进制序列化库）来定义字段的序列化布局。
• 用户定义类型（struct 和 enum）的布局生成：
  • 支持复杂的用户定义类型（如结构体、枚举等）的布局。
• 类型大小计算：
  • 根据类型计算其在内存中的大小。
• 泛型解析：
  • 支持泛型类型的解析和替换。
• 数组长度解析：
  • 处理固定长度和泛型长度的数组。

---

2. 关键模块和依赖

2.1 borsh

borsh 是一个常用的序列化库，用于将数据结构序列化为二进制格式。它支持基本类型（如 u8, i32, string）以及复杂类型（如 struct, enum, array）。

在代码中，borsh 提供了以下方法来定义字段的布局：

• 基本类型：borsh.bool, borsh.u8, borsh.i32, 等。
• 复合类型：borsh.struct, borsh.rustEnum, borsh.array, 等。
• 特殊类型：borsh.option（可选类型），borsh.vec（动态数组）。

2.2 IDL 类型

IDL 类型是基于 Anchor 的 IDL 格式定义的，支持以下类型：

• 基本类型：bool, u8, i32, string, 等。
• 复合类型：struct, enum, array, vec, 等。
• 泛型：如 option<T>, vec<T>, array<T, N>。
• 用户定义类型：defined 类型（如自定义的结构体或枚举）。

2.3 IdlField 和 IdlTypeDef

• IdlField：表示单个字段，包含字段的 name 和 type 信息。
• IdlTypeDef：表示用户定义类型（如 struct 或 enum）。

---

3. 主要方法分析

3.1 fieldLayout

• 功能：根据字段类型生成其对应的 borsh 布局。
• 实现：
  • 检查字段的 type，根据类型调用对应的 borsh 方法生成布局。
  • 处理复杂类型：
    • option：递归调用 fieldLayout 生成内部类型的布局，然后用 borsh.option 包装。
    • vec：递归调用 fieldLayout 生成元素类型的布局，然后用 borsh.vec 包装。
    • array：解析数组的长度和元素类型，递归生成布局。
    • defined：查找用户定义类型的定义（IdlTypeDef），并调用 typeDefLayout 生成布局。
    • generic：解析泛型类型，并递归生成布局。
  • 异常处理：对未实现的类型抛出 IdlError。

---

3.2 typeDefLayout

• 功能：生成用户定义类型（struct 或 enum）的布局。
• 实现：
  • struct：
    • 遍历结构体的字段，递归调用 fieldLayout 生成每个字段的布局。
    • 使用 borsh.struct 将字段布局组合成结构体布局。
  • enum：
    • 遍历枚举的变体（variants），为每个变体生成布局。
    • 使用 borsh.rustEnum 将变体布局组合成枚举布局。
  • 类型别名（type）：
    • 如果是别名类型，递归调用 fieldLayout。
  • 异常处理：对未实现的用户定义类型抛出异常。

---

3.3 typeSize

• 功能：计算给定类型的字节大小。
• 实现：
  • 对于基本类型，直接返回固定大小（如 u8 为 1 字节，u64 为 8 字节）。
  • 对于复杂类型：
    • option：大小为 1 + 内部类型大小。
    • array：大小为 元素大小 * 长度。
    • defined：
      • 如果是 struct，累加所有字段的大小。
      • 如果是 enum，取所有变体中最大的大小，并加上 1 字节的标记位。
    • vec 和 string：返回 1，表示动态大小。
  • 泛型：解析泛型后递归计算大小。
  • 异常处理：对未实现的类型抛出异常。

---

3.4 resolveArrayLen

• 功能：解析数组的长度。
• 实现：
  • 如果长度是数字，直接返回。
  • 如果长度是泛型常量，从 genericArgs 中解析其值。
  • 如果解析失败，抛出异常。

---

3.5 resolveGenericArgs

• 功能：递归解析泛型类型，替换为实际的类型或值。
• 实现：
  • 遍历 typeDef 的泛型定义，将泛型替换为 genericArgs 中的实际类型。
  • 支持嵌套泛型（如 option<vec<T>>）。
  • 如果解析失败，返回 null 或抛出异常。

---

4. 代码特点

4.1 模块化和扩展性

• 代码通过递归和分支逻辑处理复杂类型，具有良好的扩展性。
• 新增类型时，只需在 fieldLayout 和 typeSize 中添加对应的逻辑。

4.2 错误处理

• 使用自定义的 IdlError 抛出异常，便于调试和定位问题。

4.3 性能优化

• 通过缓存类型定义和泛型解析结果，避免重复计算。
• 对枚举变体计算大小时，使用 Math.max 提高效率。

---

5. 应用场景

• 智能合约开发：
  • 在 Solana 等区块链平台中，IDL 用于定义智能合约的接口。IdlCoder 可以根据 IDL 自动生成数据的序列化和反序列化逻辑。
• 数据序列化：
  • 将复杂的数据结构序列化为二进制格式，用于网络传输或存储。
• 动态类型系统：
  • 支持动态类型和泛型，适用于需要灵活数据结构的系统。

---

6. 总结

这段代码实现了一个复杂且功能强大的 IDL 编码器，适用于处理区块链和智能合约中的数据序列化需求。通过结合 borsh 和递归解析逻辑，它能够高效地处理基本类型、复杂类型和泛型类型。代码设计清晰，具有良好的扩展性和可维护性。