1const bn128 = require("snarkjs").bn128;
  2const bigInt = require("snarkjs").bigInt;
  3const createBlakeHash = require("blake-hash");
  4const assert = require("assert");
  5const F = bn128.Fr;
  6
  7const SEED = "poseidon";
  8const NROUNDSF = 8;
  9const NROUNDSP = 57;
 10const T = 6;
 11
 12function getPseudoRandom(seed, n) {
 13    const res = [];
 14    let h = createBlakeHash("blake256").update(seed).digest();
 15    while (res.length<n) {
 16        const n = F.affine(bigInt.leBuff2int(h));
 17        res.push(n);
 18        h = createBlakeHash("blake256").update(h).digest();
 19    }
 20
 21    return res;
 22}
 23
 24function allDifferent(v) {
 25    for (let i=0; i<v.length; i++) {
 26        if (v[i].isZero()) return false;
 27        for (let j=i+1; j<v.length; j++) {
 28            if (v[i].equals(v[j])) return false;
 29        }
 30    }
 31    return true;
 32}
 33
 34exports.getMatrix = (t, seed, nRounds) => {
 35    if (typeof seed === "undefined") seed = SEED;
 36    if (typeof nRounds === "undefined") nRounds = NROUNDSF + NROUNDSP;
 37    if (typeof t === "undefined") t = T;
 38    let nonce = "0000";
 39    let cmatrix = getPseudoRandom(seed+"_matrix_"+nonce, t*2);
 40    while (!allDifferent(cmatrix)) {
 41        nonce = (Number(nonce)+1)+"";
 42        while(nonce.length<4) nonce = "0"+nonce;
 43        cmatrix = getPseudoRandom(seed+"_matrix_"+nonce, t*2);
 44    }
 45
 46    const M = new Array(t);
 47    for (let i=0; i<t; i++) {
 48        M[i] = new Array(t);
 49        for (let j=0; j<t; j++) {
 50            M[i][j] = F.affine(F.inverse(F.sub(cmatrix[i], cmatrix[t+j])));
 51        }
 52    }
 53    return M;
 54};
 55
 56exports.getConstants = (t, seed, nRounds) => {
 57    if (typeof seed === "undefined") seed = SEED;
 58    if (typeof nRounds === "undefined") nRounds = NROUNDSF + NROUNDSP;
 59    if (typeof t === "undefined") t = T;
 60    const cts = getPseudoRandom(seed+"_constants", nRounds);
 61    return cts;
 62};
 63
 64function ark(state, c) {
 65    for (let j=0; j<state.length; j++ ) {
 66        state[j] = F.add(state[j], c);
 67    }
 68}
 69
 70function sigma(a) {
 71    return F.mul(a, F.square(F.square(a,a)));
 72}
 73
 74function mix(state, M) {
 75    const newState = new Array(state.length);
 76    for (let i=0; i<state.length; i++) {
 77        newState[i] = F.zero;
 78        for (let j=0; j<state.length; j++) {
 79            newState[i] = F.add(newState[i], F.mul(M[i][j], state[j]) );
 80        }
 81    }
 82    for (let i=0; i<state.length; i++) state[i] = newState[i];
 83}
 84
 85exports.createHash = (t, nRoundsF, nRoundsP, seed) => {
 86
 87    if (typeof seed === "undefined") seed = SEED;
 88    if (typeof nRoundsF === "undefined") nRoundsF = NROUNDSF;
 89    if (typeof nRoundsP === "undefined") nRoundsP = NROUNDSP;
 90    if (typeof t === "undefined") t = T;
 91
 92    assert(nRoundsF % 2 == 0);
 93    const C = exports.getConstants(t, seed, nRoundsF + nRoundsP);
 94    const M = exports.getMatrix(t, seed, nRoundsF + nRoundsP);
 95    return function(inputs) {
 96        let state = [];
 97        assert(inputs.length < t);
 98        assert(inputs.length > 0);
 99        for (let i=0; i<inputs.length; i++) state[i] = bigInt(inputs[i]);
100        for (let i=inputs.length; i<t; i++) state[i] = F.zero;
101
102        for (let i=0; i< nRoundsF + nRoundsP; i++) {
103            ark(state, C[i]);
104            if ((i<nRoundsF/2) || (i >= nRoundsF/2 + nRoundsP)) {
105                for (let j=0; j<t; j++) state[j] = sigma(state[j]);
106            } else {
107                state[0] = sigma(state[0]);
108            }
109            mix(state, M);
110        }
111        return F.affine(state[0]);
112    };
113};
114
115