1package poseidon
  2
  3import (
  4	"testing"
  5
  6	"github.com/consensys/gnark/backend/groth16"
  7	"github.com/consensys/gnark/frontend"
  8	"github.com/consensys/gnark/frontend/cs/r1cs"
  9	"github.com/consensys/gnark/test"
 10	"github.com/succinctlabs/gnark-plonky2-verifier/field"
 11	"github.com/succinctlabs/gnark-plonky2-verifier/utils"
 12)
 13
 14type TestPoseidonCircuit struct {
 15	In  [12]frontend.Variable
 16	Out [12]frontend.Variable
 17}
 18
 19func (circuit *TestPoseidonCircuit) Define(api frontend.API) error {
 20	goldilocksApi := field.NewFieldAPI(api)
 21	qeAPI := field.NewQuadraticExtensionAPI(api, goldilocksApi)
 22
 23	var input PoseidonState
 24	for i := 0; i < 12; i++ {
 25		input[i] = circuit.In[i]
 26	}
 27
 28	poseidonChip := NewPoseidonChip(api, goldilocksApi, qeAPI)
 29	output := poseidonChip.Poseidon(input)
 30
 31	for i := 0; i < 12; i++ {
 32		api.AssertIsEqual(output[i], circuit.Out[i])
 33	}
 34
 35	return nil
 36}
 37
 38func TestPoseidonWitness(t *testing.T) {
 39	assert := test.NewAssert(t)
 40
 41	testCase := func(in [12]frontend.Variable, out [12]frontend.Variable) {
 42		circuit := TestPoseidonCircuit{In: in, Out: out}
 43		witness := TestPoseidonCircuit{In: in, Out: out}
 44		err := test.IsSolved(&circuit, &witness, field.TEST_CURVE.ScalarField())
 45		assert.NoError(err)
 46	}
 47
 48	inStr := []string{"0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0"}
 49	outStr := []string{
 50		"4330397376401421145", "14124799381142128323", "8742572140681234676",
 51		"14345658006221440202", "15524073338516903644", "5091405722150716653",
 52		"15002163819607624508", "2047012902665707362", "16106391063450633726",
 53		"4680844749859802542", "15019775476387350140", "1698615465718385111",
 54	}
 55	var in [12]frontend.Variable
 56	var out [12]frontend.Variable
 57	copy(in[:], utils.StrArrayToFrontendVariableArray(inStr))
 58	copy(out[:], utils.StrArrayToFrontendVariableArray(outStr))
 59	testCase(in, out)
 60}
 61
 62func TestPoseidonProof(t *testing.T) {
 63	inStr := []string{"0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0"}
 64	outStr := []string{
 65		"4330397376401421145", "14124799381142128323", "8742572140681234676",
 66		"14345658006221440202", "15524073338516903644", "5091405722150716653",
 67		"15002163819607624508", "2047012902665707362", "16106391063450633726",
 68		"4680844749859802542", "15019775476387350140", "1698615465718385111",
 69	}
 70	var in [12]frontend.Variable
 71	var out [12]frontend.Variable
 72	copy(in[:], utils.StrArrayToFrontendVariableArray(inStr))
 73	copy(out[:], utils.StrArrayToFrontendVariableArray(outStr))
 74
 75	circuit := TestPoseidonCircuit{In: in, Out: out}
 76	assignment := TestPoseidonCircuit{In: in, Out: out}
 77
 78	r1cs, err := frontend.Compile(field.TEST_CURVE.ScalarField(), r1cs.NewBuilder, &circuit)
 79	if err != nil {
 80		panic(err)
 81	}
 82
 83	witness, err := frontend.NewWitness(&assignment, field.TEST_CURVE.ScalarField())
 84	if err != nil {
 85		panic(err)
 86	}
 87
 88	pk, vk, err := groth16.Setup(r1cs)
 89	if err != nil {
 90		panic(err)
 91	}
 92
 93	err = test.IsSolved(&circuit, &assignment, field.TEST_CURVE.ScalarField())
 94	if err != nil {
 95		panic(err)
 96	}
 97
 98	proof, err := groth16.Prove(r1cs, pk, witness)
 99	if err != nil {
100		panic(err)
101	}
102
103	publicWitness, err := witness.Public()
104	if err != nil {
105		panic(err)
106	}
107
108	err = groth16.Verify(proof, vk, publicWitness)
109	if err != nil {
110		panic(err)
111	}
112}