1package gates
  2
  3import (
  4	"regexp"
  5
  6	"github.com/consensys/gnark/frontend"
  7	gl "github.com/succinctlabs/gnark-plonky2-verifier/goldilocks"
  8	"github.com/succinctlabs/gnark-plonky2-verifier/poseidon"
  9)
 10
 11var poseidonGateRegex = regexp.MustCompile("PoseidonGate.*")
 12
 13func deserializePoseidonGate(parameters map[string]string) Gate {
 14	// Has the format "PoseidonGate(PhantomData<plonky2_field::goldilocks_field::GoldilocksField>)<WIDTH=12>"
 15	return NewPoseidonGate()
 16}
 17
 18type PoseidonGate struct {
 19}
 20
 21func NewPoseidonGate() *PoseidonGate {
 22	return &PoseidonGate{}
 23}
 24
 25func (g *PoseidonGate) Id() string {
 26	return "PoseidonGate"
 27}
 28
 29func (g *PoseidonGate) WireInput(i uint64) uint64 {
 30	return i
 31}
 32
 33func (g *PoseidonGate) WireOutput(i uint64) uint64 {
 34	return poseidon.SPONGE_WIDTH + i
 35}
 36
 37func (g *PoseidonGate) WireSwap() uint64 {
 38	return 2 * poseidon.SPONGE_WIDTH
 39}
 40
 41const START_DELTA = 2*poseidon.SPONGE_WIDTH + 1
 42
 43func (g *PoseidonGate) WireDelta(i uint64) uint64 {
 44	if i >= 4 {
 45		panic("Delta index out of range")
 46	}
 47	return START_DELTA + i
 48}
 49
 50const START_FULL_0 = START_DELTA + 4
 51
 52func (g *PoseidonGate) WireFullSBox0(round uint64, i uint64) uint64 {
 53	if round == 0 {
 54		panic("First-round S-box inputs are not stored as wires")
 55	}
 56	if round >= poseidon.HALF_N_FULL_ROUNDS {
 57		panic("S-box input round out of range")
 58	}
 59	if i >= poseidon.SPONGE_WIDTH {
 60		panic("S-box input index out of range")
 61	}
 62
 63	return START_FULL_0 + (round-1)*poseidon.SPONGE_WIDTH + i
 64}
 65
 66const START_PARTIAL = START_FULL_0 + (poseidon.HALF_N_FULL_ROUNDS-1)*poseidon.SPONGE_WIDTH
 67
 68func (g *PoseidonGate) WirePartialSBox(round uint64) uint64 {
 69	if round >= poseidon.N_PARTIAL_ROUNDS {
 70		panic("S-box input round out of range")
 71	}
 72	return START_PARTIAL + round
 73}
 74
 75const START_FULL_1 = START_PARTIAL + poseidon.N_PARTIAL_ROUNDS
 76
 77func (g *PoseidonGate) WireFullSBox1(round uint64, i uint64) uint64 {
 78	if round >= poseidon.HALF_N_FULL_ROUNDS {
 79		panic("S-box input round out of range")
 80	}
 81	if i >= poseidon.SPONGE_WIDTH {
 82		panic("S-box input index out of range")
 83	}
 84
 85	return START_FULL_1 + round*poseidon.SPONGE_WIDTH + i
 86}
 87
 88func (g *PoseidonGate) WiresEnd() uint64 {
 89	return START_FULL_1 + poseidon.HALF_N_FULL_ROUNDS*poseidon.SPONGE_WIDTH
 90}
 91
 92func (g *PoseidonGate) EvalUnfiltered(
 93	api frontend.API,
 94	glApi *gl.Chip,
 95	vars EvaluationVars,
 96) []gl.QuadraticExtensionVariable {
 97	constraints := []gl.QuadraticExtensionVariable{}
 98
 99	poseidonChip := poseidon.NewGoldilocksChip(api)
100
101	// Assert that `swap` is binary.
102	swap := vars.localWires[g.WireSwap()]
103	swapMinusOne := glApi.SubExtension(swap, gl.OneExtension())
104	constraints = append(constraints, glApi.MulExtension(swap, swapMinusOne))
105
106	// Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`.
107	for i := uint64(0); i < 4; i++ {
108		inputLhs := vars.localWires[g.WireInput(i)]
109		inputRhs := vars.localWires[g.WireInput(i+4)]
110		deltaI := vars.localWires[g.WireDelta(i)]
111		diff := glApi.SubExtension(inputRhs, inputLhs)
112		expectedDeltaI := glApi.MulExtension(swap, diff)
113		constraints = append(constraints, glApi.SubExtension(expectedDeltaI, deltaI))
114	}
115
116	// Compute the possibly-swapped input layer.
117	var state [poseidon.SPONGE_WIDTH]gl.QuadraticExtensionVariable
118	for i := uint64(0); i < 4; i++ {
119		deltaI := vars.localWires[g.WireDelta(i)]
120		inputLhs := vars.localWires[g.WireInput(i)]
121		inputRhs := vars.localWires[g.WireInput(i+4)]
122		state[i] = glApi.AddExtension(inputLhs, deltaI)
123		state[i+4] = glApi.SubExtension(inputRhs, deltaI)
124	}
125	for i := uint64(8); i < poseidon.SPONGE_WIDTH; i++ {
126		state[i] = vars.localWires[g.WireInput(i)]
127	}
128
129	roundCounter := 0
130
131	// First set of full rounds.
132	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
133		state = poseidonChip.ConstantLayerExtension(state, &roundCounter)
134		if r != 0 {
135			for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
136				sBoxIn := vars.localWires[g.WireFullSBox0(r, i)]
137				constraints = append(constraints, glApi.SubExtension(state[i], sBoxIn))
138				state[i] = sBoxIn
139			}
140		}
141		state = poseidonChip.SBoxLayerExtension(state)
142		state = poseidonChip.MdsLayerExtension(state)
143		roundCounter++
144	}
145
146	// Partial rounds.
147	state = poseidonChip.PartialFirstConstantLayerExtension(state)
148	state = poseidonChip.MdsPartialLayerInitExtension(state)
149
150	for r := uint64(0); r < poseidon.N_PARTIAL_ROUNDS-1; r++ {
151		sBoxIn := vars.localWires[g.WirePartialSBox(r)]
152		constraints = append(constraints, glApi.SubExtension(state[0], sBoxIn))
153		state[0] = poseidonChip.SBoxMonomialExtension(sBoxIn)
154		state[0] = glApi.AddExtension(state[0], gl.NewQuadraticExtensionVariable(gl.NewVariable(poseidon.FAST_PARTIAL_ROUND_CONSTANTS[r]), gl.Zero()))
155		state = poseidonChip.MdsPartialLayerFastExtension(state, int(r))
156	}
157	sBoxIn := vars.localWires[g.WirePartialSBox(poseidon.N_PARTIAL_ROUNDS-1)]
158	constraints = append(constraints, glApi.SubExtension(state[0], sBoxIn))
159	state[0] = poseidonChip.SBoxMonomialExtension(sBoxIn)
160	state = poseidonChip.MdsPartialLayerFastExtension(state, poseidon.N_PARTIAL_ROUNDS-1)
161	roundCounter += poseidon.N_PARTIAL_ROUNDS
162
163	// Second set of full rounds.
164	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
165		state = poseidonChip.ConstantLayerExtension(state, &roundCounter)
166		for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
167			sBoxIn := vars.localWires[g.WireFullSBox1(r, i)]
168			constraints = append(constraints, glApi.SubExtension(state[i], sBoxIn))
169			state[i] = sBoxIn
170		}
171		state = poseidonChip.SBoxLayerExtension(state)
172		state = poseidonChip.MdsLayerExtension(state)
173		roundCounter++
174	}
175
176	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
177		constraints = append(constraints, glApi.SubExtension(state[i], vars.localWires[g.WireOutput(i)]))
178	}
179
180	return constraints
181}