1package poseidon
  2
  3import (
  4	. "gnark-ed25519/field"
  5
  6	"github.com/consensys/gnark/frontend"
  7)
  8
  9const HALF_N_FULL_ROUNDS = 4
 10const N_FULL_ROUNDS_TOTAL = 2 * HALF_N_FULL_ROUNDS
 11const N_PARTIAL_ROUNDS = 22
 12const N_ROUNDS = N_FULL_ROUNDS_TOTAL + N_PARTIAL_ROUNDS
 13const MAX_WIDTH = 12
 14const WIDTH = 12
 15const SPONGE_WIDTH = 12
 16const SPONGE_RATE = 8
 17
 18type PoseidonState = [WIDTH]F
 19type PoseidonChip struct {
 20	api   frontend.API
 21	field frontend.API
 22}
 23
 24func NewPoseidonChip(api frontend.API, field frontend.API) *PoseidonChip {
 25	return &PoseidonChip{api: api, field: field}
 26}
 27
 28func (c *PoseidonChip) Poseidon(input PoseidonState) PoseidonState {
 29	state := input
 30	roundCounter := 0
 31	state = c.fullRounds(state, &roundCounter)
 32	state = c.partialRounds(state, &roundCounter)
 33	state = c.fullRounds(state, &roundCounter)
 34	return state
 35}
 36
 37func (c *PoseidonChip) HashNToMNoPad(input []F, nbOutputs int) []F {
 38	var state PoseidonState
 39
 40	for i := 0; i < len(input); i += SPONGE_RATE {
 41		for j := 0; j < SPONGE_RATE; j++ {
 42			if i+j < len(input) {
 43				state[j] = input[i+j]
 44			}
 45		}
 46		state = c.Poseidon(state)
 47	}
 48
 49	var outputs []F
 50
 51	for {
 52		for i := 0; i < SPONGE_RATE; i++ {
 53			outputs = append(outputs, state[i])
 54			if len(outputs) == nbOutputs {
 55				return outputs
 56			}
 57		}
 58		state = c.Poseidon(state)
 59	}
 60}
 61
 62func (c *PoseidonChip) HashNoPad(input []F) Hash {
 63	var hash Hash
 64	copy(hash[:], c.HashNToMNoPad(input, 4))
 65	return hash
 66}
 67
 68func (c *PoseidonChip) fullRounds(state PoseidonState, roundCounter *int) PoseidonState {
 69	for i := 0; i < HALF_N_FULL_ROUNDS; i++ {
 70		state = c.constantLayer(state, roundCounter)
 71		state = c.sBoxLayer(state)
 72		state = c.mdsLayer(state)
 73		if *roundCounter >= 26 && i == 3 {
 74			break
 75		}
 76		*roundCounter += 1
 77	}
 78	return state
 79}
 80
 81func (c *PoseidonChip) partialRounds(state PoseidonState, roundCounter *int) PoseidonState {
 82	state = c.partialFirstConstantLayer(state)
 83	state = c.mdsPartialLayerInit(state)
 84
 85	for i := 0; i < N_PARTIAL_ROUNDS; i++ {
 86		state[0] = c.sBoxMonomial(state[0])
 87		state[0] = c.field.Add(state[0], FAST_PARTIAL_ROUND_CONSTANTS[i]).(F)
 88		state = c.mdsPartialLayerFast(state, i)
 89	}
 90
 91	*roundCounter += N_PARTIAL_ROUNDS
 92
 93	return state
 94}
 95
 96func (c *PoseidonChip) constantLayer(state PoseidonState, roundCounter *int) PoseidonState {
 97	for i := 0; i < 12; i++ {
 98		if i < WIDTH {
 99			roundConstant := NewFieldElement(ALL_ROUND_CONSTANTS[i+WIDTH*(*roundCounter)])
100			state[i] = c.field.Add(state[i], roundConstant).(F)
101		}
102	}
103	return state
104}
105
106func (c *PoseidonChip) sBoxLayer(state PoseidonState) PoseidonState {
107	for i := 0; i < 12; i++ {
108		if i < WIDTH {
109			state[i] = c.sBoxMonomial(state[i])
110		}
111	}
112	return state
113}
114
115func (c *PoseidonChip) sBoxMonomial(x F) F {
116	x2 := c.field.Mul(x, x)
117	x4 := c.field.Mul(x2, x2)
118	x3 := c.field.Mul(x2, x)
119	return c.field.Mul(x3, x4).(F)
120}
121
122func (c *PoseidonChip) mdsRowShf(r int, v [WIDTH]frontend.Variable) frontend.Variable {
123	res := frontend.Variable(0)
124
125	for i := 0; i < 12; i++ {
126		if i < WIDTH {
127			res1 := c.api.Mul(v[(i+r)%WIDTH], frontend.Variable(MDS_MATRIX_CIRC[i]))
128			res = c.api.Add(res, res1)
129		}
130	}
131
132	res = c.api.Add(res, c.api.Mul(v[r], MDS_MATRIX_DIAG[r]))
133	return res
134}
135
136func (c *PoseidonChip) mdsLayer(state_ PoseidonState) PoseidonState {
137	var result PoseidonState
138	for i := 0; i < WIDTH; i++ {
139		result[i] = NewFieldElement(0)
140	}
141
142	var state [WIDTH]frontend.Variable
143	for i := 0; i < WIDTH; i++ {
144		state[i] = c.api.FromBinary(c.field.ToBinary(state_[i])...)
145	}
146
147	for r := 0; r < 12; r++ {
148		if r < WIDTH {
149			sum := c.mdsRowShf(r, state)
150			bits := c.api.ToBinary(sum)
151			result[r] = c.field.FromBinary(bits).(F)
152		}
153	}
154
155	return result
156}
157
158func (c *PoseidonChip) partialFirstConstantLayer(state PoseidonState) PoseidonState {
159	for i := 0; i < 12; i++ {
160		if i < WIDTH {
161			state[i] = c.field.Add(state[i], NewFieldElement(FAST_PARTIAL_FIRST_ROUND_CONSTANT[i])).(F)
162		}
163	}
164	return state
165}
166
167func (c *PoseidonChip) mdsPartialLayerInit(state PoseidonState) PoseidonState {
168	var result PoseidonState
169	for i := 0; i < 12; i++ {
170		result[i] = NewFieldElement(0)
171	}
172
173	result[0] = state[0]
174
175	for r := 1; r < 12; r++ {
176		if r < WIDTH {
177			for d := 1; d < 12; d++ {
178				if d < WIDTH {
179					t := NewFieldElement(FAST_PARTIAL_ROUND_INITIAL_MATRIX[r-1][d-1])
180					result[d] = c.field.Add(result[d], c.field.Mul(state[r], t)).(F)
181				}
182			}
183		}
184	}
185
186	return result
187}
188
189func (c *PoseidonChip) mdsPartialLayerFast(state PoseidonState, r int) PoseidonState {
190	dSum := frontend.Variable(0)
191	for i := 1; i < 12; i++ {
192		if i < WIDTH {
193			t := frontend.Variable(FAST_PARTIAL_ROUND_W_HATS[r][i-1])
194			si := c.api.FromBinary(c.field.ToBinary(state[i])...)
195			dSum = c.api.Add(dSum, c.api.Mul(si, t))
196		}
197	}
198
199	s0 := c.api.FromBinary(c.field.ToBinary(state[0])...)
200	mds0to0 := frontend.Variable(MDS_MATRIX_CIRC[0] + MDS_MATRIX_DIAG[0])
201	dSum = c.api.Add(dSum, c.api.Mul(s0, mds0to0))
202	d := c.field.FromBinary(c.api.ToBinary(dSum))
203
204	var result PoseidonState
205	for i := 0; i < WIDTH; i++ {
206		result[i] = NewFieldElement(0)
207	}
208
209	result[0] = d.(F)
210
211	for i := 1; i < 12; i++ {
212		if i < WIDTH {
213			t := NewFieldElement(FAST_PARTIAL_ROUND_VS[r][i-1])
214			result[i] = c.field.Add(state[i], c.field.Mul(state[0], t)).(F)
215		}
216	}
217
218	return result
219}