ZK WASM Hacking notes

Links

• Tech setup
• Test files https://github.com/akashin/zkwasm/tree/main/data
• Generated ZK ASM: https://github.com/near/wasmtime/tree/zkasm-hackaton/cranelift/data

Investigations

• https://github.com/near/nearcore/issues/9531

Refactoring/Features

• https://github.com/near/nearcore/issues/9529

Example module

WASM Module

(module
 (import "env" "assert_eq" (func $assert_eq (param i32) (param i32)))
 (func $main
	(local $counter i32)
	(local.set $counter (i32.const 0))
	(block
	 (loop
		(local.set $counter
		 (i32.add
			(local.get $counter)
			(i32.const 1)))
		(br_if 1
		 (i32.eq
			(local.get $counter)
			(i32.const 10)
		 )
		)
		(br 0)
	 )
	)
	(local.get $counter)
	(i32.const 10)
	call $assert_eq)
(start $main))

Corresponding Cranelift IR

; Selected as wasm start function
function u0:1(i32 vmctx) fast {
    sig0 = (i32, i32, i32 vmctx) fast
    fn0 = u0:0 sig0

                                block0(v0: i32):
@0034                               v2 = iconst.i32 0
@003a                               jump block3(v2)  ; v2 = 0

                                block3(v3: i32):
@003e                               v4 = iconst.i32 1
@0040                               v5 = iadd v3, v4  ; v4 = 1
@0045                               v6 = iconst.i32 10
@0047                               v7 = icmp eq v5, v6  ; v6 = 10
@0047                               v8 = uextend.i32 v7
@0048                               brif v8, block2, block5

                                block5:
                                    v13 = iconst.i32 1
                                    v14 = iadd.i32 v3, v13  ; v13 = 1
@004a                               jump block3(v14)

                                block2:
                                    v10 = iconst.i32 1
                                    v11 = iadd.i32 v3, v10  ; v10 = 1
                                    v12 = iconst.i32 10
@0052                               call fn0(v11, v12, v0)  ; v12 = 10
@0054                               jump block1

                                block1:
@0054                               return
}

Cranelift ZK ASM output

start:
  zkPC + 2 => RR
  :JMP(function_1)
  :JMP(finalizeExecution)

function_1:
  SP + 1 => SP
  RR :MSTORE(SP)
  0 + 0 => A
  :JMP(L1_1)
L1_1:
  0 + 1 => B
  $ => A :ADD
  0 + 10 => B
  $ => B :EQ
  B :JMPNZ(L1_3)
  :JMP(L1_1)
L1_3:
  0 + 10 => B
  B :ASSERT
  $ => RR :MLOAD(SP)
  SP - 1 => SP
  :JMP(RR)

finalizeExecution:
  ${beforeLast()}  :JMPN(finalizeExecution)
                   :JMP(start)

WASM transpiler ZK ASM output

start:
	:JMP(function_1)
function_1:
	2 :MSTORE(SP++)
	3 :MSTORE(SP++)
	SP - 1 => SP
	$ => A: MLOAD(SP)
	SP - 1 => SP
	$ => B: MLOAD(SP)
	$ => A :ADD
	A :MSTORE(SP++)
	5 :MSTORE(SP++)
	SP - 1 => SP
	$ => A: MLOAD(SP)
	SP - 1 => SP
	$ => B: MLOAD(SP)
	B :ASSERT
finalizeExecution:
	${beforeLast()}  :JMPN(finalizeExecution)
                     :JMP(start)

---

Authoring backend

https://github.com/bytecodealliance/wasmtime/issues/4126

• Implement ZK::isa_builder
• Implement ZKBackend and wrap it
• Implement compile::compile<ZKBackend>
• Implement create_reg_env
  This is relatively simple.
• Implement MachInstEmit
  Also looks relatively simple
• Implement AArch64MachineDeps
  This looks like ABI that ISA provides, so a translation for some common ops
  10-20 methods, all nicely documented, so shouldn't be a big problem
• Implement Inst pretty-printer

2. Find an interpreter

• Compile it to WASM/Cranelift
• Produce small/efficient ZK ASM out of it

Done

• Fix warnings
• Make sure that CTX register is not clobbered
• Why do we use C as intermediate in 2 + 3 => C, C => A be?
• Add function calls
• Add global label numbering
• Fix CMP implementation
• Fix useless C => E call, clobber?
• Decrease number of registers needed for counters.wat
• Fix negative number offsets
• Fix stack shifting size
• Add indentation for anything but labels
• Fix stack growth direction
• Fix +SMTH numbers
• Rewrite STORE similar to LOAD
• Run cargo fmt
• Fix register naming
• Add zkasm preamble
• Understand why we need stack pushing
• CALL: Only handle ASSERT
• Understand why we use stack allocations
• Do we need a frame pointer?
• LOAD instruction integer formatting
• Rename Addi instruction
• Fix storing values to stack
• Fix loading values from stack
• Adjust stack size
• Implement Mov
• Simonas: Int32
• Andrei: Sink put_data
• Leo: Complete TODOs with println
• Copy RiscV backend as a blueprint
• TODO(akashin): Look into ZK ASM code generation
  • Seems a bit complicated https://gist.github.com/akashin/d6f1dcafeb496029b456bed5700ea9f4
  • https://github.com/0xPolygonHermez/zkasmcom/blob/main/src/compiler.js
• Give access to zkwasm repo https://github.com/akashin/zkwasm
• Fork wasmtime https://github.com/near/wasmtime
• Update wasmtime to latest