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	// Assert that `swap` is binary.
 78	swap := vars.localWires[p.WireSwap()]
 79	notSwap := pc.qeAPI.SubExtension(pc.qeAPI.FieldToQE(ONE_F), swap)
 80	constraints = append(constraints, pc.qeAPI.MulExtension(swap, notSwap))
 81
 82	// Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`.
 83	for i := uint64(0); i < 4; i++ {
 84		inputLhs := vars.localWires[p.WireInput(i)]
 85		inputRhs := vars.localWires[p.WireInput(i+4)]
 86		deltaI := vars.localWires[p.WireDelta(i)]
 87		diff := pc.qeAPI.SubExtension(inputRhs, inputLhs)
 88		expectedDeltaI := pc.qeAPI.MulExtension(swap, diff)
 89		constraints = append(constraints, pc.qeAPI.SubExtension(expectedDeltaI, deltaI))
 90	}
 91
 92	// Compute the possibly-swapped input layer.
 93	state := make([]QuadraticExtension, poseidon.SPONGE_WIDTH)
 94	for i := uint64(0); i < 4; i++ {
 95		deltaI := vars.localWires[p.WireDelta(i)]
 96		inputLhs := vars.localWires[p.WireInput(i)]
 97		inputRhs := vars.localWires[p.WireInput(i+4)]
 98		state[i] = pc.qeAPI.AddExtension(inputLhs, deltaI)
 99		state[i+4] = pc.qeAPI.SubExtension(inputRhs, deltaI)
100	}
101	for i := uint64(8); i < poseidon.SPONGE_WIDTH; i++ {
102		state[i] = vars.localWires[p.WireInput(i)]
103	}
104
105	round_ctr := 0
106
107	// First set of full rounds.
108	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
109		// TODO: constantLayerField(state, round_ctr)
110		if r != 0 {
111			for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
112				sBoxIn := vars.localWires[p.WireFullSBox0(r, i)]
113				constraints = append(constraints, pc.qeAPI.SubExtension(state[i], sBoxIn))
114				state[i] = sBoxIn
115			}
116		}
117		// TODO: sboxLayerField(state)
118		// TODO: state = mdsLayerField(state)
119		round_ctr++
120	}
121
122	// Partial rounds.
123	// TODO: partialFirstConstantLayer(state)
124	// TODO: state = mdsParitalLayerInit(state)
125	for r := uint64(0); r < poseidon.N_PARTIAL_ROUNDS-1; r++ {
126		sBoxIn := vars.localWires[p.WirePartialSBox(r)]
127		constraints = append(constraints, pc.qeAPI.SubExtension(state[0], sBoxIn))
128		// TODO: state[0] = sBoxMonomial(sBoxIn)
129		// TODO: state[0] += NewFieldElement(FAST_PARTIAL_ROUND_CONSTANTS[r])
130		// TODO: state = mdsParitalLayerFastField(state, r)
131	}
132	sBoxIn := vars.localWires[p.WirePartialSBox(poseidon.N_PARTIAL_ROUNDS-1)]
133	constraints = append(constraints, pc.qeAPI.SubExtension(state[0], sBoxIn))
134	// TODO: state[0] = sBoxMonomial(sBoxIn)
135	// TODO: state = mdsPartialLayerLastField(state, poseidon.N_PARTIAL_ROUNDS - 1)
136	round_ctr += poseidon.N_PARTIAL_ROUNDS
137
138	// Second set of full rounds.
139	for r := uint64(0); r < poseidon.HALF_N_FULL_ROUNDS; r++ {
140		// TODO: constantLayerField(state, round_ctr)
141		for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
142			sBoxIn := vars.localWires[p.WireFullSBox1(r, i)]
143			constraints = append(constraints, pc.qeAPI.SubExtension(state[i], sBoxIn))
144			state[i] = sBoxIn
145		}
146		// TODO: sboxLayerField(state)
147		// TODO: state = mdsLayerField(state)
148		round_ctr++
149	}
150
151	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
152		constraints = append(constraints, pc.qeAPI.SubExtension(state[i], vars.localWires[p.WireOutput(i)]))
153	}
154
155	return constraints
156}