WASM WAT compiler + R1CS constraint export for snarkjs - #415
Open
VanshSahay wants to merge 76 commits into
Open
Conversation
Compiles Clean's witness-generation IR to WAT for fast witness generation. Handles all IR expression types, let-steps, mapRange loops, multi-witness circuits with correct variable mapping via VarMap tracking.
VanshSahay
marked this pull request as ready for review
July 5, 2026 10:10
VanshSahay
changed the title
mitschabaude
requested changes
Jul 15, 2026
Collaborator
There was a problem hiding this comment.
this is really promising work!!
at a high-level, the following is missing for me:
- there seem to be several operations that are not supported or cases that don't work. that's perfectly fine, but the current implementation fails silently in each such case, which I don't find acceptable: it's actually quite hard to tell from the code what exactly will work. I want the compiler to return an error when processing inputs that are not supported. the acceptance criteria is that the code either fails with a clear reason or produces something correct.
- you should add documentation, for example a README.md inside Backends/Wasm. it should document how to run the compiler and produce the R1CS and WASM files, and how to use them with snarkjs. there should also be references to the snarkjs ABI you're targeting
Comment on lines
+56
to
+57
| | .block _ _ body => encodeBlock arr resolveCall 0x02 body | ||
| | .loop _ _ body => encodeBlock arr resolveCall 0x03 body |
Collaborator
There was a problem hiding this comment.
it seems wrong to model blocks and loops with a result type, but when encoding, ignoring the result type and writing an empty result type instead. decide for one of either
- we only model blocks without result for now
- we correctly encode all result types
Comment on lines
+54
to
+55
| | .br _ => putULEB128 (arr.push 0x0C) 0 | ||
| | .brIf _ => putULEB128 (arr.push 0x0D) 0 |
Collaborator
There was a problem hiding this comment.
wait, shouldn't the emitted instruction have something to do with the given label you break to? here you ignore the label, does that mean you hard-code breaking to the innermost parent scope? if yes then you should not make the blocks get a label at all. or (better), properly implement translation of labels into numbers
| let arr := match result with | ||
| | [] => arr.push 0x40 -- empty block type | ||
| | [t] => arr.push (vtOpc t) -- single result type | ||
| | _ => arr.push 0x40 -- multi-value: placeholder (needs type section ref) |
Collaborator
There was a problem hiding this comment.
this is another placeholder with currently incorrect behavior.. again, the best way to handle missing support is to not model multivalue returns in the type for now
Comment on lines
+68
to
+71
| | .memLoad .i32 off _ => putULEB128 (putULEB128 (arr.push 0x28) 2) off | ||
| | .memLoad .i64 off _ => putULEB128 (putULEB128 (arr.push 0x29) 3) off | ||
| | .memStore .i32 off _ => putULEB128 (putULEB128 (arr.push 0x36) 2) off | ||
| | .memStore .i64 off _ => putULEB128 (putULEB128 (arr.push 0x37) 3) off |
Collaborator
There was a problem hiding this comment.
you don't use the align value, looks incorrect to me?
| else List.range nw |>.foldl (fun cb' w => cb'.push (local.get (base + w))) cb | ||
|
|
||
| mutual | ||
| partial def compileFExpr (vm : VarMap) : FExpr F → CodeBuilder → CodeBuilder |
Collaborator
There was a problem hiding this comment.
would be nice to get rid of the partials here so that you prove the recursion terminates
| | op :: rest => flattenOp op ++ flattenOps rest | ||
|
|
||
| /-- Compile to a WASM Module. This is the main entry point. -/ | ||
| def compileModule (fieldPrime numInputs : ℕ) (ops : List (Operation F)) (numWords : ℕ := 1) : String := |
Collaborator
There was a problem hiding this comment.
it seems very easy to accidentally call this with the numWords = 1 default for any fieldPrime, which is a footgun
Comment on lines
+81
to
+89
| def genSingleWordArith (p : ℕ) : List Func := | ||
| let pVal : ℕ := p | ||
| let pm2 : ℕ := p - 2 | ||
| [ | ||
| { name := "$fadd" | ||
| params := [("", .i64), ("", .i64)] | ||
| results := [.i64] | ||
| body := [.localGet 0, .localGet 1, .binop .i64 .add, | ||
| .const .i64 pVal, .binop .i64 .rem_u] } |
Collaborator
There was a problem hiding this comment.
would be nicer to use the prettier instruction API you introduced just above this method
Comment on lines
+91
to
+95
| { name := "$fmul" | ||
| params := [("", .i64), ("", .i64)] | ||
| results := [.i64] | ||
| body := [.localGet 0, .localGet 1, .binop .i64 .mul, | ||
| .const .i64 pVal, .binop .i64 .rem_u] } |
Collaborator
There was a problem hiding this comment.
this method is incorrect for primes of bit size between 32 and 64. the multiplication of two 64-bit numbers will overflow the uint64 range before being reduced modulo the prime
| { computeFunc with name := "$witness", exportName := some "witness" }] | ||
| ++ abiFuncs | ||
| } | ||
| Module.toString module |
Collaborator
There was a problem hiding this comment.
can you compile to binary instead? it's the much nicer target IMO
Collaborator
There was a problem hiding this comment.
once that works, can you also remove the redundant WAT compiler? WASM -> WAT is probably not needed or can be done with separate tooling
Comment on lines
+26
to
+28
| let (lc, st1) := flattenExpr p vm e st | ||
| let constr : Constraint := (lc, [(0, 1)], []) | ||
| let st2 := { st1 with constraints := constr :: st1.constraints } |
Collaborator
There was a problem hiding this comment.
the way we generate constraints is pretty inefficient right now. for example, any R1CS-shaped constraint, even one as simple as x * y === z, will cause you to generate two constraints
Collaborator
|
@VanshSahay one feedback I forgot in #415 (review) is that there should be tests documenting both working and failing input circuits |
- Define ~60 named WASM opcodes in Binary.lean (eliminating raw hex) - Define memory/signal/layout constants in Compile.lean: alignmentI32/I64, bytesPerI32/I64, srwmBaseAddress, wasmPageSize, limbBits/limbModulus, low32Mask/hiWordShift, singleWordPrimeMax, getWitnessFixedLocals, r1csSignalOffset, snarkjs versions - Fix R1CS hardcoded n8, nOutputs, nPrvInputs (now computed/parameterized) - Fix putSLEB128: document it never receives negative values - Implement multi-word feq (pairwise limb compare + AND-reduce) - Implement multi-word ofNat (zero-extend to nw limbs, reduce via fadd) - Implement multi-word val (compile FExpr, keep lowest limb, drop rest) - Implement append in compileVExpr - Change interactions from error to skip (like lookups) in witness gen Co-Authored-By: Claude <noreply@anthropic.com>
…ed-zkEVM#26) - Add label field to ifElse AST constructor (matching block/loop) - Update WAT emitter, binary encoder, and all call sites - Remove numWords := 1 default from VarMap.init - Specs/Poseidon.lean: BN254_PRIME is abbrev (defeq), Fact uses CompPoly proof directly (reviewer comment Verified-zkEVM#1 was already addressed in prior commits)
Inline the goGeneral helper at each call site so Lean's structural termination checker can directly see that 'rest' is the tail of 'ops'.
- Bug A: multi-word feq shared temp-locals corrupt captured values when second operand contains nested ite. Fix by compiling a first (tmpBase free for nested temp use), capturing a to upper half, then compiling e to lower half. - Bug B: compileModuleBinary incorrectly computed witnessCount and startSignal using finalVm.nextLocal (includes let-step locals). Fix by using finalVarIdx from processFlatOps and finalVm.lookup for output-store local indices, matching the WAT emitter. - Update README: append, ofNat, val, feq are now supported; interactions skip in witness gen (not error).
Root cause: missing i64.lt_s (0x53) in encodeInstr relop table. The catch-all | .relop _ _ => arr silently dropped it, corrupting the stack before if in . Also fixed wrong opcode values for i64.lt_u (was 0x65, correct 0x54) and added all missing relop encodings with verified opcodes from wat2wasm output. Additional fix: output store local index in compileModuleBinary was using R1CS signal number (offset by 1) as circuit variable index. Now correctly uses numInputs+i for local lookup and 1+numInputs+i for signal address, matching WAT emitter. Tests: empty circuit, addOps (with let-steps), and Poseidon1 (BN254, 4-limb) all validate via compileModuleBinary → wasm-validate.
- Delete the WAT text compileModule (removes Module.toString dependency) - Rename compileModuleBinary to compileModule (only entry point) - Update all tests to use binary validation (wasm-validate + wasm2wat) - Remove 180-line duplicate between WAT and binary emitters
The Circom 2 ABI calls getInputSignalSize(hMSB, hLSB) to determine how many field elements an input signal expects. For our circuits, each input is a single scalar (not an array), so this must return 1. Previously returned n32/2 (= numWords), which snarkjs interpreted as requiring 4 field elements per input, triggering 'Not enough values' even for correct single-value inputs. Verified: Poseidon1 circuit compiles → wasm-validate → snarkjs wtns calculate produces 19KB valid witness file.
- Test Poseidon1 with inputs 0, 1, 5 — all produce different outputs confirming correct hash function behavior - Fix string escape syntax in test file for JSON generation
- Remove genFsub from genMultiWordArith: multi-word subtraction is never called by the compiler (only fadd/fmul/finv are emitted). Keeping it as compiled-but-unused code is misleading. - Add custom name section (section 0) to binary encoder, mapping function indices to their internal names. This makes wasm2wat output much more readable for debugging (e.g., '' instead of 'func[1]').
# Conflicts: # lakefile.lean
Merge origin/main (u64 IR Verified-zkEVM#442 + Lean 4.32.2 Verified-zkEVM#443), then adapt: - NExpr → U64Expr, compileNExpr → compileU64Expr - .ofNat → .ofU64 (drop now-vacuous 2^64 const guard) - Step.letN → Step.letU - Remove FExpr.envGet (deleted from IR); listGet/dataGet/hintGet still rejected with .error - New BExpr.flt: field-sorted < — i64.lt_u for single-word, nested-ifElse limb-wise unsigned compare for multi-word - New BExpr.bit: bit test via and/eqz on limb i/64, bit i%64 - New VExpr.envRange: witness env cells from WASM locals - New VExpr.bitsOf: compile x once into scratch, per-bit tests - Reserve 2*nw scratch locals at every allocation site so multi-word flt/feq/ite never overrun declared locals (fixes 'local variable out of range' + wrong results) - Drop CompPoly dependency (removed upstream); BN254 Fact instance restored as by-sorry in Circomlib/Poseidon.lean - Update tests: envGet→listGet rejection, letN→letU, new flt/bit/bitsOf/envRange correctness tests (snarkjs-verified) - Update README for the new IR constructors
Re-import CompPoly at v4.32.0 (Lean 4.32-compatible tag), which provides a kernel-verified Pratt certificate for the BN254 scalar field prime. Replace the by-sorry Fact instance with BN254.ScalarField_is_prime directly: - Clean/Specs/Poseidon.lean: BN254_PRIME is abbrev = BN254.scalarFieldSize, Fact instance uses the Pratt certificate proof - Clean/Circomlib/Poseidon.lean: remove the by-sorry instance (now inferred from Specs.Poseidon) - lakefile.lean: require CompPoly @ v4.32.0 Verified: no sorrys in Poseidon/WASM files; full build --wfail clean; Poseidon1 E2E (wasm-validate + snarkjs) still passes.
Switch multi-word field arithmetic to CIOS Montgomery multiplication
with 64-bit limbs (HAC 14.36), keeping the snarkjs ABI (32-byte elements):
- genFmul: full 64×64-limb product via schoolbook, then N Montgomery
reduction steps using n' = -p^{-1} mod 2^64 (Newton iteration),
single conditional subtraction (result < 2p)
- Keep everything in Montgomery form internally; convert at boundaries:
pushConst/pushCoeffF emit c·R mod p, loadSignal/toMont use montMul
by R², outputStores/fromMont use montMul by 1, inputs converted in
getWitness, constants R/R² precomputed in Lean
- genFadd unchanged (works on Montgomery forms); genFinv init is R mod p
so Fermat exponentiation yields the Montgomery-form inverse
- flt/bit/bitsOf convert operands back from Montgomery before comparing
- setInputSignal: write all nw limbs per element (SRWM is LSW-first)
Critical fix: Binary.lean binopOffset had shr_u/shr_s swapped, emitting
i64.shr_s where shr_u was intended — sign-extending in and
corrupting every 128-bit product. All Montgomery arithmetic depends on it.
Verified end-to-end: identity, x², x+5, finv, FullRound, and Poseidon1
all produce witnesses matching Lean ground truth via snarkjs (Poseidon1
output at signal 402 = poseidon1Opt for inputs 0, 1, 5).
…utSignalSize=numInputs
- scratchReserve returns 0 for nw=1: single-word expressions need no
scratch (.flt/.feq use direct i64 ops, .ite/.bit reuse the output
slot transiently). Previously each witness reserved 2nw extra locals,
inflating Keccak's 31K witnesses to 129K locals — over the WASM
50K limit ('local count too large'). Now Keccak fits in 30,914.
- getInputSignalSize returns numInputs: snarkjs validates the input
array length against it; circuits with multi-element inputs (e.g.
Keccak's 200 field elements) previously failed 'Too many values'.
- Verified Keccak-f[1600] end-to-end: 31,113 signals, all 25 output
lanes' limbs present in the witness (correct permutation of zero
state), compiled to 2.1MB WASM.
- Replace per-witness scratchReserve (which inflated local counts by 2nw per witness) with a SINGLE shared scratch region above all witness/step locals, computed by a pre-pass over the flat ops. Keccak-f[1600] (31K witnesses) now fits in 30,914 locals. - Implement listGet: compile the runtime u64 index to a temp local, emit a select-sum chain (elem_k if i==k else 0) — matches FExpr.eval semantics (out-of-range reads 0). - getInputSignalSize returns numInputs (was hardcoded 1, breaking multi-element inputs like Keccak's 200 field elements). Verified: Keccak-f[1600] of zero state produces the correct permutation (all 25 output lanes' limbs present in the witness); Poseidon1 still matches ground truth; SHA256Compress compiles but exceeds the WASM 50K-local limit at 80K 2-limb witnesses (documented limitation).
Addresses reviewer standards after the Montgomery/IR migration: Comments: - Module header: binary-only (no WAT), CIOS Montgomery description - montNPrime: fix docstring (returns n' = -p^-1 mod 2^64, not p^-1) - genFinv/genFadd: correct limb-order descriptions (lowest first, deepest); genFinv init comment matches pushCoeff(montR) - Remove stale Barrett references Dead code removed: - Multi-word genFsub (never emitted) and single-word $fsub (no callers — the IR has no subtraction), plus now-unused i64.lt_s - WAT text emitter in Ast.lean (Module/Func/Instr.toString) — binary-only - ifNone, numScratchLocals, r1csSignalOffset (unused) Named constants: - w*8 → bytesPerI64, w*4 → bytesPerI32, signalBaseRaw uses srwmBaseAddress, alignment uses bytesPerI64, w*32 → hiWordShift - toLimbs/pushCoeff/pushConst use limbBits/limbModulus - getMinorVersion/getPatchVersion use snarkjsMinor/PatchVersion Reviewer ask Verified-zkEVM#1 (no silent failures): - Binary.lean unop/relop catch-alls now emit unreachable (0x00) so unsupported instructions fail WASM validation loudly instead of silently emitting malformed code R1CS: error prefix processOps: (was compileR1CS:) README: binary-only, Montgomery arithmetic, correct Except types, compileR1CS signature, listGet supported, local-limit note
…and outputs-first layout
VanshSahay
force-pushed
the
wasm-ir-compiler-clean
branch
from
August 12, 2026 06:13
25f8f0d to
48ba979
Compare
|
Hi, this PR is extremely cool, thanks a lot! I'm using it as the example for Clean tutorial (work in progress at https://clean-intro.zksec.workers.dev/) However, I/Codex, uncovered some performance issues: https://gist.github.com/rot256/082e5f8311741da2b0be828329225157 TLDR: appending to an immutable list inside |
Address the reviewer's performance findings: - thread CodeBuilder (O(1) cons) through the six quadratic `acc ++ chunk` loops; VarMap.lookup is now the direct `i * numWords` formula, with let-steps in a dedicated `stepNext` region so the layout invariant is exact - write `e2 * 2` (not `e2 + e2`) for power-of-two accumulators in Num2Bits / Bits2Num / BinSub: flattenExpr/compileExpr are sharing-free structural recursions, so doubling a shared subtree was 2^N - add performance regression tests (2000-witness ops through WASM and R1CS, 3000-element mapRange, Num2Bits 128 gadget) and a README Performance section - CI: install snarkjs and wabt so the wasm-validate/snarkjs-backed checks run instead of SKIP
Contributor
Author
|
Hey @rot256 , thanks for the report. I've fixed the WASM-side exponential flattening one finding of yours is still open and it's the same acc ++ chunk pattern, |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
WASM backend: binary WASM + R1CS export for snarkjs
Compiles Clean circuits to snarkjs-compatible
.wasmwitness generators (Circom 2 ABI) and.r1csconstraint files. Verified end-to-end:snarkjs wtns calculate,snarkjs r1cs info, and the full Groth16 pipeline (setup → prove → verify).Architecture
Two compilers sharing a common flattening pass in
Compile.lean:Compile.lean→Ast.lean→Binary.lean): fully typed WASM AST with a LEB128 binary encoder (type/function/memory/export/code/data/name sections). Compiles witness-generation IR to WASM with the full snarkjs Circom 2 ABI (init,setInputSignal,getWitness,witness, ...). No.rawinstructions.R1CS.lean): constraint extraction shared with the WASM side, serialized both as pretty-printed JSON (compileR1CS) and as binary.r1csin the r1csfile format (compileR1CSBin, consumed directly bygroth16 setup; small primes are padded to snarkjs's field-word width).Field arithmetic
i64.mul/i64.rem_u; larger primes withnumWords = 1are rejected with an error.x·R mod p,R = 2^(N·64)):$mul64x64: 64×64→128 with carry detection$fmul: N×N schoolbook + CIOS Montgomery reduction (HAC Algorithm 14.36)$fadd: limb-wise addition with two-operand carry checks + conditional mod p$finv: Fermat square-and-multiply via$fmulite,flt,feq,bit,bitsOf,listGetwith multi-word support; intermediate signals, let-steps, and a shared scratch region; outputs-first signal layoutAPI
inputNames(one per input): strict FNV-1a input-key validation like circom — unknown keys are rejected ("Signal not found"). Empty = lenient mode.outputVarIdx: outputs-first signal layout (1, outputs…, inputs…, witnesses…) sogroth16 public.jsoncontains the circuit's real outputs and inputs.Except String; unsupported constructs fail at compile time with a message — no silent fallbacks.Verification
CleanTestschecks pass, including: Poseidon1 compiled to WASM, run throughsnarkjs wtns calculate, output matched against Lean ground truth (Specs.PoseidonOptimized.poseidon1Opt) on 3 inputs; every module passeswasm-validate; binary.r1csvalidated bysnarkjs r1cs info; negative tests (native witnesses, lookups,numWords = 1on BN254, unknown input keys) all rejected.powersoftau→groth16 setup→zkey export verificationkey→groth16 prove→groth16 verify= OK, withpublic.json = [poseidon1(0), input]correct under the outputs-first layout.$faddcarry-corner bug, shared-scratch sizing, multi-wordfeq/fltcodegen,listGetindex and Montgomery-form bugs, R1CS signal numbering sync, strict input names, outputs-first layout, thewitnessABI signature, and the data section.Verified circuits
FormalCircuitwith soundness/completeness proofs (BN254, 620 signals, 618 constraints)Clean/Examples/WasmDemo.leanClean/Utils/Test/TestWasmCompile.leanKey changes in this PR
$faddinputNames) and outputs-first signal layout (outputVarIdx).r1csexport (compileR1CSBin) in the r1csfile formatwitnessexport with the correct zero-arg ABI signature.rawinstructions; dynamic memory sizing; LEB128 binary encoderWasmDemo.leanwalkthrough; README rewriteKnown limitations (documented in the README)
init(sanityCheck)is accepted but ignored; nestedlistGetis rejected at compile time;nativewitnesses anddataGet/hintGetare not representable; lookups/interactions are rejected by R1CS export.Progress on #420.