440 lines
16 KiB
JavaScript
440 lines
16 KiB
JavaScript
|
|
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");
|
|
} |