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