1package poseidon
  2
  3import (
  4	"math/big"
  5
  6	"github.com/consensys/gnark/frontend"
  7	"github.com/succinctlabs/gnark-plonky2-verifier/field"
  8)
  9
 10const fullRounds = 8
 11const partialRounds = 56
 12const spongeWidth = 4
 13const spongeRate = 3
 14
 15type PoseidonBN128Chip struct {
 16	api      frontend.API   `gnark:"-"`
 17	fieldAPI field.FieldAPI `gnark:"-"`
 18}
 19
 20type PoseidonBN128State = [spongeWidth]frontend.Variable
 21type PoseidonBN128HashOut = frontend.Variable
 22
 23// This implementation is based on the following implementation:
 24// https://github.com/iden3/go-iden3-crypto/blob/e5cf066b8be3da9a3df9544c65818df189fdbebe/poseidon/poseidon.go
 25func NewPoseidonBN128Chip(api frontend.API, fieldAPI field.FieldAPI) *PoseidonBN128Chip {
 26	return &PoseidonBN128Chip{api: api, fieldAPI: fieldAPI}
 27}
 28
 29func (c *PoseidonBN128Chip) Poseidon(state PoseidonBN128State) PoseidonBN128State {
 30	state = c.ark(state, 0)
 31	state = c.fullRounds(state, true)
 32	state = c.partialRounds(state)
 33	state = c.fullRounds(state, false)
 34	return state
 35}
 36
 37func (c *PoseidonBN128Chip) HashNoPad(input []field.F) PoseidonBN128HashOut {
 38	state := PoseidonBN128State{
 39		frontend.Variable(0),
 40		frontend.Variable(0),
 41		frontend.Variable(0),
 42		frontend.Variable(0),
 43	}
 44
 45	for i := 0; i < len(input); i += spongeRate * 3 {
 46		endI := c.min(len(input), i+spongeRate*3)
 47		rateChunk := input[i:endI]
 48		for j, stateIdx := 0, 0; j < len(rateChunk); j, stateIdx = j+3, stateIdx+1 {
 49			endJ := c.min(len(rateChunk), j+3)
 50			bn128Chunk := rateChunk[j:endJ]
 51
 52			bits := []frontend.Variable{}
 53			for k := 0; k < len(bn128Chunk); k++ {
 54				bn128Chunk[k] = c.fieldAPI.Reduce(bn128Chunk[k])
 55				bits = append(bits, c.fieldAPI.ToBits(bn128Chunk[k])...)
 56			}
 57
 58			state[stateIdx+1] = c.api.FromBinary(bits...)
 59		}
 60
 61		state = c.Poseidon(state)
 62	}
 63
 64	return PoseidonBN128HashOut(state[0])
 65}
 66
 67func (c *PoseidonBN128Chip) HashOrNoop(input []field.F) PoseidonBN128HashOut {
 68	if len(input) <= 3 {
 69		returnVal := frontend.Variable(0)
 70
 71		alpha := new(big.Int).SetInt64(1 << 32)
 72		for i, inputElement := range input {
 73			returnVal = c.api.Add(returnVal, c.api.Mul(inputElement, alpha.Exp(alpha, big.NewInt(int64(i)), nil)))
 74		}
 75
 76		return PoseidonBN128HashOut(returnVal)
 77	} else {
 78		return c.HashNoPad(input)
 79	}
 80}
 81
 82func (c *PoseidonBN128Chip) TwoToOne(left PoseidonBN128HashOut, right PoseidonBN128HashOut) PoseidonBN128HashOut {
 83	var inputs PoseidonBN128State
 84	inputs[0] = frontend.Variable(0)
 85	inputs[1] = frontend.Variable(0)
 86	inputs[2] = left
 87	inputs[3] = right
 88	state := c.Poseidon(inputs)
 89	return state[0]
 90}
 91
 92func (c *PoseidonBN128Chip) ToVec(hash PoseidonBN128HashOut) []field.F {
 93	bits := c.api.ToBinary(hash)
 94
 95	returnElements := []field.F{}
 96
 97	// Split into 7 byte chunks, since 8 byte chunks can result in collisions
 98	chunkSize := 56
 99	for i := 0; i < len(bits); i += chunkSize {
100		maxIdx := c.min(len(bits), i+chunkSize)
101		bitChunk := bits[i:maxIdx]
102		returnElements = append(returnElements, c.fieldAPI.FromBits(bitChunk...))
103	}
104
105	return returnElements
106}
107
108func (c *PoseidonBN128Chip) min(x, y int) int {
109	if x < y {
110		return x
111	}
112
113	return y
114}
115
116func (c *PoseidonBN128Chip) fullRounds(state PoseidonBN128State, isFirst bool) PoseidonBN128State {
117	for i := 0; i < fullRounds/2-1; i++ {
118		state = c.exp5state(state)
119		if isFirst {
120			state = c.ark(state, (i+1)*spongeWidth)
121		} else {
122			state = c.ark(state, (fullRounds/2+1)*spongeWidth+partialRounds+i*spongeWidth)
123		}
124		state = c.mix(state, mMatrix)
125	}
126
127	state = c.exp5state(state)
128	if isFirst {
129		state = c.ark(state, (fullRounds/2)*spongeWidth)
130		state = c.mix(state, pMatrix)
131	} else {
132		state = c.mix(state, mMatrix)
133	}
134
135	return state
136}
137
138func (c *PoseidonBN128Chip) partialRounds(state PoseidonBN128State) PoseidonBN128State {
139	for i := 0; i < partialRounds; i++ {
140		state[0] = c.exp5(state[0])
141		state[0] = c.api.Add(state[0], cConstants[(fullRounds/2+1)*spongeWidth+i])
142
143		var mul frontend.Variable
144		newState0 := frontend.Variable(0)
145		for j := 0; j < spongeWidth; j++ {
146			mul = c.api.Mul(sConstants[(spongeWidth*2-1)*i+j], state[j])
147			newState0 = c.api.Add(newState0, mul)
148		}
149
150		for k := 1; k < spongeWidth; k++ {
151			mul = c.api.Mul(state[0], sConstants[(spongeWidth*2-1)*i+spongeWidth+k-1])
152			state[k] = c.api.Add(state[k], mul)
153		}
154		state[0] = newState0
155	}
156
157	return state
158}
159
160func (c *PoseidonBN128Chip) ark(state PoseidonBN128State, it int) PoseidonBN128State {
161	var result PoseidonBN128State
162
163	for i := 0; i < len(state); i++ {
164		result[i] = c.api.Add(state[i], cConstants[it+i])
165	}
166
167	return result
168}
169
170func (c *PoseidonBN128Chip) exp5(x frontend.Variable) frontend.Variable {
171	x2 := c.api.Mul(x, x)
172	x4 := c.api.Mul(x2, x2)
173	return c.api.Mul(x4, x)
174}
175
176func (c *PoseidonBN128Chip) exp5state(state PoseidonBN128State) PoseidonBN128State {
177	for i := 0; i < spongeWidth; i++ {
178		state[i] = c.exp5(state[i])
179	}
180	return state
181}
182
183func (c *PoseidonBN128Chip) mix(state_ PoseidonBN128State, constantMatrix [][]*big.Int) PoseidonBN128State {
184	var mul frontend.Variable
185	var result PoseidonBN128State
186
187	for i := 0; i < spongeWidth; i++ {
188		result[i] = frontend.Variable(0)
189	}
190
191	for i := 0; i < spongeWidth; i++ {
192		for j := 0; j < spongeWidth; j++ {
193			mul = c.api.Mul(constantMatrix[j][i], state_[j])
194			result[i] = c.api.Add(result[i], mul)
195		}
196	}
197
198	return result
199}