circomlibjs-hinkal-fork 0.0.1

package/src/poseidon_wasm.js

@@ -0,0 +1,440 @@
+
+import poseidonConstants from "./poseidon_constants_opt.js";
+
+import { getCurveFromName, Scalar, F1Field } from "ffjavascript";
+
+export async function buildPoseidon() {
+    const bn128 = await getCurveFromName("bn128", true, buildPoseidonWasm);
+
+    const F = bn128.Fr;
+
+    const pState = bn128.tm.alloc(32);
+    const pIn = bn128.tm.alloc(32*16);
+    const pOut = bn128.tm.alloc(32*17);
+
+    const poseidon = (arr, state, nOut) => {
+        let buff;
+        let n;
+        if (Array.isArray(arr)) {
+            n = arr.length;
+            buff = new Uint8Array(n*32);
+            for (let i=0; i<n; i++) buff.set(F.e(arr[i]), i*32);
+        } else {
+            buff = arr;
+            n = buff.byteLength / 32;
+            if (n*32 != buff.byteLength) throw new Error("Invalid iput buff size");
+        }
+        bn128.tm.setBuff(pIn, buff);
+
+        if ((n<1)||(n>16)) throw new Error("Invalid poseidon size");
+
+        if (typeof state == "undefined") {
+            state = F.zero;
+        } else {
+            state = F.e(state);
+        }
+        bn128.tm.setBuff(pState, state);
+        nOut = nOut || 1;
+
+        bn128.tm.instance.exports.poseidon(pState, pIn, n, pOut, nOut);
+        if (nOut == 1) {
+            return bn128.tm.getBuff(pOut, 32);
+        } else {
+            const out = [];
+            for (let i=0; i<nOut; i++) {
+                out.push(bn128.tm.getBuff(pOut+i*32, 32));
+            }
+            return out
+        }
+    }
+
+    poseidon.F = F;
+    return poseidon;
+}
+
+
+export function buildPoseidonWasm(module) {
+
+    const F = new F1Field(Scalar.e("21888242871839275222246405745257275088548364400416034343698204186575808495617"));
+
+    const pointers = {
+        C: [],
+        S: [],
+        M: [],
+        P: [],
+    }
+    const N_ROUNDS_P = [56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68];
+    const NSizes = poseidonConstants.C.length;
+    const buffIdx = new Uint8Array(NSizes*5*4);
+    const buffIdx32 = new Uint32Array(buffIdx.buffer);
+    for (let i=0; i<NSizes; i++) {
+        buffIdx32[i*5] = N_ROUNDS_P[i];
+        const buffC = new Uint8Array(32*poseidonConstants.C[i].length);
+        for (let j=0; j<poseidonConstants.C[i].length; j++) {
+            F.toRprLEM(buffC, j*32, F.e(poseidonConstants.C[i][j]))
+        }
+        buffIdx32[i*5 + 1] = module.alloc(buffC);
+
+        const buffS = new Uint8Array(32*poseidonConstants.S[i].length);
+        for (let j=0; j<poseidonConstants.S[i].length; j++) {
+            F.toRprLEM(buffS, j*32, F.e(poseidonConstants.S[i][j]))
+        }
+        buffIdx32[i*5 + 2]  = module.alloc(buffS);
+
+        const N = poseidonConstants.M[i].length;
+        const buffM = new Uint8Array(32*N*N);
+        for (let j=0; j<N; j++) {
+            for (let k=0; k<N; k++) {
+                F.toRprLEM(buffM, (j*N+k)*32, F.e(poseidonConstants.M[i][k][j]));
+            }
+        }
+        buffIdx32[i*5 + 3]  = module.alloc(buffM);
+
+        const buffP = new Uint8Array(32*N*N);
+        for (let j=0; j<N; j++) {
+            for (let k=0; k<N; k++) {
+                F.toRprLEM(buffP, (j*N+k)*32, F.e(poseidonConstants.P[i][k][j]));
+            }
+        }
+        buffIdx32[i*5 + 4]  = module.alloc(buffP);
+    }
+
+    const pConstants = module.alloc(buffIdx);
+    const pState = module.alloc(32*((NSizes+1)*32));
+
+    function buildAddConstant() {
+        const f = module.addFunction("poseidon_addConstant");
+        f.addParam("t", "i32");
+        f.addParam("pC", "i32");
+        f.setReturnType("i32");
+        f.addLocal("i", "i32");
+        f.addLocal("pState", "i32");
+
+        const c = f.getCodeBuilder();
+
+        f.addCode(
+            c.setLocal("pState", c.i32_const(pState)),
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("t") )),
+                c.call("frm_add", c.getLocal("pC"), c.getLocal("pState"), c.getLocal("pState")),
+                c.setLocal("pC", c.i32_add(c.getLocal("pC"), c.i32_const(32))),
+                c.setLocal("pState", c.i32_add(c.getLocal("pState"), c.i32_const(32))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+            c.ret(c.getLocal("pC"))
+        );
+    }
+
+    function buildPower5() {
+
+        const f = module.addFunction("poseidon_power5");
+        f.addParam("p", "i32");
+
+        const c = f.getCodeBuilder();
+
+        const AUX = c.i32_const(module.alloc(32));
+
+        f.addCode(
+            c.call("frm_square", c.getLocal("p"), AUX),
+            c.call("frm_square", AUX, AUX),
+            c.call("frm_mul", c.getLocal("p"), AUX, c.getLocal("p")),
+        );
+
+    }
+
+    function buildPower5all() {
+        const f = module.addFunction("poseidon_power5all");
+        f.addParam("t", "i32");
+        f.addLocal("i", "i32");
+        f.addLocal("pState", "i32");
+
+        const c = f.getCodeBuilder();
+
+        f.addCode(
+            c.setLocal("pState", c.i32_const(pState)),
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("t") )),
+                c.call("poseidon_power5", c.getLocal("pState")),
+                c.setLocal("pState", c.i32_add(c.getLocal("pState"), c.i32_const(32))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            ))
+        );
+    }
+
+
+    function buildApplyMatrix() {
+        const f = module.addFunction("poseidon_applyMatrix");
+        f.addParam("t", "i32");
+        f.addParam("pM", "i32");
+        f.addLocal("i", "i32");
+        f.addLocal("j", "i32");
+        f.addLocal("pState", "i32");
+        f.addLocal("pStateAux", "i32");
+
+
+        const c = f.getCodeBuilder();
+
+        const pStateAux = module.alloc(32*((NSizes+1)*32));
+
+        const pAux = module.alloc(32);
+
+        f.addCode(
+            c.setLocal("pStateAux", c.i32_const(pStateAux)),
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("t") )),
+                c.call("frm_zero", c.getLocal("pStateAux")),
+                c.setLocal("pState", c.i32_const(pState)),
+                c.setLocal("j", c.i32_const(0)),
+                c.block(c.loop(
+                    c.br_if(1, c.i32_eq ( c.getLocal("j"), c.getLocal("t") )),
+                    c.call(
+                        "frm_mul",
+                        c.getLocal("pState"),
+                        c.getLocal("pM"),
+                        c.i32_const(pAux)
+                    ),
+                    c.call(
+                        "frm_add",
+                        c.i32_const(pAux),
+                        c.getLocal("pStateAux"),
+                        c.getLocal("pStateAux"),
+                    ),
+                    c.setLocal("pM", c.i32_add(c.getLocal("pM"), c.i32_const(32))),
+                    c.setLocal("pState", c.i32_add(c.getLocal("pState"), c.i32_const(32))),
+                    c.setLocal("j", c.i32_add(c.getLocal("j"), c.i32_const(1))),
+                    c.br(0)
+                )),
+                c.setLocal("pStateAux", c.i32_add(c.getLocal("pStateAux"), c.i32_const(32))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+            c.setLocal("pStateAux", c.i32_const(pStateAux)),
+            c.setLocal("pState", c.i32_const(pState)),
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("t") )),
+                c.call("frm_copy", c.getLocal("pStateAux"), c.getLocal("pState")),
+                c.setLocal("pState", c.i32_add(c.getLocal("pState"), c.i32_const(32))),
+                c.setLocal("pStateAux", c.i32_add(c.getLocal("pStateAux"), c.i32_const(32))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            ))
+        );
+    }
+
+    function buildApplySMatrix() {
+        const f = module.addFunction("poseidon_applySMatrix");
+        f.addParam("t", "i32");
+        f.addParam("pS", "i32");
+        f.setReturnType("i32");
+        f.addLocal("i", "i32");
+        f.addLocal("pState", "i32");
+
+        const c = f.getCodeBuilder();
+
+        const pS0 = module.alloc(32);
+        const pAux = module.alloc(32);
+
+        f.addCode(
+            c.call("frm_zero", c.i32_const(pS0)),
+            c.setLocal("pState", c.i32_const(pState)),
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("t") )),
+                c.call(
+                    "frm_mul",
+                    c.getLocal("pState"),
+                    c.getLocal("pS"),
+                    c.i32_const(pAux),
+                ),
+                c.call(
+                    "frm_add",
+                    c.i32_const(pS0),
+                    c.i32_const(pAux),
+                    c.i32_const(pS0),
+                ),
+                c.setLocal("pS", c.i32_add(c.getLocal("pS"), c.i32_const(32))),
+                c.setLocal("pState", c.i32_add(c.getLocal("pState"), c.i32_const(32))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+
+            c.setLocal("pState", c.i32_const(pState+32)),
+            c.setLocal("i", c.i32_const(1)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("t") )),
+                c.call(
+                    "frm_mul",
+                    c.i32_const(pState),
+                    c.getLocal("pS"),
+                    c.i32_const(pAux),
+                ),
+                c.call(
+                    "frm_add",
+                    c.getLocal("pState"),
+                    c.i32_const(pAux),
+                    c.getLocal("pState"),
+                ),
+                c.setLocal("pS", c.i32_add(c.getLocal("pS"), c.i32_const(32))),
+                c.setLocal("pState", c.i32_add(c.getLocal("pState"), c.i32_const(32))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+            c.call("frm_copy", c.i32_const(pS0), c.i32_const(pState)),
+            c.ret(c.getLocal("pS"))
+        );
+    }
+
+    function buildPoseidon() {
+        const f = module.addFunction("poseidon");
+        f.addParam("pInitState", "i32");
+        f.addParam("pIn", "i32");
+        f.addParam("n", "i32");
+        f.addParam("pOut", "i32");
+        f.addParam("nOut", "i32");
+        f.addLocal("pC", "i32");
+        f.addLocal("pS", "i32");
+        f.addLocal("pM", "i32");
+        f.addLocal("pP", "i32");
+        f.addLocal("t", "i32");
+        f.addLocal("i", "i32");
+        f.addLocal("nRoundsP", "i32");
+        f.addLocal("pAux", "i32");
+
+        const c = f.getCodeBuilder();
+
+        f.addCode(
+            c.setLocal(
+                "t",
+                c.i32_add(c.getLocal("n"), c.i32_const(1))
+            ),
+            c.setLocal(
+                "pAux",
+                c.i32_add(
+                    c.i32_const(pConstants),
+                    c.i32_mul(
+                        c.i32_sub(c.getLocal("n"), c.i32_const(1)),
+                        c.i32_const(20)
+                    )
+                )
+            ),
+            c.setLocal(
+                "nRoundsP",
+                c.i32_load(c.getLocal("pAux"))
+            ),
+            c.setLocal(
+                "pC",
+                c.i32_load(c.i32_add(c.getLocal("pAux"), c.i32_const(4)))
+            ),
+            c.setLocal(
+                "pS",
+                c.i32_load(c.i32_add(c.getLocal("pAux"), c.i32_const(8)))
+            ),
+            c.setLocal(
+                "pM",
+                c.i32_load(c.i32_add(c.getLocal("pAux"), c.i32_const(12)))
+            ),
+            c.setLocal(
+                "pP",
+                c.i32_load(c.i32_add(c.getLocal("pAux"), c.i32_const(16)))
+            ),
+
+            // Initialize state
+            c.call("frm_zero", c.i32_const(pState)),
+            c.call(
+                "frm_copy",
+                c.getLocal("pInitState"),
+                c.i32_const(pState),
+            ),
+            c.setLocal("i", c.i32_const(1)),
+            c.block(c.loop(
+                c.call(
+                    "frm_copy",
+                    c.i32_add(
+                        c.getLocal("pIn"),
+                        c.i32_mul(c.i32_sub(c.getLocal("i"), c.i32_const(1)), c.i32_const(32))
+                    ),
+                    c.i32_add(
+                        c.i32_const(pState),
+                        c.i32_mul(c.getLocal("i"), c.i32_const(32))
+                    )
+                ),
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("n") )),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+
+            // Initialize state
+            c.setLocal("pC", c.call("poseidon_addConstant", c.getLocal("t"), c.getLocal("pC"))),
+            // First full rounds
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.i32_const(3) )),
+                c.call("poseidon_power5all", c.getLocal("t")),
+                c.setLocal("pC", c.call("poseidon_addConstant", c.getLocal("t"), c.getLocal("pC"))),
+                c.call("poseidon_applyMatrix", c.getLocal("t"), c.getLocal("pM")),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+
+            c.call("poseidon_power5all", c.getLocal("t")),
+            c.setLocal("pC", c.call("poseidon_addConstant", c.getLocal("t"), c.getLocal("pC"))),
+            c.call("poseidon_applyMatrix", c.getLocal("t"), c.getLocal("pP")),
+
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("nRoundsP") )),
+                c.call("poseidon_power5", c.i32_const(pState)),
+                c.call("frm_add", c.i32_const(pState), c.getLocal("pC"), c.i32_const(pState)),
+                c.setLocal("pC", c.i32_add(c.getLocal("pC"), c.i32_const(32))),
+                c.setLocal("pS", c.call("poseidon_applySMatrix", c.getLocal("t"), c.getLocal("pS"))),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.i32_const(3) )),
+                c.call("poseidon_power5all", c.getLocal("t")),
+                c.setLocal("pC", c.call("poseidon_addConstant", c.getLocal("t"), c.getLocal("pC"))),
+                c.call("poseidon_applyMatrix", c.getLocal("t"), c.getLocal("pM")),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+            c.call("poseidon_power5all", c.getLocal("t")),
+            c.call("poseidon_applyMatrix", c.getLocal("t"), c.getLocal("pM")),
+
+            c.setLocal("i", c.i32_const(0)),
+            c.block(c.loop(
+                c.br_if(1, c.i32_eq ( c.getLocal("i"), c.getLocal("nOut") )),
+                c.call("frm_copy",
+                    c.i32_add(
+                        c.i32_const(pState),
+                        c.i32_mul(c.getLocal("i"), c.i32_const(32))
+                    ),
+                    c.i32_add(
+                        c.getLocal("pOut"),
+                        c.i32_mul(c.getLocal("i"), c.i32_const(32))
+                    )
+                ),
+                c.setLocal("i", c.i32_add(c.getLocal("i"), c.i32_const(1))),
+                c.br(0)
+            )),
+
+        );
+    }
+
+    buildAddConstant();
+    buildPower5();
+    buildPower5all();
+    buildApplyMatrix();
+    buildApplySMatrix();
+
+    buildPoseidon();
+    module.exportFunction("poseidon");
+}