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