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