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