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