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 PoseidonChip struct {
 22	api   frontend.API
 23	field frontend.API
 24}
 25
 26func Poseidon(api frontend.API, field frontend.API, input PoseidonState) PoseidonState {
 27	chip := &PoseidonChip{api: api, field: field}
 28	return chip.Poseidon(input)
 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) fullRounds(state PoseidonState, roundCounter *int) PoseidonState {
 41	for i := 0; i < HALF_N_FULL_ROUNDS; i++ {
 42		state = c.constantLayer(state, roundCounter)
 43		state = c.sBoxLayer(state)
 44		state = c.mdsLayer(state)
 45		if *roundCounter >= 26 && i == 3 {
 46			break
 47		}
 48		*roundCounter += 1
 49	}
 50	return state
 51}
 52
 53func (c *PoseidonChip) partialRounds(state PoseidonState, roundCounter *int) PoseidonState {
 54	state = c.partialFirstConstantLayer(state)
 55	state = c.mdsPartialLayerInit(state)
 56
 57	for i := 0; i < N_PARTIAL_ROUNDS; i++ {
 58		state[0] = c.sBoxMonomial(state[0])
 59		state[0] = c.field.Add(state[0], FAST_PARTIAL_ROUND_CONSTANTS[i]).(GoldilocksElement)
 60		state = c.mdsPartialLayerFast(state, i)
 61	}
 62
 63	*roundCounter += N_PARTIAL_ROUNDS
 64
 65	return state
 66}
 67
 68func (c *PoseidonChip) constantLayer(state PoseidonState, roundCounter *int) PoseidonState {
 69	for i := 0; i < 12; i++ {
 70		if i < WIDTH {
 71			roundConstant := NewGoldilocksElement(ALL_ROUND_CONSTANTS[i+WIDTH*(*roundCounter)])
 72			state[i] = c.field.Add(state[i], roundConstant).(GoldilocksElement)
 73		}
 74	}
 75	return state
 76}
 77
 78func (c *PoseidonChip) sBoxLayer(state PoseidonState) PoseidonState {
 79	for i := 0; i < 12; i++ {
 80		if i < WIDTH {
 81			state[i] = c.sBoxMonomial(state[i])
 82		}
 83	}
 84	return state
 85}
 86
 87func (c *PoseidonChip) sBoxMonomial(x GoldilocksElement) GoldilocksElement {
 88	x2 := c.field.Mul(x, x)
 89	x4 := c.field.Mul(x2, x2)
 90	x3 := c.field.Mul(x2, x)
 91	return c.field.Mul(x3, x4).(GoldilocksElement)
 92}
 93
 94func (c *PoseidonChip) mdsRowShf(r int, v [WIDTH]frontend.Variable) frontend.Variable {
 95	res := frontend.Variable(0)
 96
 97	for i := 0; i < 12; i++ {
 98		if i < WIDTH {
 99			res1 := c.api.Mul(v[(i+r)%WIDTH], frontend.Variable(MDS_MATRIX_CIRC[i]))
100			res = c.api.Add(res, res1)
101		}
102	}
103
104	res = c.api.Add(res, c.api.Mul(v[r], MDS_MATRIX_DIAG[r]))
105	return res
106}
107
108func (c *PoseidonChip) mdsLayer(state_ PoseidonState) PoseidonState {
109	var result PoseidonState
110	for i := 0; i < WIDTH; i++ {
111		result[i] = NewGoldilocksElement(0)
112	}
113
114	var state [WIDTH]frontend.Variable
115	for i := 0; i < WIDTH; i++ {
116		state[i] = c.api.FromBinary(c.field.ToBinary(state_[i])...)
117	}
118
119	for r := 0; r < 12; r++ {
120		if r < WIDTH {
121			sum := c.mdsRowShf(r, state)
122			bits := c.api.ToBinary(sum)
123			result[r] = c.field.FromBinary(bits).(GoldilocksElement)
124		}
125	}
126
127	return result
128}
129
130func (c *PoseidonChip) partialFirstConstantLayer(state PoseidonState) PoseidonState {
131	for i := 0; i < 12; i++ {
132		if i < WIDTH {
133			state[i] = c.field.Add(state[i], NewGoldilocksElement(FAST_PARTIAL_FIRST_ROUND_CONSTANT[i])).(GoldilocksElement)
134		}
135	}
136	return state
137}
138
139func (c *PoseidonChip) mdsPartialLayerInit(state PoseidonState) PoseidonState {
140	var result PoseidonState
141	for i := 0; i < 12; i++ {
142		result[i] = NewGoldilocksElement(0)
143	}
144
145	result[0] = state[0]
146
147	for r := 1; r < 12; r++ {
148		if r < WIDTH {
149			for d := 1; d < 12; d++ {
150				if d < WIDTH {
151					t := NewGoldilocksElement(FAST_PARTIAL_ROUND_INITIAL_MATRIX[r-1][d-1])
152					result[d] = c.field.Add(result[d], c.field.Mul(state[r], t)).(GoldilocksElement)
153				}
154			}
155		}
156	}
157
158	return result
159}
160
161func (c *PoseidonChip) mdsPartialLayerFast(state PoseidonState, r int) PoseidonState {
162	dSum := frontend.Variable(0)
163	for i := 1; i < 12; i++ {
164		if i < WIDTH {
165			t := frontend.Variable(FAST_PARTIAL_ROUND_W_HATS[r][i-1])
166			si := c.api.FromBinary(c.field.ToBinary(state[i])...)
167			dSum = c.api.Add(dSum, c.api.Mul(si, t))
168		}
169	}
170
171	s0 := c.api.FromBinary(c.field.ToBinary(state[0])...)
172	mds0to0 := frontend.Variable(MDS_MATRIX_CIRC[0] + MDS_MATRIX_DIAG[0])
173	dSum = c.api.Add(dSum, c.api.Mul(s0, mds0to0))
174	d := c.field.FromBinary(c.api.ToBinary(dSum))
175
176	var result PoseidonState
177	for i := 0; i < WIDTH; i++ {
178		result[i] = NewGoldilocksElement(0)
179	}
180
181	result[0] = d.(GoldilocksElement)
182
183	for i := 1; i < 12; i++ {
184		if i < WIDTH {
185			t := NewGoldilocksElement(FAST_PARTIAL_ROUND_VS[r][i-1])
186			result[i] = c.field.Add(state[i], c.field.Mul(state[0], t)).(GoldilocksElement)
187		}
188	}
189
190	return result
191}