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