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 poseidonMdsGateRegex = regexp.MustCompile("PoseidonMdsGate.*")
12
13func deserializePoseidonMdsGate(parameters map[string]string) Gate {
14	// Has the format "PoseidonMdsGate(PhantomData<plonky2_field::goldilocks_field::GoldilocksField>)<WIDTH=12>"
15	return NewPoseidonMdsGate()
16}
17
18type PoseidonMdsGate struct {
19}
20
21func NewPoseidonMdsGate() *PoseidonMdsGate {
22	return &PoseidonMdsGate{}
23}
24
25func (g *PoseidonMdsGate) Id() string {
26	return "PoseidonMdsGate"
27}
28
29func (g *PoseidonMdsGate) WireInput(i uint64) Range {
30	if i >= poseidon.SPONGE_WIDTH {
31		panic("Input less than sponge width")
32	}
33	return Range{i * field.D, (i + 1) * field.D}
34}
35
36func (g *PoseidonMdsGate) WireOutput(i uint64) Range {
37	if i >= poseidon.SPONGE_WIDTH {
38		panic("Input less than sponge width")
39	}
40	return Range{(poseidon.SPONGE_WIDTH + i) * field.D, (poseidon.SPONGE_WIDTH + i + 1) * field.D}
41}
42
43func (g *PoseidonMdsGate) mdsRowShfAlgebra(r uint64, v [poseidon.SPONGE_WIDTH]field.QEAlgebra, qeAPI *field.QuadraticExtensionAPI) field.QEAlgebra {
44	if r >= poseidon.SPONGE_WIDTH {
45		panic("MDS row index out of range")
46	}
47
48	res := qeAPI.ZERO_QE_ALGEBRA
49	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
50		coeff := qeAPI.FieldToQE(poseidon.MDS_MATRIX_CIRC[i])
51		res = qeAPI.AddExtensionAlgebra(res, qeAPI.ScalarMulExtensionAlgebra(coeff, v[(i+r)%poseidon.SPONGE_WIDTH]))
52	}
53
54	coeff := qeAPI.FieldToQE(poseidon.MDS_MATRIX_DIAG[r])
55	res = qeAPI.AddExtensionAlgebra(res, qeAPI.ScalarMulExtensionAlgebra(coeff, v[r]))
56
57	return res
58}
59
60func (g *PoseidonMdsGate) mdsLayerAlgebra(state [poseidon.SPONGE_WIDTH]field.QEAlgebra, qeAPI *field.QuadraticExtensionAPI) [poseidon.SPONGE_WIDTH]field.QEAlgebra {
61	var result [poseidon.SPONGE_WIDTH]field.QEAlgebra
62	for r := uint64(0); r < poseidon.SPONGE_WIDTH; r++ {
63		result[r] = g.mdsRowShfAlgebra(r, state, qeAPI)
64	}
65
66	return result
67}
68
69func (g *PoseidonMdsGate) EvalUnfiltered(api frontend.API, qeAPI *field.QuadraticExtensionAPI, vars EvaluationVars) []field.QuadraticExtension {
70	constraints := []field.QuadraticExtension{}
71
72	var inputs [poseidon.SPONGE_WIDTH]field.QEAlgebra
73	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
74		inputs[i] = vars.GetLocalExtAlgebra(g.WireInput(i))
75	}
76
77	computed_outputs := g.mdsLayerAlgebra(inputs, qeAPI)
78
79	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
80		output := vars.GetLocalExtAlgebra(g.WireOutput(i))
81		diff := qeAPI.SubExtensionAlgebra(output, computed_outputs[i])
82		constraints = append(constraints, diff[0], diff[1])
83	}
84
85	return constraints
86}