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