1pragma circom 2.0.0;
  2
  3include "./poseidon_constants.circom";
  4
  5template Sigma() {
  6    signal input in;
  7    signal output out;
  8
  9    signal in2;
 10    signal in4;
 11
 12    in2 <== in*in;
 13    in4 <== in2*in2;
 14
 15    out <== in4*in;
 16}
 17
 18template Ark(t, C, r) {
 19    signal input in[t];
 20    signal output out[t];
 21
 22    for (var i=0; i<t; i++) {
 23        out[i] <== in[i] + C[i + r];
 24    }
 25}
 26
 27template Mix(t, M) {
 28    signal input in[t];
 29    signal output out[t];
 30
 31    var lc;
 32    for (var i=0; i<t; i++) {
 33        lc = 0;
 34        for (var j=0; j<t; j++) {
 35            lc += M[j][i]*in[j];
 36        }
 37        out[i] <== lc;
 38    }
 39}
 40
 41template MixLast(t, M, s) {
 42    signal input in[t];
 43    signal output out;
 44
 45    var lc = 0;
 46    for (var j=0; j<t; j++) {
 47        lc += M[j][s]*in[j];
 48    }
 49    out <== lc;
 50}
 51
 52template MixS(t, S, r) {
 53    signal input in[t];
 54    signal output out[t];
 55
 56
 57    var lc = 0;
 58    for (var i=0; i<t; i++) {
 59        lc += S[(t*2-1)*r+i]*in[i];
 60    }
 61    out[0] <== lc;
 62    for (var i=1; i<t; i++) {
 63        out[i] <== in[i] +  in[0] * S[(t*2-1)*r + t + i -1];
 64    }
 65}
 66
 67template PoseidonEx(nInputs, nOuts) {
 68    signal input inputs[nInputs];
 69    signal input initialState;
 70    signal output out[nOuts];
 71
 72    // Using recommended parameters from whitepaper https://eprint.iacr.org/2019/458.pdf (table 2, table 8)
 73    // Generated by https://extgit.iaik.tugraz.at/krypto/hadeshash/-/blob/master/code/calc_round_numbers.py
 74    // And rounded up to nearest integer that divides by t
 75    var N_ROUNDS_P[16] = [56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68];
 76    var t = nInputs + 1;
 77    var nRoundsF = 8;
 78    var nRoundsP = N_ROUNDS_P[t - 2];
 79    var C[t*nRoundsF + nRoundsP] = POSEIDON_C(t);
 80    var S[  N_ROUNDS_P[t-2]  *  (t*2-1)  ]  = POSEIDON_S(t);
 81    var M[t][t] = POSEIDON_M(t);
 82    var P[t][t] = POSEIDON_P(t);
 83
 84    component ark[nRoundsF];
 85    component sigmaF[nRoundsF][t];
 86    component sigmaP[nRoundsP];
 87    component mix[nRoundsF-1];
 88    component mixS[nRoundsP];
 89    component mixLast[nOuts];
 90
 91
 92    ark[0] = Ark(t, C, 0);
 93    for (var j=0; j<t; j++) {
 94        if (j>0) {
 95            ark[0].in[j] <== inputs[j-1];
 96        } else {
 97            ark[0].in[j] <== initialState;
 98        }
 99    }
100
101    for (var r = 0; r < nRoundsF\2-1; r++) {
102        for (var j=0; j<t; j++) {
103            sigmaF[r][j] = Sigma();
104            if(r==0) {
105                sigmaF[r][j].in <== ark[0].out[j];
106            } else {
107                sigmaF[r][j].in <== mix[r-1].out[j];
108            }
109        }
110
111        ark[r+1] = Ark(t, C, (r+1)*t);
112        for (var j=0; j<t; j++) {
113            ark[r+1].in[j] <== sigmaF[r][j].out;
114        }
115
116        mix[r] = Mix(t,M);
117        for (var j=0; j<t; j++) {
118            mix[r].in[j] <== ark[r+1].out[j];
119        }
120
121    }
122
123    for (var j=0; j<t; j++) {
124        sigmaF[nRoundsF\2-1][j] = Sigma();
125        sigmaF[nRoundsF\2-1][j].in <== mix[nRoundsF\2-2].out[j];
126    }
127
128    ark[nRoundsF\2] = Ark(t, C, (nRoundsF\2)*t );
129    for (var j=0; j<t; j++) {
130        ark[nRoundsF\2].in[j] <== sigmaF[nRoundsF\2-1][j].out;
131    }
132
133    mix[nRoundsF\2-1] = Mix(t,P);
134    for (var j=0; j<t; j++) {
135        mix[nRoundsF\2-1].in[j] <== ark[nRoundsF\2].out[j];
136    }
137
138
139    for (var r = 0; r < nRoundsP; r++) {
140        sigmaP[r] = Sigma();
141        if (r==0) {
142            sigmaP[r].in <== mix[nRoundsF\2-1].out[0];
143        } else {
144            sigmaP[r].in <== mixS[r-1].out[0];
145        }
146
147        mixS[r] = MixS(t, S, r);
148        for (var j=0; j<t; j++) {
149            if (j==0) {
150                mixS[r].in[j] <== sigmaP[r].out + C[(nRoundsF\2+1)*t + r];
151            } else {
152                if (r==0) {
153                    mixS[r].in[j] <== mix[nRoundsF\2-1].out[j];
154                } else {
155                    mixS[r].in[j] <== mixS[r-1].out[j];
156                }
157            }
158        }
159    }
160
161    for (var r = 0; r < nRoundsF\2-1; r++) {
162        for (var j=0; j<t; j++) {
163            sigmaF[nRoundsF\2 + r][j] = Sigma();
164            if (r==0) {
165                sigmaF[nRoundsF\2 + r][j].in <== mixS[nRoundsP-1].out[j];
166            } else {
167                sigmaF[nRoundsF\2 + r][j].in <== mix[nRoundsF\2+r-1].out[j];
168            }
169        }
170
171        ark[ nRoundsF\2 + r + 1] = Ark(t, C,  (nRoundsF\2+1)*t + nRoundsP + r*t );
172        for (var j=0; j<t; j++) {
173            ark[nRoundsF\2 + r + 1].in[j] <== sigmaF[nRoundsF\2 + r][j].out;
174        }
175
176        mix[nRoundsF\2 + r] = Mix(t,M);
177        for (var j=0; j<t; j++) {
178            mix[nRoundsF\2 + r].in[j] <== ark[nRoundsF\2 + r + 1].out[j];
179        }
180
181    }
182
183    for (var j=0; j<t; j++) {
184        sigmaF[nRoundsF-1][j] = Sigma();
185        sigmaF[nRoundsF-1][j].in <== mix[nRoundsF-2].out[j];
186    }
187
188    for (var i=0; i<nOuts; i++) {
189        mixLast[i] = MixLast(t,M,i);
190        for (var j=0; j<t; j++) {
191            mixLast[i].in[j] <== sigmaF[nRoundsF-1][j].out;
192        }
193        out[i] <== mixLast[i].out;
194    }
195
196}
197
198template Poseidon(nInputs) {
199    signal input inputs[nInputs];
200    signal output out;
201
202    component pEx = PoseidonEx(nInputs, 1);
203    pEx.initialState <== 0;
204    for (var i=0; i<nInputs; i++) {
205        pEx.inputs[i] <== inputs[i];
206    }
207    out <== pEx.out[0];
208}