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