1package poseidon
  2
  3import (
  4	. "gnark-ed25519/goldilocks"
  5	"math/big"
  6	"testing"
  7
  8	"github.com/consensys/gnark-crypto/ecc"
  9	"github.com/consensys/gnark/backend/groth16"
 10	"github.com/consensys/gnark/frontend"
 11	"github.com/consensys/gnark/frontend/cs/r1cs"
 12	"github.com/consensys/gnark/test"
 13)
 14
 15var testCurve = ecc.BN254
 16
 17type TestPoseidonCircuit struct {
 18	In  [12]frontend.Variable
 19	Out [12]frontend.Variable
 20}
 21
 22func (circuit *TestPoseidonCircuit) Define(api frontend.API) error {
 23	goldilocksApi := NewGoldilocksAPI(api)
 24
 25	// BN254 -> Binary(64) -> GoldilocksElement
 26	var input PoseidonState
 27	for i := 0; i < 12; i++ {
 28		input[i] = goldilocksApi.FromBinary(api.ToBinary(circuit.In[i], 64)).(GoldilocksElement)
 29	}
 30
 31	output := Poseidon(api, goldilocksApi, input)
 32
 33	// Check that output is correct
 34	for i := 0; i < 12; i++ {
 35		goldilocksApi.AssertIsEqual(
 36			output[i],
 37			goldilocksApi.FromBinary(api.ToBinary(circuit.Out[i])).(GoldilocksElement),
 38		)
 39	}
 40
 41	return nil
 42}
 43
 44func TestPoseidonWitness(t *testing.T) {
 45	assert := test.NewAssert(t)
 46
 47	testCase := func(inBigInt [12]big.Int, outBigInt [12]big.Int) {
 48		var in [12]frontend.Variable
 49		var out [12]frontend.Variable
 50
 51		for i := 0; i < 12; i++ {
 52			in[i] = inBigInt[i]
 53			out[i] = outBigInt[i]
 54		}
 55
 56		circuit := TestPoseidonCircuit{In: in, Out: out}
 57		witness := TestPoseidonCircuit{In: in, Out: out}
 58		err := test.IsSolved(&circuit, &witness, testCurve.ScalarField())
 59		assert.NoError(err)
 60	}
 61
 62	inStr := [12]string{"0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0"}
 63	outStr := [12]string{
 64		"4330397376401421145", "14124799381142128323", "8742572140681234676",
 65		"14345658006221440202", "15524073338516903644", "5091405722150716653",
 66		"15002163819607624508", "2047012902665707362", "16106391063450633726",
 67		"4680844749859802542", "15019775476387350140", "1698615465718385111",
 68	}
 69
 70	var inBigInt [12]big.Int
 71	var outBigInt [12]big.Int
 72
 73	for i := 0; i < 12; i++ {
 74		inTmp := new(big.Int)
 75		inTmp, _ = inTmp.SetString(inStr[i], 10)
 76		inBigInt[i] = *inTmp
 77
 78		outTmp := new(big.Int)
 79		outTmp, _ = outTmp.SetString(outStr[i], 10)
 80		outBigInt[i] = *outTmp
 81	}
 82
 83	testCase(inBigInt, outBigInt)
 84}
 85
 86func TestPoseidonProof(t *testing.T) {
 87	inStr := [12]string{"0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0"}
 88	outStr := [12]string{
 89		"4330397376401421145", "14124799381142128323", "8742572140681234676",
 90		"14345658006221440202", "15524073338516903644", "5091405722150716653",
 91		"15002163819607624508", "2047012902665707362", "16106391063450633726",
 92		"4680844749859802542", "15019775476387350140", "1698615465718385111",
 93	}
 94
 95	var in [12]frontend.Variable
 96	var out [12]frontend.Variable
 97
 98	for i := 0; i < 12; i++ {
 99		inTmp := new(big.Int)
100		inTmp, _ = inTmp.SetString(inStr[i], 10)
101		in[i] = *inTmp
102
103		outTmp := new(big.Int)
104		outTmp, _ = outTmp.SetString(outStr[i], 10)
105		out[i] = *outTmp
106	}
107
108	circuit := TestPoseidonCircuit{In: in, Out: out}
109	assignment := TestPoseidonCircuit{In: in, Out: out}
110
111	r1cs, err := frontend.Compile(testCurve.ScalarField(), r1cs.NewBuilder, &circuit)
112	if err != nil {
113		panic(err)
114	}
115
116	witness, err := frontend.NewWitness(&assignment, testCurve.ScalarField())
117	if err != nil {
118		panic(err)
119	}
120
121	pk, vk, err := groth16.Setup(r1cs)
122	if err != nil {
123		panic(err)
124	}
125
126	err = test.IsSolved(&circuit, &assignment, testCurve.ScalarField())
127	if err != nil {
128		panic(err)
129	}
130
131	proof, err := groth16.Prove(r1cs, pk, witness)
132	if err != nil {
133		panic(err)
134	}
135
136	publicWitness, err := witness.Public()
137	if err != nil {
138		panic(err)
139	}
140
141	err = groth16.Verify(proof, vk, publicWitness)
142	if err != nil {
143		panic(err)
144	}
145}