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