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 C = new Contract();
 27
 28    function saveM() {
 29        for (let i=0; i<t; i++) {
 30            for (let j=0; j<t; j++) {
 31                C.push(toHex256(M[t-2][j][i]));
 32                C.push((1+i*t+j)*32);
 33                C.mstore();
 34            }
 35        }
 36    }
 37
 38    function ark(r) {   // st, q
 39        for (let i=0; i<t; i++) {
 40            C.dup(t); // q, st, q
 41            C.push(toHex256(K[t-2][r*t+i]));  // K, q, st, q
 42            C.dup(2+i); // st[i], K, q, st, q
 43            C.addmod(); // newSt[i], st, q
 44            C.swap(1 + i); // xx, st, q
 45            C.pop();
 46        }
 47    }
 48
 49    function sigma(p) {
 50        // sq, q
 51        C.dup(t);   // q, st, q
 52        C.dup(1+p); // st[p] , q , st, q
 53        C.dup(1);   // q, st[p] , q , st, q
 54        C.dup(0);   // q, q, st[p] , q , st, q
 55        C.dup(2);   // st[p] , q, q, st[p] , q , st, q
 56        C.dup(0);   // st[p] , st[p] , q, q, st[p] , q , st, q
 57        C.mulmod(); // st2[p], q, st[p] , q , st, q
 58        C.dup(0);   // st2[p], st2[p], q, st[p] , q , st, q
 59        C.mulmod(); // st4[p], st[p] , q , st, q
 60        C.mulmod(); // st5[p], st, q
 61        C.swap(1+p);
 62        C.pop();      // newst, q
 63    }
 64
 65    function mix() {
 66        C.label("mix");
 67        for (let i=0; i<t; i++) {
 68            for (let j=0; j<t; j++) {
 69                if (j==0) {
 70                    C.dup(i+t);      // q, newSt, oldSt, q
 71                    C.push((1+i*t+j)*32);
 72                    C.mload();      // M, q, newSt, oldSt, q
 73                    C.dup(2+i+j);    // oldSt[j], M, q, newSt, oldSt, q
 74                    C.mulmod();      // acc, newSt, oldSt, q
 75                } else {
 76                    C.dup(1+i+t);    // q, acc, newSt, oldSt, q
 77                    C.push((1+i*t+j)*32);
 78                    C.mload();      // M, q, acc, newSt, oldSt, q
 79                    C.dup(3+i+j);    // oldSt[j], M, q, acc, newSt, oldSt, q
 80                    C.mulmod();      // aux, acc, newSt, oldSt, q
 81                    C.dup(2+i+t);    // q, aux, acc, newSt, oldSt, q
 82                    C.swap(2);       // acc, aux, q, newSt, oldSt, q
 83                    C.addmod();      // acc, newSt, oldSt, q
 84                }
 85            }
 86        }
 87        for (let i=0; i<t; i++) {
 88            C.swap((t -i) + (t -i-1));
 89            C.pop();
 90        }
 91        C.push(0);
 92        C.mload();
 93        C.jmp();
 94    }
 95
 96
 97    // Check selector
 98    C.push("0x0100000000000000000000000000000000000000000000000000000000");
 99    C.push(0);
100    C.calldataload();
101    C.div();
102    C.push("0xc4420fb4"); // poseidon(uint256[])
103    C.eq();
104    C.jmpi("start");
105    C.invalid();
106
107    C.label("start");
108
109    saveM();
110
111    C.push("0x30644e72e131a029b85045b68181585d2833e84879b9709143e1f593f0000001");  // q
112
113    // Load t values from the call data.
114    // The function has a single array param param
115    // [Selector (4)] [Pointer (32)][Length (32)] [data1 (32)] ....
116    // We ignore the pointer and the length and just load t values to the state
117    // (Stack positions 0-{t-1}) If the array is shorter, we just set zeros.
118    for (let i=0; i<t; i++) {
119        C.push(0x44+(0x20*(t-1-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
158module.exports.abi = [
159    {
160        "constant": true,
161        "inputs": [
162            {
163                "name": "input",
164                "type": "uint256[]"
165            }
166        ],
167        "name": "poseidon",
168        "outputs": [
169            {
170                "name": "",
171                "type": "uint256"
172            }
173        ],
174        "payable": false,
175        "stateMutability": "pure",
176        "type": "function"
177    }
178];
179
180module.exports.createCode = createCode;
181
182