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