wasmati: 你应该用 TypeScript 编写 WebAssembly - ZKSECURITY

如果你了解 WebAssembly 或 Wasm，你可能听说过一些不同的理由来解释为什么你需要它：

1. 它是一种用于 Web 的底层语言。因此，它比 JS 快。用它来加速事情。
2. 这是一个构建目标。因此，它使你能够在 Web 上使用 JS 以外的语言。你可以用 Rust 或 Go 编写一个 Web 库！
3. 你实际上不在 Web 上使用它，而是在后端使用！对于不受信任的代码来说，它是一个完美的 VM，因为它只能通过主机运行时提供的功能访问外部世界。

我感觉到第 2 和第 3 点是人们使用 Wasm 的主要原因，而第 1 点是我个人感到兴奋的原因。我发现 Wasm 可以给你带来比 JS 更大的速度提升，尤其是在浏览器中进行繁重的密码学操作时，这是我们在 zkSecurity 喜欢做的事情。请参阅本文末尾的真实示例。

因此，如果你像我一样，想使用 Wasm 来加速你的 JS，你将面临如何编写 Wasm 以及如何将其集成到你的堆栈中的问题。两年来，无论是在专业上还是在业余项目中，我都一直在思考这个问题，最终我提出了一个令我满意的解决方案。

该解决方案是 wasmati，一个用于在指令级别编写 Wasm 的 TypeScript 库。wasmati API 看起来与 Webassembly text format (WAT) 完全一样。这是一个例子：

import { i64, func, local } from "wasmati";

const myMultiply = func({ in: [i64, i64], out: [i64] }, ([x, y]) => {
  local.get(x); // 将输入 x 放入堆栈
  local.get(y); // 将输入 y 放入堆栈
  i64.mul(); // 从堆栈中弹出最后两个值，将它们相乘，将结果放入堆栈
});

作为参考，这将是一个等效的 WAT 代码片段：

(func $myMultiply (param $x i64) (param $y i64) (result i64)
  (local.get $x)
  (local.get $y)
  i64.mul
)

在你因为要处理像堆栈变量这样的底层细节的想法而尖叫逃跑之前，我必须强调一件事：手动编写 Wasm 是可行的，因为如果你使用它来提高性能，你不需要很多。JS 在 99% 的情况下都可以，为了使事情尽可能快，你只需要在库的关键位置撒上一些 Wasm（但是在那里，你可能非常需要它）。

与编写 .wat 文件形成鲜明对比的是，使用嵌入在 TS 中的 WAT 是我非常喜欢的事情。你可以使用原始指令，没有任何抽象阻止你编写最优化的代码。同时，拥有一个高级语言来组合指令可以释放编译时编程，稍后会详细介绍。

本文的其余部分是我创建 wasmati 的故事，设计时考虑的权衡，以及如何使用它的示例。最后，我们将实现有限域加法（许多密码学中的性能关键操作），并将我们的 wasmati 代码与使用 BigInt 的纯 JS 实现进行基准测试。

我为什么创建 wasmati

在我的日常工作中，我维护 SnarkyJS，这是一个 TypeScript 框架，可让你在 Web 浏览器中创建任意语句的 零知识证明。创建 zk 证明的重型机械位于 Rust 库 中。Rust 被编译为 Wasm，以将该机械引入 JS 运行时。顺便说一句，我们也使用相当数量的编译为 JS 的 OCaml。

SnarkyJS 的多语言架构植根于可靠的工程选择。Rust 是所有加密人才所在地，并且它是我参与其中的加密生态系统中占主导地位的语言。你可以在本机和 JS 运行时之间重用 Rust 代码库。尽管如此，作为负责 TypeScript 方面事情的开发人员，我对这种阻抗不匹配、糟糕的调试体验以及这种架构造成的总体复杂性产生了厌恶。我渴望拥有一种语言的简单性，只运行一个构建步骤，只维护一个依赖性文件，你懂的。当我 ctrl-click 一个函数签名时，我想看到该函数的源代码，而不是代替不透明的已编译 Wasm blob 的 .d.ts 文件。我不想上下文切换到 Rust 代码库中，在那里我首先必须涉水通过一个丑陋的 wasm-bindgen 粘合层，直到我找到实际的实现。

去年，我参加了 ZPrize，在一个名为 加速椭圆曲线运算和有限域算术（WASM） 的类别中。这意味着我有机会为速度至上的用例重新发明 JS 中的 Wasm 堆栈。我将我的很多空闲时间都投入到了这个奖项中，最终 获得了第二名：

比较 ZPrize 参与者采取的不同方法很有趣：

• 获胜团队 Yrrid 由 Niall Emmart 领导，他们用 C 语言从头开始编写代码，并 使用 clang 将其编译为 Wasm。他们使用 Java 自动生成有限域算术的 C 代码 使用 Java。
• 我的提交主要使用 JS，加上使用一个被黑客攻击的、类型错误的 JS 库生成的原始 Wasm，用于将 WAT 语法转储到字符串中。
• 获得第三名的团队使用了 Jordi Baylina 的 wasmcurves 的一个分支，它使用一个名为 wasmbuilder 的 JS 包生成 Wasm。这在精神上与 wasmati 相似，但没有 API 润色和类型。
• 获得第四名的团队使用了编译为 Wasm 的 Go（gnark）
• 获得第五名的团队使用了编译为 Wasm 的 Rust（arkworks）
• 获得第六名的团队使用了编译为 Wasm 的 C（MIRACL Core）
• 参考实现是 Rust（arkworks）

这可能是一个巧合，但有趣的是，与从 Go 或 Rust 编译 Wasm 并构建在现成的库上的“更干净”的解决方案相比，最好的解决方案都有些hacky。

另一个结论是，对于大多数代码来说，使用 JS 而不是 Wasm 不会损害性能。我的提交是唯一一个这样做的，仅使用 Wasm 进行低级算术，而其他提交则 100% 使用 Wasm。

一种可扩展的编写原始 WebAssembly 的方法

当我开始在 ZPrize 上工作时，我知道我想主要使用 JS，并撒上一些 Wasm。我已经在去年直接编写 Wasm 的 一个业余项目 中沉浸了自己，在那里我开始为 WAT 开发一个成熟的捆绑器，称为 watever。我当时的论文是基于 WAT 作为开发环境的。

在我开始尝试用于 ZPrize 的有限域乘法的不同变体后，很明显我错了：WAT 方法无法扩展。为什么？好吧，有限域算法由数百条指令组成，其中包含数十个烘焙常量，如果你要更改算法的某些参数，每个常量都必须更改。很明显，手动编写所有内容将是疯狂的。我需要的是使用元编程或编译时编程的能力：我需要代码来为我生成所有这些指令。

因此，编写 Wasm 的 TS 包的想法诞生了。对于这个奖项，我以快速而肮脏的方式完成了，但几个月后，我决定重做它，这次是正确地：一个类型良好、维护良好的库，支持完整的 WebAssembly 规范。它将为开发人员提供一个干净的 API，镜像 WAT 语法，并创建 Wasm 字节码，而无需编译步骤。Wasm 编码将使用 解析器-打印器组合器模式 在 TS 中从头开始实现，而无需像 wabt 或 binaryen 这样的重依赖项。

我很自豪地宣布，这个名为 wasmati 的库已经可以使用了！在我使用它重新实现了我所有的 ZPrize 代码时，它经历了足够的试用周期。

一个真实的例子：有限域运算

为了了解 wasmati，并理解编译时编程的概念，让我们做一个端到端的例子：我们将实现一些有限域运算。如果你不知道什么是有限域也没关系，因为我们只会做一些基本的事情并解释一切。

对于我们的目的而言，“域元素”是数字 0,…,p−1 之一，其中 p 是素数。在实际应用中，p 通常太大而无法容纳在 64 位中。在椭圆曲线密码术中，它通常约为 256 位。因此，在软件中，我们必须使用本机整数数组 x[0],…,x[n−1] 表示域元素 x。每个整数 x[i] 称为“字”或“limb”，并且被理解为具有一定的最大位数，我们称之为 w（“字长”或“limb 大小”）。大小 n 是固定的 —— 它取决于 p 和 limb 大小 w ——，因此域元素将存储为整数的静态大小数组。这是域元素 x 如何由其 limbs 表示的数学公式：

x=x[0]+x[1]2w+x[2]22w+…+x[n−1]2(n−1)w

对域元素的运算通常涉及遍历 n 个 limbs 的循环。

作为第一个代码示例，让我们尝试实现运算 x > y。我们将使用 wasmati 在 TypeScript 中编写这些代码；Wasm 指令足够直观，你可以当场掌握它们。如果你想一起编写代码，请使用 npm init -y 在空文件夹中创建一个新的 npm 项目，并安装 npm i wasmati typescript。创建一个最小的 tsconfig：

// tsconfig.json
{
  "compilerOptions": {
    "target": "esnext",
    "module": "nodenext",
    "strict": true
  }
}

另外，将 package.json 更改为 "type": "module"，因为我们不是野蛮人。

最后，创建一个包含以下内容的 greater.ts 文件。你可以随时使用 ts-node-esm greater.ts 执行该文件（假设你已全局安装 ts-node）。

import { i32, func } from "wasmati";

const isGreater = func({ in: [i32, i32], out: [i32] }, ([x, y]) => {
  // TODO
});

你也可以在 我们的 github 上找到完整的示例代码。

wasmati 基础

你可以使用 func 声明 Wasm 函数。func 构造函数接受一个描述函数签名的对象，以及一个你将在其中编写函数代码的回调。我们的函数接受两个 i32 作为输入 —— 它们是指向内存中 x 和 y 的指针。输出是另一个 i32，它是 1 或 0（表示 true 或 false），具体取决于 x 是否大于 y。Wasm 没有布尔值。

x > y 的算法是这样的：如果 x 的最高 limb 大于 y 的最高 limb，那么肯定 x > y，我们可以返回 true。

在我们的实现中，limbs 都是 32 位整数，而加载 i32 的 Wasm 指令是 i32.load()：

local.get(x); // 将内存地址 x 压入堆栈
i32.load({ offset: 0 }); // 从地址 x 处的内存加载一个 i32，偏移量为 0

请注意，Wasm 是一种基于堆栈的语言：每个操作都会从隐式堆栈中弹出其参数，然后将其结果再次推送到该堆栈。在此示例中，我们首先通过调用 local.get(x) 将 x （内存地址）推送到堆栈上。然后，i32.load() 弹出 x 并推送它从内存中加载的 i32。

wasmati 库允许你以命令式地方式“调用” Wasm 指令，例如 local.get() 和 i32.load()。调用指令会将它们添加到当前创建的 Wasm 函数的指令列表中。

像许多指令一样，i32.load() 具有一个参数，该参数确定其在编译时的行为：它采用一个可选的偏移量（以字节为单位），该偏移量在从该位置加载之前添加到堆栈参数 x 中。

为了使指令序列更易于阅读，WAT 有一种语法糖形式，可让你将它们分组为 S 表达式，以便指令的堆栈参数看起来像函数参数：

(i32.load offset=0 (local.get $x))

## 与以下内容相同

local.get $x
i32.load offset=0

我们允许你在 wasmati 中执行类似的操作。你可以写：

i32.load({ offset: 0 }, local.get(x));

这工作很简单：local.get(x) 在 i32.load() 之前执行，因此它在指令列表中排在第一位。作为副作用，这使你具有类型安全性。local.get() 使用 x 的类型（Local<i32>）来推断其返回类型为 StackVar<i32>；另一方面，i32.load() 仅接受 i32 类型的输入，因此如果你意外地使 x 成为不同的类型，则会收到类型错误。

由于我们在编写 Wasm 时需要大量的 local.get()，因此 wasmati 为你提供了另一项便利：你可以直接传入 x，并且会自动添加 local.get()。同样，这是类型安全的。

i32.load({ offset: 0 }, x);

请记住，我们想从内存中加载域元素的 limbs。上面的代码加载了 x[0], 因为这是我们存储在内存位置 x 的第一个 i32。下一个 limb x[1] 将是存储在 4 字节偏移量处的另一个 i32（因为 32 位 = 4 * 8 位 = 4 字节）：

i32.load({ offset: 4 }, x); // 加载 x[1]

要加载最终的 limb x[n-1], 我们需要一个指针运算。为了本文的其余部分起见，我们假设 n=9。那么 x[n-1] 位于地址 x 的 4(n−1)=4⋅8=32 字节偏移量处的内存中。

i32.load({ offset: 32 }, x); // 加载 x[n-1] 其中 n=9

在这一点上，你可能会提出抗议并说：这太疯狂了，我们为什么要硬编码这些字节偏移量？让我们创建一个帮助函数来加载第 i 个 limb，这样我们就不必再看这个指针算术了！你是对的，让我们这样做：

function loadLimb(x: Local<i32>, i: number) {
  return i32.load({ offset: 4 * i }, x);
}

在这一点上，我们正在做一些在编写原始 WAT 文件时不可能的事情：我们创建一个抽象，以帮助我们更轻松地编写指令，而不会以任何方式影响 Wasm 运行时。

我们的第一个 Wasm 函数

有了一种从内存中加载 x[i] 的便捷方法，让我们开始编写 isGreater()：如果 x[n-1] > y[n-1], 则返回 true。

import { i32, func, Local, if_, return_ } from "wasmati";

const n = 9; // limb 的数量

const isGreater = func({ in: [i32, i32], out: [i32] }, ([x, y]) => {
  loadLimb(x, n - 1); // 将 x[n-1] 放入堆栈
  loadLimb(y, n - 1); // 将 y[n-1] 放入堆栈
  i32.gt_s(); // 将 (x[n-1] > y[n-1]) 放入堆栈

  // 如果 (x[n-1] > y[n-1]) 则执行此 if 块
  if_(null, () => {
    i32.const(1); // 将 1 (true) 放入堆栈
    return_(); // 提前返回
  });

  // TODO: 剩余的逻辑
});

我在该代码上添加了注释，以便你了解每一行所做的事情。再次注意 Wasm 基于堆栈的特性：if_ 块期望条件（一个 i32）已经在堆栈上；它弹出该 i32，如果 i32 非零，则执行 if 分支。

当你调用 func() 时，wasmati 将急切地执行其回调并构造一个表示 Wasm 函数的对象。如果你现在运行你的文件，你将看到它抛出一个错误：

Error: expected i32 on the stack, got nothing
  ...
  at file:///<path>/greater.ts:5:19

堆栈跟踪将我们指向我们定义函数的行。这是因为 isGreater 的返回签名是 [i32]，但是在条件 if 块之后，我们不返回任何内容：到达函数末尾时，堆栈为空。为了匹配函数签名，让我们将 0 (false) 放在堆栈上作为最后一行：

  // TODO: 剩余的逻辑
  i32.const(0);
});

继续。到目前为止，如果最高的 x limb 大于最高的 y limb，我们的函数将返回 true。相反，如果最高的 y limb 大于，那么我们知道 x < y，我们可以立即返回 false。为了能够在 x[n-1] 和 y[n-1] 上执行第二次操作而无需再次从内存中加载它们，我们引入了两个局部变量，我们在其中存储这些值。局部变量在 func 构造函数的另一个参数中声明，并且它们也是回调的输入：

const isGreater = func(
  { in: [i32, i32], locals: [i32, i32], out: [i32] },
  ([x, y], [xi, yi]) => {
    // ...
  }
);

我们使用 local.set() 将这些局部变量设置为 loadLimb() 返回的值。然后，我们将它们传递给 i32.gt_s()，再次避免写出 local.get() 指令：

const isGreater = func(
  { in: [i32, i32], locals: [i32, i32], out: [i32] },
  ([x, y], [xi, yi]) => {
    // 设置 xi = x[n-1] 并且 yi = y[n-1]
    local.set(xi, loadLimb(x, n - 1));
    local.set(yi, loadLimb(y, n - 1));

    i32.gt_s(xi, yi); // 将 (xi > yi) 放入堆栈

    // ... 与之前一样 ...
  }
);

现在，我们可以重用 xi 和 yi 来测试 x[n-1] < y[n-1]，在这种情况下，我们可以立即返回 false：

const isGreater = func(
  { in: [i32, i32], locals: [i32, i32], out: [i32] },
  ([x, y], [xi, yi]) => {
    // 设置 xi = x[n-1] 并且 yi = y[n-1]
    local.set(xi, loadLimb(x, n - 1));
    local.set(yi, loadLimb(y, n - 1));

    // 如果 (xi > yi) 则返回 true
    i32.gt_s(xi, yi);
    if_(null, () => {
      i32.const(1);
      return_();
    });

    // 如果 (xi != yi) 则返回 false
    i32.ne(xi, yi);
    if_(null, () => {
      i32.const(0);
      return_();
    });

    // TODO 在返回之前进行更多逻辑
    i32.const(0); // 返回 0 (false)
  }
);

请注意，我们可以使用 x[n-1] ≠ y[n-1] 而不是 x[n-1] < y[n-1]，因为我们之前已经测试过 x[n-1] > y[n-1]。

如果 x[n-1] = y[n-1]，则两个 if 条件都不为真，我们必须查看第二个最高的 limb 以确定 x > y。

在这一点上，很明显我们可以为所有 limbs 复制相同的代码，直到在最终情况下 x[0] = y[0]，我们知道 x = y，因此我们返回 false。

出于性能原因，为每个 limb 重复所有指令正是我们想要的。几乎可以肯定，它的性能会比使用 loop 指令在 Wasm 中编写运行时循环更好。但是，我们不需要一个接一个地重复这些指令 —— 毕竟，我们正在使用 TS 来组合 Wasm 代码，因此我们可以只使用 TS for 循环来重复指令：

const isGreater = func(
  { in: [i32, i32], locals: [i32, i32], out: [i32] },
  ([x, y], [xi, yi]) => {
    for (let i = n - 1; i >= 0; i--) {
      // 设置 xi = x[i] 并且 yi = y[i]
      local.set(xi, loadLimb(x, i));
      local.set(yi, loadLimb(y, i));

      // 如果 (xi > yi) 则返回 true
      i32.gt_s(xi, yi);
      if_(null, () => {
        i32.const(1);
        return_();
      });

      // 如果 (xi != yi) 则返回 false
      i32.ne(xi, yi);
      if_(null, () => {
        i32.const(0);
        return_();
      });
    }

    // 贯穿的情况：如果 x = y 则返回 false
    i32.const(0);
  }
);

这里发生了一些有趣的事情：感觉我们同时用两种语言编写代码。运行时代码是 Wasm，但也有编译时逻辑，例如 for 循环，或者将 loadLimb() 分解为一个帮助程序，它是用 TS 编写的。

底层编程是编译时编程

我认为，如果你查看任何的低级、对性能至关重要的事物的实现，运行时编程和编译时编程之间的张力就会变得明显。在对性能至关重要的代码中，如果只是根据需要评估 if 语句或从内存中加载静态已知的常量，就会浪费时间。你需要尽可能多地硬编码内容。静态已知的常量应该烘焙到你的代码中，应该将依赖于静态已知参数的分支的解析烘焙到你的代码中。循环通常应该展开。

因此，此级别的代码几乎永远不会读起来很好。你将看到大量编译器注释、宏和内联汇编，与正常的运行时代码混合在一起。查看 arkworks 或 barrettenberg 中的有限域代码，以了解我的意思。在这些示例中，一方面有一个合适的运行时语言（Rust 和 C），另一方面有一个分层的“编译时语言”，它由注释和宏组成，并且往往相当丑陋。

另一方面是代码生成方案，其中好的一部分是编译时语言，它可以是某种脚本语言，而运行时部分只是表示为字符串的代码片段，由脚本语言拼接在一起，未类型化且不受欢迎。

很少能找到一个以可读和惯用的方式将两个编程层结合在一起的示例。我认为 wasmati 方法非常接近该理想。编译时语言是 TS。运行时语言是你编写指令的库，并且它与该库的设计一样好。对于 wasmati，我努力使 TS 类型反映 Wasm 具有的运行时类型：i64、i32 等。目标是编译时语言 (TS) 中的类型阻止你在运行时语言 (Wasm) 中编写无效代码。这有时是可能的，但并非总是如此，因为可悲的是，你无法以保留简单命令式 API 的方式在类型中编码对隐式堆栈的操作的有效性。在类型无法提供帮助的情况下，必须介入良好且尽可能本地化的错误消息。

当然，存在权衡。避免使用编译器意味着你正在接受编译器本来会处理的许多复杂性，例如内存管理。

Wasm 与 JS 相比有多快？

我仍然必须证明通过转移到 Wasm 我们获得了任何东西。JS 具有本机 bigint —— 我们真的会比它们有所改进吗？

为了找出答案，我们将实现有限域加法，并将性能与 JS bigint 进行比较。这是我们想要击败的基线：

// 将两个数字模 p 相加
function addBigint(x: bigint, y: bigint, p: bigint) {
  let z = x + y;
  return z < p ? z : z - p;
}

我们添加两个数字，假定它们是非负的并且小于模数 p。如果结果仍然小于 p，则它是有效的域元素，我们将其返回。否则，我们知道结果必须小于 2p，因此我们减去 p 一次以将其带回到 0,…,p−1 的范围内。

现在用 Wasm 实现相同的逻辑。我将在此处转储完整的实现，并告诉你它在高层次上做什么：

// 定义常量

// 模数
let p = 0x40000000000000000000000000000000224698fc094cf91b992d30ed00000001n;

let w = 29; // limb 大小
let wordMax = (1 << w) - 1;
let n = 9; // limb 的数量，因此 2^(n*w) > p

let P: Uint32Array = bigintToLimbs(p); // p 作为 n 个 limbs 的数组

/**
 * 添加存储在内存中的两个域元素
 */
const add = func(
  {
    in: [i32, i32, i32],
    locals: [i32],
    out: [],
  },
  ([z, x, y], [zi]) => {
    // z = x + y
    for (let i = 0; i < n; i++) {
      // zi = x[i] + y[i] + 进位
      i32.add(loadLimb(x, i), loadLimb(y, i));
      if (i > 0) i32.add(); // 添加进位
      local.set(zi);

      // 对 zi 执行进位并将其存储在 z[i] 中；
      // 进位位留在堆栈上供下一个 i 使用
      if (i < n - 1) i32.shr_s(zi, w); // 进位位
      storeLimb(z, i, i32.and(zi, wordMax));
    }

    // 如果 (z < p) 返回;
    isLower(z, zi, P);
    if_(null, () => return_());

    // z -= p
    for (let i = 0; i < n; i++) {
      // zi = z[i] - p[i] + 进位
      i32.sub(loadLimb(z, i), P[i]);
      if (i > 0) i32.add(); // 添加进位
      local.set(zi);

      // 对 zi 执行进位并将其存储在 z[i] 中；
      // 进位“位”（0 或 -1）留在堆栈上供下一个 i 使用
      if (i < n - 1) i32.shr_s(zi, w); // 进位“位”
      storeLimb(z, i, i32.and(zi, wordMax));
    }
  }
);

/**
 * 用于检查 x < y 的辅助函数，其中 x 存储在内存中，y 是
 * 存储为 n 个 limbs 数组的常量
 */
function isLower(x: Local<i32>, xi: Local<i32>, y: Uint32Array) {
  block({ out: [i32] }, (block) => {
    for (let i = n - 1; i >= 0; i--) {
      // 设置 xi = x[i]
      local.set(xi, loadLimb(x, i));

      // 如果 (xi < yi) 则返回 true
      i32.lt_s(xi, y[i]);
      if_(null, () => {
        i32.const(1);
        br(block);
      });

      // 如果 (xi != yi) 则返回 false
      i32.ne(xi, y[i]);
      if_(null, () => {
        i32.const(0);
        br(block);
      });
    }

    // 贯穿的情况：如果 z = p 则返回 false
    i32.const(0);
  });
}

function loadLimb(x: Local<i32>, i: number) {
  return i32.load({ offset: 4 * i }, x);
}

function storeLimb(x: Local<i32>, i: number, xi: Input<i32>) {
  return i32.store({ offset: 4 * i }, x, xi);
}

function bigintToLimbs(x: bigint, limbs = new Uint32Array(n)) {
  for (let i = 0; i < n; i++) {
    limbs[i] = Number(x & BigInt(wordMax));
    x >>= BigInt(w);
  }
  return limbs;
}

在开始时，我们声明我们的素数 p（作为一个 bigint），以及确定 limb 表示的参数 w 和 n。我们还调用一个方法 bigintToLimbs(p) 以将 p 存储为长度为 n 的 Uint32Array。它的 limbs 在我们的代码中用作硬编码常量。

底部的函数是辅助函数。storeLimb() 类似于 loadLimb()，用于将域元素 limb 存储回内存。

isLower() 与我们之前在 isGreater() 中具有的逻辑几乎相同，只是它的一个参数现在是一个以 limb 数组形式存储的常量域元素，例如 p。此外，isLower() 不是它自己的 func，而只是可以调用以添加到其他函数的代码内联块。它使用返回类型为 [i32] 的 block 指令，以内联方式实现提前返回逻辑。

add() 函数的工作方式大致如下：

• 有一个初始循环计算 z = x + y

  • 对于每个 limb，我们执行 z[i] = x[i] + y[i]
  • 在每个步骤中，我们对结果执行进位操作，使其恢复到 w 位的大小。进位位通过右移 z[i] >> w 找到，并且简单地留在堆栈上并在下一个循环迭代中添加。
  • 余数（没有进位位的 z[i]）通过与位掩码 2^w - 1 进行按位与运算找到，并存储在内存中。

• 我们测试是否 z < p，如果是，则返回

• 如果 z >= p，我们减去 p 得到 z - p

  • 此循环类似于加法循环，不同之处在于这次，其中一个输入是常量：i32.sub(loadLimb(z, i), P[i])。这里，P[i] 是一个 number。当将纯数字传递给指令时，wasmati会填充一个 i32.const(number) 指令。
  • 对于负值，进位的工作方式相同，只要我们使用带符号的右移 i32.shr_s 和 limb 大小 w<32 即可。在这种情况下，进位“位”为 0 或 -1。

使用和运行 Wasm 函数

剩下的唯一事情是将我们的函数编码为 Wasm 模块并运行它。与编写原始 WAT 相比，这是使用 wasmati 更方便的另一个步骤：

let module = Module({
  exports: { add, memory: memory({ min: 1 }) },
});
let { instance } = await module.instantiate();
let wasm = instance.exports;

只需声明模块的 exports 即可；任何传递依赖项以及导入都会自动收集。

Exports 是正确类型的 TS 函数：

wasm.add satisfies (z: number, x: number, y: number) => void;

但是，要在 bigints 上运行我们的 add() 函数，我们首先必须将这些 bigints 写入 wasm 内存以获取指针（我们的函数期望的 number）。最后，我们想将结果读回 bigint。我们有 helpers fromBigint() 和 toBigint() 来执行该 IO；你可以在完整的代码示例中看到它们。

哪个更快？

以下是进行基准测试的两个实现。我们只是在一个循环中将一个数字加到自身：

function benchWasm(x0: bigint, N: number) {
  let z = fromBigint(0n);
  let x = fromBigint(x0);
  for (let i = 0; i < N; i++) {
    wasm.add(z, z, x);
  }
  return toBigint(z);
}

function benchBigint(x: bigint, N: number) {
  let z = 0n;
  for (let i = 0; i < N; i++) {
    z = addBigint(z, x, p);
  }
  return z;
}

这就是我们运行它们的方式：

let x = randomField();
let N = 1e7;

// warm up
for (let i = 0; i < 5; i++) {
  benchWasm(x, N);
  benchBigint(x, N);
}

// benchmark
console.time("wasm");
let z0 = benchWasm(x, N);
console.timeEnd("wasm");

console.time("bigint");
let z1 = benchBigint(x, N);
console.timeEnd("bigint");

// make sure that results are the same
if (z0 !== z1) throw Error("wrong result");

输出几乎解决了 Wasm 与 bigint 的问题：

wasm: 103.538ms
bigint: 459.061ms

正如你所看到的，bigint 版本慢了 4 倍以上。如果性能很重要，那么花费额外的工作来编写优化的 Wasm 绝对是值得的。对于字段乘法，它比加法对性能的要求更高，我发现差异甚至更大，因为有更多的优化向量。

仅此而已！

我希望你喜欢这篇文章，并且它让你对向 TS 代码库添加底层 Wasm 产生了一些兴趣。在 zkSecurity，我们将继续改进 wasmati 并在此基础上构建，以提供快速且正确的密码算法 Wasm 实现，作为 zk 社区的公共产品。

如果你想利用我们的专业知识来审核你的加密堆栈，请通过 hello@zksecurity.xyz 与我们联系！

• 原文链接： blog.zksecurity.xyz/post...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～