1package gates
 2
 3import (
 4	"regexp"
 5
 6	"github.com/consensys/gnark/frontend"
 7	gl "github.com/succinctlabs/gnark-plonky2-verifier/goldilocks"
 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 * gl.D, (i + 1) * gl.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) * gl.D, (poseidon.SPONGE_WIDTH + i + 1) * gl.D}
41}
42
43func (g *PoseidonMdsGate) mdsRowShfAlgebra(
44	r uint64,
45	v [poseidon.SPONGE_WIDTH]gl.QuadraticExtensionAlgebraVariable,
46	api frontend.API,
47) gl.QuadraticExtensionAlgebraVariable {
48	glApi := gl.New(api)
49	if r >= poseidon.SPONGE_WIDTH {
50		panic("MDS row index out of range")
51	}
52
53	res := gl.ZeroExtensionAlgebra()
54	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
55		coeff := gl.NewQuadraticExtensionVariable(gl.NewVariable(poseidon.MDS_MATRIX_CIRC[i]), gl.Zero())
56		res = glApi.AddExtensionAlgebra(res, glApi.ScalarMulExtensionAlgebra(coeff, v[(i+r)%poseidon.SPONGE_WIDTH]))
57	}
58
59	coeff := gl.NewQuadraticExtensionVariable(gl.NewVariable(poseidon.MDS_MATRIX_DIAG[r]), gl.Zero())
60	res = glApi.AddExtensionAlgebra(res, glApi.ScalarMulExtensionAlgebra(coeff, v[r]))
61
62	return res
63}
64
65func (g *PoseidonMdsGate) mdsLayerAlgebra(
66	state [poseidon.SPONGE_WIDTH]gl.QuadraticExtensionAlgebraVariable,
67	api frontend.API,
68) [poseidon.SPONGE_WIDTH]gl.QuadraticExtensionAlgebraVariable {
69	var result [poseidon.SPONGE_WIDTH]gl.QuadraticExtensionAlgebraVariable
70	for r := uint64(0); r < poseidon.SPONGE_WIDTH; r++ {
71		result[r] = g.mdsRowShfAlgebra(r, state, api)
72	}
73	return result
74}
75
76func (g *PoseidonMdsGate) EvalUnfiltered(
77	api frontend.API,
78	glApi gl.Chip,
79	vars EvaluationVars,
80) []gl.QuadraticExtensionVariable {
81	constraints := []gl.QuadraticExtensionVariable{}
82
83	var inputs [poseidon.SPONGE_WIDTH]gl.QuadraticExtensionAlgebraVariable
84	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
85		inputs[i] = vars.GetLocalExtAlgebra(g.WireInput(i))
86	}
87
88	computed_outputs := g.mdsLayerAlgebra(inputs, api)
89
90	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
91		output := vars.GetLocalExtAlgebra(g.WireOutput(i))
92		diff := glApi.SubExtensionAlgebra(output, computed_outputs[i])
93		for i := 0; i < gl.D; i++ {
94			constraints = append(constraints, diff[i])
95		}
96	}
97
98	return constraints
99}