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