1package plonky2_verifier
  2
  3import (
  4	. "gnark-plonky2-verifier/field"
  5	"gnark-plonky2-verifier/poseidon"
  6)
  7
  8func (g *PoseidonGate) WireInput(i uint64) uint64 {
  9	return i
 10}
 11
 12func (g *PoseidonGate) WireOutput(i uint64) uint64 {
 13	return poseidon.SPONGE_WIDTH + i
 14}
 15
 16func (g *PoseidonGate) WireSwap() uint64 {
 17	return 2 * poseidon.SPONGE_WIDTH
 18}
 19
 20const START_DELTA = 2*poseidon.SPONGE_WIDTH + 1
 21
 22func (g *PoseidonGate) WireDelta(i uint64) uint64 {
 23	if i >= 4 {
 24		panic("Delta index out of range")
 25	}
 26	return START_DELTA + i
 27}
 28
 29const START_FULL_0 = START_DELTA + 4
 30
 31func (g *PoseidonGate) WireFullSBox0(round uint64, i uint64) uint64 {
 32	if round == 0 {
 33		panic("First-round S-box inputs are not stored as wires")
 34	}
 35	if round >= poseidon.HALF_N_FULL_ROUNDS {
 36		panic("S-box input round out of range")
 37	}
 38	if i >= poseidon.SPONGE_WIDTH {
 39		panic("S-box input index out of range")
 40	}
 41
 42	return START_FULL_0 + (round-1)*poseidon.SPONGE_WIDTH + i
 43}
 44
 45const START_PARTIAL = START_FULL_0 + (poseidon.HALF_N_FULL_ROUNDS-1)*poseidon.SPONGE_WIDTH
 46
 47func (g *PoseidonGate) WirePartialSBox(round uint64) uint64 {
 48	if round >= poseidon.N_PARTIAL_ROUNDS {
 49		panic("S-box input round out of range")
 50	}
 51	return START_PARTIAL + round
 52}
 53
 54const START_FULL_1 = START_PARTIAL + poseidon.N_PARTIAL_ROUNDS
 55
 56func (g *PoseidonGate) WireFullSBox1(round uint64, i uint64) uint64 {
 57	if round >= poseidon.HALF_N_FULL_ROUNDS {
 58		panic("S-box input round out of range")
 59	}
 60	if i >= poseidon.SPONGE_WIDTH {
 61		panic("S-box input index out of range")
 62	}
 63
 64	return START_FULL_1 + round*poseidon.SPONGE_WIDTH + i
 65}
 66
 67func (g *PoseidonGate) WiresEnd() uint64 {
 68	return START_FULL_1 + poseidon.HALF_N_FULL_ROUNDS*poseidon.SPONGE_WIDTH
 69}
 70
 71type PoseidonGate struct {
 72}
 73
 74func (p *PoseidonGate) EvalUnfiltered(pc *PlonkChip, vars EvaluationVars) []QuadraticExtension {
 75	constraints := []QuadraticExtension{}
 76
 77	poseidonChip := poseidon.NewPoseidonChip(pc.api, NewFieldAPI(pc.api), pc.qeAPI)
 78
 79	// Assert that `swap` is binary.
 80	swap := vars.localWires[p.WireSwap()]
 81	notSwap := pc.qeAPI.SubExtension(pc.qeAPI.FieldToQE(ONE_F), swap)
 82	constraints = append(constraints, pc.qeAPI.MulExtension(swap, notSwap))
 83
 84	// Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`.
 85	for i := uint64(0); i < 4; i++ {
 86		inputLhs := vars.localWires[p.WireInput(i)]
 87		inputRhs := vars.localWires[p.WireInput(i+4)]
 88		deltaI := vars.localWires[p.WireDelta(i)]
 89		diff := pc.qeAPI.SubExtension(inputRhs, inputLhs)
 90		expectedDeltaI := pc.qeAPI.MulExtension(swap, diff)
 91		constraints = append(constraints, pc.qeAPI.SubExtension(expectedDeltaI, deltaI))
 92	}
 93
 94	// Compute the possibly-swapped input layer.
 95	var state [poseidon.SPONGE_WIDTH]QuadraticExtension
 96	for i := uint64(0); i < 4; i++ {
 97		deltaI := vars.localWires[p.WireDelta(i)]
 98		inputLhs := vars.localWires[p.WireInput(i)]
 99		inputRhs := vars.localWires[p.WireInput(i+4)]
100		state[i] = pc.qeAPI.AddExtension(inputLhs, deltaI)
101		state[i+4] = pc.qeAPI.SubExtension(inputRhs, deltaI)
102	}
103	for i := uint64(8); i < poseidon.SPONGE_WIDTH; i++ {
104		state[i] = vars.localWires[p.WireInput(i)]
105	}
106
107	round_ctr := 0
108
109	// First set of full rounds.
110	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
111		state = poseidonChip.ConstantLayerExtension(state, &round_ctr)
112		if r != 0 {
113			for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
114				sBoxIn := vars.localWires[p.WireFullSBox0(r, i)]
115				constraints = append(constraints, pc.qeAPI.SubExtension(state[i], sBoxIn))
116				state[i] = sBoxIn
117			}
118		}
119		state = poseidonChip.SBoxLayerExtension(state)
120		state = poseidonChip.MdsLayerExtension(state)
121		round_ctr++
122	}
123
124	// Partial rounds.
125	state = poseidonChip.PartialFirstConstantLayerExtension(state)
126	state = poseidonChip.MdsPartialLayerInitExtension(state)
127	for r := uint64(0); r < poseidon.N_PARTIAL_ROUNDS-1; r++ {
128		sBoxIn := vars.localWires[p.WirePartialSBox(r)]
129		constraints = append(constraints, pc.qeAPI.SubExtension(state[0], sBoxIn))
130		state[0] = poseidonChip.SBoxMonomialExtension(sBoxIn)
131		state[0] = pc.qeAPI.AddExtension(state[0], pc.qeAPI.FieldToQE(NewFieldElement(poseidon.FAST_PARTIAL_ROUND_CONSTANTS[r])))
132		state = poseidonChip.MdsPartialLayerFastExtension(state, int(r))
133	}
134	sBoxIn := vars.localWires[p.WirePartialSBox(poseidon.N_PARTIAL_ROUNDS-1)]
135	constraints = append(constraints, pc.qeAPI.SubExtension(state[0], sBoxIn))
136	state[0] = poseidonChip.SBoxMonomialExtension(sBoxIn)
137	state = poseidonChip.MdsPartialLayerFastExtension(state, poseidon.N_PARTIAL_ROUNDS-1)
138	round_ctr += poseidon.N_PARTIAL_ROUNDS
139
140	// Second set of full rounds.
141	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
142		poseidonChip.ConstantLayerExtension(state, &round_ctr)
143		for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
144			sBoxIn := vars.localWires[p.WireFullSBox1(r, i)]
145			constraints = append(constraints, pc.qeAPI.SubExtension(state[i], sBoxIn))
146			state[i] = sBoxIn
147		}
148		state = poseidonChip.MdsLayerExtension(state)
149		state = poseidonChip.SBoxLayerExtension(state)
150		round_ctr++
151	}
152
153	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
154		constraints = append(constraints, pc.qeAPI.SubExtension(state[i], vars.localWires[p.WireOutput(i)]))
155	}
156
157	return constraints
158}