1package gates
  2
  3import (
  4	"regexp"
  5
  6	"github.com/consensys/gnark/frontend"
  7	"github.com/succinctlabs/gnark-plonky2-verifier/field"
  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(api frontend.API, qeAPI *field.QuadraticExtensionAPI, vars EvaluationVars) []field.QuadraticExtension {
 93	constraints := []field.QuadraticExtension{}
 94
 95	poseidonChip := poseidon.NewPoseidonChip(api, field.NewFieldAPI(api), qeAPI)
 96
 97	// Assert that `swap` is binary.
 98	swap := vars.localWires[g.WireSwap()]
 99	swapMinusOne := qeAPI.SubExtension(swap, qeAPI.FieldToQE(field.ONE_F))
100	constraints = append(constraints, qeAPI.MulExtension(swap, swapMinusOne))
101
102	// Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`.
103	for i := uint64(0); i < 4; i++ {
104		inputLhs := vars.localWires[g.WireInput(i)]
105		inputRhs := vars.localWires[g.WireInput(i+4)]
106		deltaI := vars.localWires[g.WireDelta(i)]
107		diff := qeAPI.SubExtension(inputRhs, inputLhs)
108		expectedDeltaI := qeAPI.MulExtension(swap, diff)
109		constraints = append(constraints, qeAPI.SubExtension(expectedDeltaI, deltaI))
110	}
111
112	// Compute the possibly-swapped input layer.
113	var state [poseidon.SPONGE_WIDTH]field.QuadraticExtension
114	for i := uint64(0); i < 4; i++ {
115		deltaI := vars.localWires[g.WireDelta(i)]
116		inputLhs := vars.localWires[g.WireInput(i)]
117		inputRhs := vars.localWires[g.WireInput(i+4)]
118		state[i] = qeAPI.AddExtension(inputLhs, deltaI)
119		state[i+4] = qeAPI.SubExtension(inputRhs, deltaI)
120	}
121	for i := uint64(8); i < poseidon.SPONGE_WIDTH; i++ {
122		state[i] = vars.localWires[g.WireInput(i)]
123	}
124
125	roundCounter := 0
126
127	// First set of full rounds.
128	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
129		state = poseidonChip.ConstantLayerExtension(state, &roundCounter)
130		if r != 0 {
131			for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
132				sBoxIn := vars.localWires[g.WireFullSBox0(r, i)]
133				constraints = append(constraints, qeAPI.SubExtension(state[i], sBoxIn))
134				state[i] = sBoxIn
135			}
136		}
137		state = poseidonChip.SBoxLayerExtension(state)
138		state = poseidonChip.MdsLayerExtension(state)
139		roundCounter++
140	}
141
142	// Partial rounds.
143	state = poseidonChip.PartialFirstConstantLayerExtension(state)
144	state = poseidonChip.MdsPartialLayerInitExtension(state)
145
146	for r := uint64(0); r < poseidon.N_PARTIAL_ROUNDS-1; r++ {
147		sBoxIn := vars.localWires[g.WirePartialSBox(r)]
148		constraints = append(constraints, qeAPI.SubExtension(state[0], sBoxIn))
149		state[0] = poseidonChip.SBoxMonomialExtension(sBoxIn)
150		state[0] = qeAPI.AddExtension(state[0], qeAPI.FieldToQE(poseidon.FAST_PARTIAL_ROUND_CONSTANTS[r]))
151		state = poseidonChip.MdsPartialLayerFastExtension(state, int(r))
152	}
153	sBoxIn := vars.localWires[g.WirePartialSBox(poseidon.N_PARTIAL_ROUNDS-1)]
154	constraints = append(constraints, qeAPI.SubExtension(state[0], sBoxIn))
155	state[0] = poseidonChip.SBoxMonomialExtension(sBoxIn)
156	state = poseidonChip.MdsPartialLayerFastExtension(state, poseidon.N_PARTIAL_ROUNDS-1)
157	roundCounter += poseidon.N_PARTIAL_ROUNDS
158
159	// Second set of full rounds.
160	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
161		state = poseidonChip.ConstantLayerExtension(state, &roundCounter)
162		for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
163			sBoxIn := vars.localWires[g.WireFullSBox1(r, i)]
164			constraints = append(constraints, qeAPI.SubExtension(state[i], sBoxIn))
165			state[i] = sBoxIn
166		}
167		state = poseidonChip.SBoxLayerExtension(state)
168		state = poseidonChip.MdsLayerExtension(state)
169		roundCounter++
170	}
171
172	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
173		constraints = append(constraints, qeAPI.SubExtension(state[i], vars.localWires[g.WireOutput(i)]))
174	}
175
176	return constraints
177}