1// Copyright (c) 2018 Jordi Baylina
  2// License: LGPL-3.0+
  3//
  4
  5const Contract = require("./evmasm");
  6const { unstringifyBigInts } = require("ffjavascript").utils;
  7const Web3Utils = require("web3-utils");
  8
  9const { C:K, M } = unstringifyBigInts(require("./poseidon_constants.json"));
 10
 11const N_ROUNDS_F = 8;
 12const N_ROUNDS_P = [56, 57, 56, 60, 60, 63, 64, 63];
 13
 14function toHex256(a) {
 15    let S = a.toString(16);
 16    while (S.length < 64) S="0"+S;
 17    return "0x" + S;
 18}
 19
 20function createCode(nInputs) {
 21
 22    if (( nInputs<1) || (nInputs>8)) throw new Error("Invalid number of inputs. Must be 1<=nInputs<=8");
 23    const t = nInputs + 1;
 24    const nRoundsF = N_ROUNDS_F;
 25    const nRoundsP = N_ROUNDS_P[t - 2];
 26
 27    const C = new Contract();
 28
 29    function saveM() {
 30        for (let i=0; i<t; i++) {
 31            for (let j=0; j<t; j++) {
 32                C.push(toHex256(M[t-2][j][i]));
 33                C.push((1+i*t+j)*32);
 34                C.mstore();
 35            }
 36        }
 37    }
 38
 39    function ark(r) {   // st, q
 40        for (let i=0; i<t; i++) {
 41            C.dup(t); // q, st, q
 42            C.push(toHex256(K[t-2][r*t+i]));  // K, q, st, q
 43            C.dup(2+i); // st[i], K, q, st, q
 44            C.addmod(); // newSt[i], st, q
 45            C.swap(1 + i); // xx, st, q
 46            C.pop();
 47        }
 48    }
 49
 50    function sigma(p) {
 51        // sq, q
 52        C.dup(t);   // q, st, q
 53        C.dup(1+p); // st[p] , q , st, q
 54        C.dup(1);   // q, st[p] , q , st, q
 55        C.dup(0);   // q, q, st[p] , q , st, q
 56        C.dup(2);   // st[p] , q, q, st[p] , q , st, q
 57        C.dup(0);   // st[p] , st[p] , q, q, st[p] , q , st, q
 58        C.mulmod(); // st2[p], q, st[p] , q , st, q
 59        C.dup(0);   // st2[p], st2[p], q, st[p] , q , st, q
 60        C.mulmod(); // st4[p], st[p] , q , st, q
 61        C.mulmod(); // st5[p], st, q
 62        C.swap(1+p);
 63        C.pop();      // newst, q
 64    }
 65
 66    function mix() {
 67        C.label("mix");
 68        for (let i=0; i<t; i++) {
 69            for (let j=0; j<t; j++) {
 70                if (j==0) {
 71                    C.dup(i+t);      // q, newSt, oldSt, q
 72                    C.push((1+i*t+j)*32);
 73                    C.mload();      // M, q, newSt, oldSt, q
 74                    C.dup(2+i+j);    // oldSt[j], M, q, newSt, oldSt, q
 75                    C.mulmod();      // acc, newSt, oldSt, q
 76                } else {
 77                    C.dup(1+i+t);    // q, acc, newSt, oldSt, q
 78                    C.push((1+i*t+j)*32);
 79                    C.mload();      // M, q, acc, newSt, oldSt, q
 80                    C.dup(3+i+j);    // oldSt[j], M, q, acc, newSt, oldSt, q
 81                    C.mulmod();      // aux, acc, newSt, oldSt, q
 82                    C.dup(2+i+t);    // q, aux, acc, newSt, oldSt, q
 83                    C.swap(2);       // acc, aux, q, newSt, oldSt, q
 84                    C.addmod();      // acc, newSt, oldSt, q
 85                }
 86            }
 87        }
 88        for (let i=0; i<t; i++) {
 89            C.swap((t -i) + (t -i-1));
 90            C.pop();
 91        }
 92        C.push(0);
 93        C.mload();
 94        C.jmp();
 95    }
 96
 97
 98    // Check selector
 99    C.push("0x0100000000000000000000000000000000000000000000000000000000");
100    C.push(0);
101    C.calldataload();
102    C.div();
103    C.push(Web3Utils.keccak256(`poseidon(uint256[${nInputs}])`).slice(0, 10)); // poseidon(uint256[])
104    C.eq();
105    C.jmpi("start");
106    C.invalid();
107
108    C.label("start");
109
110    saveM();
111
112    C.push("0x30644e72e131a029b85045b68181585d2833e84879b9709143e1f593f0000001");  // q
113
114    // Load t values from the call data.
115    // The function has a single array param param
116    // [Selector (4)] [item1 (32)] [item2 (32)] ....
117    // Stack positions 0-nInputs.
118    for (let i=0; i<t; i++) {
119        C.push(0x04+(0x20*(nInputs-i)));
120        C.calldataload();
121    }
122
123    for (let i=0; i<nRoundsF+nRoundsP-1; i++) {
124        ark(i);
125        if ((i<nRoundsF/2) || (i>=nRoundsP+nRoundsF/2)) {
126            for (let j=0; j<t; j++) {
127                sigma(j);
128            }
129        } else {
130            sigma(0);
131        }
132        const strLabel = "aferMix"+i;
133        C._pushLabel(strLabel);
134        C.push(0);
135        C.mstore();
136        C.jmp("mix");
137        C.label(strLabel);
138    }
139
140    C.push(toHex256(K[t-2][(nRoundsF+nRoundsP-1)*t]));  // K, st, q
141    C.dup(t+1); // q, K, st, q
142    C.swap(2);  // st[0], K, q, st\st[0]
143    C.addmod();  // st q
144
145    sigma(0);
146
147    C.push("0x00");
148    C.mstore();     // Save it to pos 0;
149    C.push("0x20");
150    C.push("0x00");
151    C.return();
152
153    mix();
154
155    return C.createTxData();
156}
157
158function generateABI(nInputs) {
159    return [
160        {
161            "constant": true,
162            "inputs": [
163                {
164                    "internalType": `uint256[${nInputs}]`,
165                    "name": "input",
166                    "type": `uint256[${nInputs}]`
167                }
168            ],
169            "name": "poseidon",
170            "outputs": [
171                {
172                    "internalType": "uint256",
173                    "name": "",
174                    "type": "uint256"
175                }
176            ],
177            "payable": false,
178            "stateMutability": "pure",
179            "type": "function"
180        }
181    ];
182}
183
184module.exports.generateABI = generateABI;
185module.exports.createCode = createCode;
186
187