fix poseidon tests (#16)

Author: zhenfeizhang

poseidon/param.go

@@ -23,6 +23,8 @@ type PoseidonParams struct {
 
 // TODOs: the parameters are not secure. use a better way to generate the constants
 func NewPoseidonParams() *PoseidonParams {
+	r := rand.New(rand.NewSource(42))
+
 	num_full_rounds := 8
 	num_part_rounds := 14
 	num_states := 16
@@ -31,13 +33,13 @@ func NewPoseidonParams() *PoseidonParams {
 	for i := 0; i < num_states; i++ {
 		external_round_constant[i] = make([]uint32, num_full_rounds)
 		for j := 0; j < num_full_rounds; j++ {
-			external_round_constant[i][j] = randomM31()
+			external_round_constant[i][j] = randomM31(r)
 		}
 	}
 
 	internal_round_constant := make([]uint32, num_part_rounds)
 	for i := 0; i < num_part_rounds; i++ {
-		internal_round_constant[i] = randomM31()
+		internal_round_constant[i] = randomM31(r)
 	}
 
 	// mds parameters adopted from Plonky3
@@ -66,8 +68,8 @@ func NewPoseidonParams() *PoseidonParams {
 	}
 }
 
-func randomM31() uint32 {
-	t := rand.Uint32() & 0x7FFFFFFF
+func randomM31(r *rand.Rand) uint32 {
+	t := r.Uint32() & 0x7FFFFFFF
 
 	for t == 0x7fffffff {
 		t = rand.Uint32() & 0x7FFFFFFF

poseidon/poseidon_circuit_test.go

@@ -1,71 +1,71 @@
 package poseidon
 
 import (
-	"fmt"
 	"testing"
 
+	"github.com/PolyhedraZK/ExpanderCompilerCollection"
 	"github.com/PolyhedraZK/ExpanderCompilerCollection/field/m31"
-	"github.com/consensys/gnark-crypto/ecc"
+	"github.com/PolyhedraZK/ExpanderCompilerCollection/test"
 	"github.com/consensys/gnark/constraint"
 	"github.com/consensys/gnark/frontend"
-	"github.com/consensys/gnark/frontend/cs/r1cs"
-	"github.com/consensys/gnark/test"
 )
 
 type MockPoseidonCircuit struct {
-	State  []frontend.Variable `gnark:",public"`
-	Output frontend.Variable   `gnark:",public"`
+	State  [16]frontend.Variable `gnark:",public"`
+	Output frontend.Variable     `gnark:",public"`
 }
 
 func (c *MockPoseidonCircuit) Define(api frontend.API) (err error) {
-	// Define the circuit
 	param := NewPoseidonParams()
 	engine := m31.Field{}
-	t := PoseidonCircuit(api, engine, param, c.State, false)
+	t := PoseidonCircuit(api, engine, param, c.State[:], false)
 	api.AssertIsEqual(t, c.Output)
 
 	return
 }
 
 func TestPoseidonCircuit(t *testing.T) {
-	assert := test.NewAssert(t)
-
 	param := NewPoseidonParams()
 
-	state := make([]constraint.Element, 16)
-	stateVar := make([]frontend.Variable, 16)
-	var internalStateVars PoseidonInternalStateVar
+	var states [16]constraint.Element
+	var stateVars [16]frontend.Variable
+	var outputVar frontend.Variable
 
-	for i := 0; i < 16; i++ {
-		state[i] = constraint.Element{uint64(i)}
-		stateVar[i] = frontend.Variable(uint64(i))
+	for j := 0; j < 16; j++ {
+		states[j] = constraint.Element{uint64(j)}
+		stateVars[j] = frontend.Variable(uint64(j))
 	}
-	internalState, output := PoseidonM31WithInternalStates(param, state, true)
-	outputVar := frontend.Variable(output[0])
+	output := PoseidonM31(param, states[:])
+	outputVar = frontend.Variable(output[0])
 
-	fmt.Println("internal state", internalState)
-
-	for j := 0; j < 16; j++ {
-		internalStateVars.AfterHalfFullRound[j] = frontend.Variable(internalState.AfterHalfFullRound[j][0])
-		internalStateVars.AfterHalfPartialRound[j] = frontend.Variable(internalState.AfterHalfPartialRound[j][0])
-		internalStateVars.AfterPartialRound[j] = frontend.Variable(internalState.AfterPartialRound[j][0])
+	assignment := &MockPoseidonCircuit{
+		State:  stateVars,
+		Output: outputVar,
 	}
 
-	c := MockPoseidonCircuit{
-		stateVar,
-		outputVar,
+	// Gnark test disabled as it does not support randomness and custom gates
+	// err := test.IsSolved(&MockPoseidonCircuit{}, assignment, m31.ScalarField)
+	// if err != nil {
+	// 	panic(err)
+	// }
+	// fmt.Println("Gnark test passed")
+
+	// Ecc test
+	circuit, err := ExpanderCompilerCollection.Compile(m31.ScalarField, &MockPoseidonCircuit{}, frontend.WithCompressThreshold(32))
+	if err != nil {
+		panic(err)
 	}
 
-	w, _ := frontend.NewWitness(&c, m31.ScalarField)
-	fmt.Println("witness", w)
+	layered_circuit := circuit.GetLayeredCircuit()
+	// circuit.GetCircuitIr().Print()
 
-	err := test.IsSolved(&c, &c, m31.ScalarField)
-	assert.NoError(err)
+	inputSolver := circuit.GetInputSolver()
+	witness, err := inputSolver.SolveInputAuto(assignment)
+	if err != nil {
+		panic(err)
+	}
 
-	r1cs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &c)
-	assert.NoError(err)
-	fmt.Println("num constraints:", r1cs.GetNbConstraints())
-	fmt.Println("num coefficients:", r1cs.GetNbCoefficients())
-	i, p, s := r1cs.GetNbVariables()
-	fmt.Println("num variables:", i, p, s)
+	if !test.CheckCircuit(layered_circuit, witness) {
+		panic("verification failed")
+	}
 }