1package poseidon
  2
  3import (
  4	"github.com/consensys/gnark/frontend"
  5	"github.com/succinctlabs/gnark-plonky2-verifier/field"
  6)
  7
  8const HALF_N_FULL_ROUNDS = 4
  9const N_PARTIAL_ROUNDS = 22
 10const MAX_WIDTH = 12
 11const SPONGE_WIDTH = 12
 12const SPONGE_RATE = 8
 13
 14type PoseidonState = [SPONGE_WIDTH]frontend.Variable
 15type PoseidonStateExtension = [SPONGE_WIDTH]field.QuadraticExtension
 16type PoseidonHashOut = [4]field.F
 17
 18type PoseidonChip struct {
 19	api      frontend.API                 `gnark:"-"`
 20	fieldAPI field.FieldAPI               `gnark:"-"`
 21	qeAPI    *field.QuadraticExtensionAPI `gnark:"-"`
 22}
 23
 24func NewPoseidonChip(api frontend.API, fieldAPI field.FieldAPI, qeAPI *field.QuadraticExtensionAPI) *PoseidonChip {
 25	return &PoseidonChip{api: api, fieldAPI: fieldAPI, qeAPI: qeAPI}
 26}
 27
 28// The permutation function.
 29// The input state MUST have all it's elements be within Goldilocks field (e.g. this function will not reduce the input elements).
 30// The returned state's elements will all be within Goldilocks field.
 31func (c *PoseidonChip) Poseidon(input PoseidonState) PoseidonState {
 32	state := input
 33	roundCounter := 0
 34	state = c.fullRounds(state, &roundCounter)
 35	state = c.partialRounds(state, &roundCounter)
 36	state = c.fullRounds(state, &roundCounter)
 37	return state
 38}
 39
 40// The input elements MUST have all it's elements be within Goldilocks field.
 41// The returned slice's elements will all be within Goldilocks field.
 42func (c *PoseidonChip) HashNToMNoPad(input []frontend.Variable, nbOutputs int) []frontend.Variable {
 43	var state PoseidonState
 44
 45	for i := 0; i < SPONGE_WIDTH; i++ {
 46		state[i] = frontend.Variable(0)
 47	}
 48
 49	for i := 0; i < len(input); i += SPONGE_RATE {
 50		for j := 0; j < SPONGE_RATE; j++ {
 51			if i+j < len(input) {
 52				state[j] = input[i+j]
 53			}
 54		}
 55		state = c.Poseidon(state)
 56	}
 57
 58	var outputs []frontend.Variable
 59
 60	for {
 61		for i := 0; i < SPONGE_RATE; i++ {
 62			outputs = append(outputs, state[i])
 63			if len(outputs) == nbOutputs {
 64				return outputs
 65			}
 66		}
 67		state = c.Poseidon(state)
 68	}
 69}
 70
 71// The input elements can be outside of the Goldilocks field.
 72// The returned slice's elements will all be within Goldilocks field.
 73func (c *PoseidonChip) HashNoPad(input []field.F) PoseidonHashOut {
 74	var hash PoseidonHashOut
 75	inputVars := []frontend.Variable{}
 76
 77	for i := 0; i < len(input); i++ {
 78		inputVars = append(inputVars, c.fieldAPI.Reduce(input[i]).Limbs[0])
 79	}
 80
 81	outputVars := c.HashNToMNoPad(inputVars, 4)
 82	for i := 0; i < 4; i++ {
 83		hash[i] = c.fieldAPI.NewElement(outputVars[i])
 84	}
 85
 86	return hash
 87}
 88
 89func (c *PoseidonChip) ToVec(hash PoseidonHashOut) []field.F {
 90	return hash[:]
 91}
 92
 93func (c *PoseidonChip) fullRounds(state PoseidonState, roundCounter *int) PoseidonState {
 94	for i := 0; i < HALF_N_FULL_ROUNDS; i++ {
 95		state = c.constantLayer(state, roundCounter)
 96		state = c.sBoxLayer(state)
 97		state = c.mdsLayer(state)
 98		*roundCounter += 1
 99	}
100	return state
101}
102
103func (c *PoseidonChip) partialRounds(state PoseidonState, roundCounter *int) PoseidonState {
104	state = c.partialFirstConstantLayer(state)
105	state = c.mdsPartialLayerInit(state)
106
107	for i := 0; i < N_PARTIAL_ROUNDS; i++ {
108		state[0] = c.sBoxMonomial(state[0])
109		state[0] = field.GoldilocksMulAdd(c.api, frontend.Variable(1), state[0], FAST_PARTIAL_ROUND_CONSTANTS[i])
110		state = c.mdsPartialLayerFast(state, i)
111	}
112
113	*roundCounter += N_PARTIAL_ROUNDS
114
115	return state
116}
117
118func (c *PoseidonChip) constantLayer(state PoseidonState, roundCounter *int) PoseidonState {
119	for i := 0; i < 12; i++ {
120		if i < SPONGE_WIDTH {
121			roundConstant := ALL_ROUND_CONSTANTS[i+SPONGE_WIDTH*(*roundCounter)]
122			state[i] = field.GoldilocksMulAdd(c.api, frontend.Variable(1), state[i], roundConstant)
123		}
124	}
125	return state
126}
127
128func (c *PoseidonChip) ConstantLayerExtension(state PoseidonStateExtension, roundCounter *int) PoseidonStateExtension {
129	for i := 0; i < 12; i++ {
130		if i < SPONGE_WIDTH {
131			roundConstant := c.qeAPI.VarToQE(ALL_ROUND_CONSTANTS[i+SPONGE_WIDTH*(*roundCounter)])
132			state[i] = c.qeAPI.AddExtension(state[i], roundConstant)
133		}
134	}
135	return state
136}
137
138func (c *PoseidonChip) sBoxMonomial(x frontend.Variable) frontend.Variable {
139	x2 := field.GoldilocksMulAdd(c.api, x, x, frontend.Variable(0))
140	x4 := field.GoldilocksMulAdd(c.api, x2, x2, frontend.Variable(0))
141	x6 := field.GoldilocksMulAdd(c.api, x4, x2, frontend.Variable(0))
142	return field.GoldilocksMulAdd(c.api, x6, x, frontend.Variable(0))
143}
144
145func (c *PoseidonChip) SBoxMonomialExtension(x field.QuadraticExtension) field.QuadraticExtension {
146	x2 := c.qeAPI.SquareExtension(x)
147	x4 := c.qeAPI.SquareExtension(x2)
148	x3 := c.qeAPI.MulExtension(x, x2)
149	return c.qeAPI.MulExtension(x3, x4)
150}
151
152func (c *PoseidonChip) sBoxLayer(state PoseidonState) PoseidonState {
153	for i := 0; i < 12; i++ {
154		if i < SPONGE_WIDTH {
155			state[i] = c.sBoxMonomial(state[i])
156		}
157	}
158	return state
159}
160
161func (c *PoseidonChip) SBoxLayerExtension(state PoseidonStateExtension) PoseidonStateExtension {
162	for i := 0; i < 12; i++ {
163		if i < SPONGE_WIDTH {
164			state[i] = c.SBoxMonomialExtension(state[i])
165		}
166	}
167	return state
168}
169
170func (c *PoseidonChip) mdsRowShf(r int, v [SPONGE_WIDTH]frontend.Variable) frontend.Variable {
171	res := ZERO_VAR
172
173	for i := 0; i < 12; i++ {
174		if i < SPONGE_WIDTH {
175			res1 := c.api.Mul(v[(i+r)%SPONGE_WIDTH], MDS_MATRIX_CIRC_VARS[i])
176			res = c.api.Add(res, res1)
177		}
178	}
179
180	res = c.api.Add(res, c.api.Mul(v[r], MDS_MATRIX_DIAG_VARS[r]))
181	return res
182}
183
184func (c *PoseidonChip) MdsRowShfExtension(r int, v [SPONGE_WIDTH]field.QuadraticExtension) field.QuadraticExtension {
185	res := c.qeAPI.FieldToQE(field.ZERO_F)
186
187	for i := 0; i < 12; i++ {
188		if i < SPONGE_WIDTH {
189			matrixVal := c.qeAPI.VarToQE(MDS_MATRIX_CIRC[i])
190			res1 := c.qeAPI.MulExtension(v[(i+r)%SPONGE_WIDTH], matrixVal)
191			res = c.qeAPI.AddExtension(res, res1)
192		}
193	}
194
195	matrixVal := c.qeAPI.VarToQE(MDS_MATRIX_DIAG[r])
196	res = c.qeAPI.AddExtension(res, c.qeAPI.MulExtension(v[r], matrixVal))
197	return res
198}
199
200func (c *PoseidonChip) mdsLayer(state_ PoseidonState) PoseidonState {
201	var result PoseidonState
202	for i := 0; i < SPONGE_WIDTH; i++ {
203		result[i] = frontend.Variable(0)
204	}
205
206	var state [SPONGE_WIDTH]frontend.Variable
207	for i := 0; i < SPONGE_WIDTH; i++ {
208		state[i] = field.GoldilocksReduce(c.api, state_[i])
209	}
210
211	for r := 0; r < 12; r++ {
212		if r < SPONGE_WIDTH {
213			sum := c.mdsRowShf(r, state)
214			result[r] = field.GoldilocksReduce(c.api, sum)
215		}
216	}
217
218	return result
219}
220
221func (c *PoseidonChip) MdsLayerExtension(state_ PoseidonStateExtension) PoseidonStateExtension {
222	var result PoseidonStateExtension
223
224	for r := 0; r < 12; r++ {
225		if r < SPONGE_WIDTH {
226			sum := c.MdsRowShfExtension(r, state_)
227			result[r] = sum
228		}
229	}
230
231	return result
232}
233
234func (c *PoseidonChip) partialFirstConstantLayer(state PoseidonState) PoseidonState {
235	for i := 0; i < 12; i++ {
236		if i < SPONGE_WIDTH {
237			state[i] = field.GoldilocksMulAdd(c.api, frontend.Variable(1), state[i], FAST_PARTIAL_FIRST_ROUND_CONSTANT[i])
238		}
239	}
240	return state
241}
242
243func (c *PoseidonChip) PartialFirstConstantLayerExtension(state PoseidonStateExtension) PoseidonStateExtension {
244	for i := 0; i < 12; i++ {
245		if i < SPONGE_WIDTH {
246			state[i] = c.qeAPI.AddExtension(state[i], c.qeAPI.VarToQE((FAST_PARTIAL_FIRST_ROUND_CONSTANT[i])))
247		}
248	}
249	return state
250}
251
252func (c *PoseidonChip) mdsPartialLayerInit(state PoseidonState) PoseidonState {
253	var result PoseidonState
254	for i := 0; i < 12; i++ {
255		result[i] = frontend.Variable(0)
256	}
257
258	result[0] = state[0]
259
260	for r := 1; r < 12; r++ {
261		if r < SPONGE_WIDTH {
262			for d := 1; d < 12; d++ {
263				if d < SPONGE_WIDTH {
264					t := FAST_PARTIAL_ROUND_INITIAL_MATRIX[r-1][d-1]
265					result[d] = field.GoldilocksMulAdd(c.api, state[r], t, result[d])
266				}
267			}
268		}
269	}
270
271	return result
272}
273
274func (c *PoseidonChip) MdsPartialLayerInitExtension(state PoseidonStateExtension) PoseidonStateExtension {
275	var result PoseidonStateExtension
276	for i := 0; i < 12; i++ {
277		result[i] = c.qeAPI.FieldToQE(field.ZERO_F)
278	}
279
280	result[0] = state[0]
281
282	for r := 1; r < 12; r++ {
283		if r < SPONGE_WIDTH {
284			for d := 1; d < 12; d++ {
285				if d < SPONGE_WIDTH {
286					t := c.qeAPI.VarToQE(FAST_PARTIAL_ROUND_INITIAL_MATRIX[r-1][d-1])
287					result[d] = c.qeAPI.AddExtension(result[d], c.qeAPI.MulExtension(state[r], t))
288				}
289			}
290		}
291	}
292
293	return result
294}
295
296func (c *PoseidonChip) mdsPartialLayerFast(state PoseidonState, r int) PoseidonState {
297	dSum := ZERO_VAR
298	for i := 1; i < 12; i++ {
299		if i < SPONGE_WIDTH {
300			t := FAST_PARTIAL_ROUND_W_HATS_VARS[r][i-1]
301			reducedState := field.GoldilocksReduce(c.api, state[i])
302			dSum = c.api.Add(dSum, c.api.Mul(reducedState, t))
303		}
304	}
305
306	s0 := field.GoldilocksReduce(c.api, state[0])
307	dSum = c.api.Add(dSum, c.api.Mul(s0, MDS0TO0_VAR))
308	d := field.GoldilocksReduce(c.api, dSum)
309
310	var result PoseidonState
311	for i := 0; i < SPONGE_WIDTH; i++ {
312		result[i] = frontend.Variable(0)
313	}
314
315	result[0] = d
316
317	for i := 1; i < 12; i++ {
318		if i < SPONGE_WIDTH {
319			t := FAST_PARTIAL_ROUND_VS[r][i-1]
320			result[i] = field.GoldilocksMulAdd(c.api, state[0], t, state[i])
321		}
322	}
323
324	return result
325}
326
327func (c *PoseidonChip) MdsPartialLayerFastExtension(state PoseidonStateExtension, r int) PoseidonStateExtension {
328	s0 := state[0]
329	mds0to0 := c.qeAPI.VarToQE(MDS0TO0)
330	d := c.qeAPI.MulExtension(s0, mds0to0)
331	for i := 1; i < 12; i++ {
332		if i < SPONGE_WIDTH {
333			t := c.qeAPI.VarToQE(FAST_PARTIAL_ROUND_W_HATS[r][i-1])
334			d = c.qeAPI.AddExtension(d, c.qeAPI.MulExtension(state[i], t))
335		}
336	}
337
338	var result PoseidonStateExtension
339	result[0] = d
340	for i := 1; i < 12; i++ {
341		if i < SPONGE_WIDTH {
342			t := c.qeAPI.VarToQE(FAST_PARTIAL_ROUND_VS[r][i-1])
343			result[i] = c.qeAPI.AddExtension(c.qeAPI.MulExtension(state[0], t), state[i])
344		}
345	}
346
347	return result
348}