# ZK WASM Hacking notes

## Links

- [Tech setup](https://hackmd.io/xlfXSX0qQGe2npk4S0Kf3A?view)
- 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
- [x] Fix warnings
- [x] Make sure that CTX register is not clobbered
- [x] Why do we use `C` as intermediate in `2 + 3 => C, C => A` be?
- [x] Add function calls
- [x] Add global label numbering
- [x] Fix CMP implementation
- [x] Fix useless `C => E` call, clobber?
- [x] Decrease number of registers needed for counters.wat
- [x] Fix negative number offsets
- [x] Fix stack shifting size
- [x] Add indentation for anything but labels
- [x] Fix stack growth direction
- [x] Fix +SMTH numbers
- [x] Rewrite STORE similar to LOAD
- [x] Run `cargo fmt`
- [x] Fix register naming
- [x] Add zkasm preamble
- [x] Understand why we need stack pushing
- [x] CALL: Only handle ASSERT
- [x] Understand why we use stack allocations
- [x] Do we need a frame pointer?
- [x] LOAD instruction integer formatting
- [x] Rename Addi instruction
- [x] Fix storing values to stack
- [x] Fix loading values from stack
- [x] Adjust stack size
- [x] Implement Mov
- [x] Simonas: Int32
- [x] Andrei: Sink put_data
- [X] Leo: Complete TODOs with println
- [x] Copy RiscV backend as a blueprint
- [x] 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
- [X] Give access to zkwasm repo https://github.com/akashin/zkwasm
- [X] Fork wasmtime https://github.com/near/wasmtime
- [X] Update wasmtime to latest