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