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.New(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	two_to_32 := new(big.Int).SetInt64(1 << 32)
 51	two_to_64 := new(big.Int).Mul(two_to_32, two_to_32)
 52
 53	for i := 0; i < len(input); i += BN254_SPONGE_RATE * 3 {
 54		endI := c.min(len(input), i+BN254_SPONGE_RATE*3)
 55		rateChunk := input[i:endI]
 56		for j, stateIdx := 0, 0; j < len(rateChunk); j, stateIdx = j+3, stateIdx+1 {
 57			endJ := c.min(len(rateChunk), j+3)
 58			bn254Chunk := rateChunk[j:endJ]
 59
 60			inter := frontend.Variable(0)
 61			for k := 0; k < len(bn254Chunk); k++ {
 62				inter = c.api.MulAcc(inter, bn254Chunk[k].Limb, new(big.Int).Exp(two_to_64, big.NewInt(int64(k)), nil))
 63			}
 64
 65			state[stateIdx+1] = inter
 66		}
 67
 68		state = c.Poseidon(state)
 69	}
 70
 71	return BN254HashOut(state[0])
 72}
 73
 74func (c *BN254Chip) HashOrNoop(input []gl.Variable) BN254HashOut {
 75	if len(input) <= 3 {
 76		returnVal := frontend.Variable(0)
 77
 78		alpha := new(big.Int).SetInt64(1 << 32)
 79		alpha = new(big.Int).Mul(alpha, alpha)
 80		for i, inputElement := range input {
 81			mulFactor := new(big.Int).Exp(alpha, big.NewInt(int64(i)), nil)
 82			returnVal = c.api.MulAcc(returnVal, inputElement, mulFactor)
 83		}
 84
 85		return BN254HashOut(returnVal)
 86	} else {
 87		return c.HashNoPad(input)
 88	}
 89}
 90
 91func (c *BN254Chip) TwoToOne(left BN254HashOut, right BN254HashOut) BN254HashOut {
 92	var inputs BN254State
 93	inputs[0] = frontend.Variable(0)
 94	inputs[1] = frontend.Variable(0)
 95	inputs[2] = left
 96	inputs[3] = right
 97	state := c.Poseidon(inputs)
 98	return state[0]
 99}
100
101func (c *BN254Chip) ToVec(hash BN254HashOut) []gl.Variable {
102	bits := c.api.ToBinary(hash)
103
104	returnElements := []gl.Variable{}
105
106	// Split into 7 byte chunks, since 8 byte chunks can result in collisions
107	chunkSize := 56
108	for i := 0; i < len(bits); i += chunkSize {
109		maxIdx := c.min(len(bits), i+chunkSize)
110		bitChunk := bits[i:maxIdx]
111		returnElements = append(returnElements, gl.NewVariable(c.api.FromBinary(bitChunk...)))
112	}
113
114	return returnElements
115}
116
117func (c *BN254Chip) min(x, y int) int {
118	if x < y {
119		return x
120	}
121
122	return y
123}
124
125func (c *BN254Chip) fullRounds(state BN254State, isFirst bool) BN254State {
126	for i := 0; i < BN254_FULL_ROUNDS/2-1; i++ {
127		state = c.exp5state(state)
128		if isFirst {
129			state = c.ark(state, (i+1)*BN254_SPONGE_WIDTH)
130		} else {
131			state = c.ark(state, (BN254_FULL_ROUNDS/2+1)*BN254_SPONGE_WIDTH+BN254_PARTIAL_ROUNDS+i*BN254_SPONGE_WIDTH)
132		}
133		state = c.mix(state, mMatrix)
134	}
135
136	state = c.exp5state(state)
137	if isFirst {
138		state = c.ark(state, (BN254_FULL_ROUNDS/2)*BN254_SPONGE_WIDTH)
139		state = c.mix(state, pMatrix)
140	} else {
141		state = c.mix(state, mMatrix)
142	}
143
144	return state
145}
146
147func (c *BN254Chip) partialRounds(state BN254State) BN254State {
148	for i := 0; i < BN254_PARTIAL_ROUNDS; i++ {
149		state[0] = c.exp5(state[0])
150		state[0] = c.api.Add(state[0], cConstants[(BN254_FULL_ROUNDS/2+1)*BN254_SPONGE_WIDTH+i])
151
152		newState0 := frontend.Variable(0)
153		for j := 0; j < BN254_SPONGE_WIDTH; j++ {
154			newState0 = c.api.MulAcc(newState0, sConstants[(BN254_SPONGE_WIDTH*2-1)*i+j], state[j])
155		}
156
157		for k := 1; k < BN254_SPONGE_WIDTH; k++ {
158			state[k] = c.api.MulAcc(state[k], state[0], sConstants[(BN254_SPONGE_WIDTH*2-1)*i+BN254_SPONGE_WIDTH+k-1])
159		}
160		state[0] = newState0
161	}
162
163	return state
164}
165
166func (c *BN254Chip) ark(state BN254State, it int) BN254State {
167	var result BN254State
168
169	for i := 0; i < len(state); i++ {
170		result[i] = c.api.Add(state[i], cConstants[it+i])
171	}
172
173	return result
174}
175
176func (c *BN254Chip) exp5(x frontend.Variable) frontend.Variable {
177	x2 := c.api.Mul(x, x)
178	x4 := c.api.Mul(x2, x2)
179	return c.api.Mul(x4, x)
180}
181
182func (c *BN254Chip) exp5state(state BN254State) BN254State {
183	for i := 0; i < BN254_SPONGE_WIDTH; i++ {
184		state[i] = c.exp5(state[i])
185	}
186	return state
187}
188
189func (c *BN254Chip) mix(state_ BN254State, constantMatrix [][]*big.Int) BN254State {
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			result[i] = c.api.MulAcc(result[i], constantMatrix[j][i], state_[j])
199		}
200	}
201
202	return result
203}