1package circuitcompiler
  2
  3import (
  4	"math/big"
  5	"strings"
  6	"testing"
  7
  8	"github.com/stretchr/testify/assert"
  9)
 10
 11func TestCircuitParser(t *testing.T) {
 12	// y = x^3 + x + 5
 13	flat := `
 14	func main(private s0, public s1):
 15		s2 = s0 * s0
 16		s3 = s2 * s0
 17		s4 = s3 + s0
 18		s5 = s4 + 5
 19		equals(s1, s5)
 20		out = 1 * 1
 21	`
 22	parser := NewParser(strings.NewReader(flat))
 23	circuit, err := parser.Parse()
 24	assert.Nil(t, err)
 25
 26	// flat code to R1CS
 27	a, b, c := circuit.GenerateR1CS()
 28	assert.Equal(t, "s0", circuit.PrivateInputs[0])
 29	assert.Equal(t, "s1", circuit.PublicInputs[0])
 30
 31	assert.Equal(t, []string{"one", "s1", "s0", "s2", "s3", "s4", "s5", "out"}, circuit.Signals)
 32
 33	// expected result
 34	b0 := big.NewInt(int64(0))
 35	b1 := big.NewInt(int64(1))
 36	b5 := big.NewInt(int64(5))
 37	aExpected := [][]*big.Int{
 38		[]*big.Int{b0, b0, b1, b0, b0, b0, b0, b0},
 39		[]*big.Int{b0, b0, b0, b1, b0, b0, b0, b0},
 40		[]*big.Int{b0, b0, b1, b0, b1, b0, b0, b0},
 41		[]*big.Int{b5, b0, b0, b0, b0, b1, b0, b0},
 42		[]*big.Int{b0, b0, b0, b0, b0, b0, b1, b0},
 43		[]*big.Int{b0, b1, b0, b0, b0, b0, b0, b0},
 44		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
 45	}
 46	bExpected := [][]*big.Int{
 47		[]*big.Int{b0, b0, b1, b0, b0, b0, b0, b0},
 48		[]*big.Int{b0, b0, b1, b0, b0, b0, b0, b0},
 49		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
 50		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
 51		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
 52		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
 53		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
 54	}
 55	cExpected := [][]*big.Int{
 56		[]*big.Int{b0, b0, b0, b1, b0, b0, b0, b0},
 57		[]*big.Int{b0, b0, b0, b0, b1, b0, b0, b0},
 58		[]*big.Int{b0, b0, b0, b0, b0, b1, b0, b0},
 59		[]*big.Int{b0, b0, b0, b0, b0, b0, b1, b0},
 60		[]*big.Int{b0, b1, b0, b0, b0, b0, b0, b0},
 61		[]*big.Int{b0, b0, b0, b0, b0, b0, b1, b0},
 62		[]*big.Int{b0, b0, b0, b0, b0, b0, b0, b1},
 63	}
 64
 65	assert.Equal(t, aExpected, a)
 66	assert.Equal(t, bExpected, b)
 67	assert.Equal(t, cExpected, c)
 68
 69	b3 := big.NewInt(int64(3))
 70	privateInputs := []*big.Int{b3}
 71	b35 := big.NewInt(int64(35))
 72	publicInputs := []*big.Int{b35}
 73	// Calculate Witness
 74	w, err := circuit.CalculateWitness(privateInputs, publicInputs)
 75	assert.Nil(t, err)
 76	b9 := big.NewInt(int64(9))
 77	b27 := big.NewInt(int64(27))
 78	b30 := big.NewInt(int64(30))
 79	wExpected := []*big.Int{b1, b35, b3, b9, b27, b30, b35, b1}
 80	assert.Equal(t, wExpected, w)
 81
 82	// circuitJson, _ := json.Marshal(circuit)
 83	// fmt.Println("circuit:", string(circuitJson))
 84
 85	assert.Equal(t, circuit.NPublic, 1)
 86	assert.Equal(t, len(circuit.PublicInputs), 1)
 87	assert.Equal(t, len(circuit.PrivateInputs), 1)
 88}
 89
 90func TestCircuitWithFuncCallsParser(t *testing.T) {
 91	// y = x^3 + x + 5
 92	code := `
 93		func exp3(private a):
 94			b = a * a
 95			c = a * b
 96			return c
 97		func sum(private a, private b):
 98			c = a + b
 99			return c
100
101		func main(private s0, public s1):
102			s3 = exp3(s0)
103			s4 = sum(s3, s0)
104			s5 = s4 + 5
105			equals(s1, s5)
106			out = 1 * 1
107	`
108	parser := NewParser(strings.NewReader(code))
109	circuit, err := parser.Parse()
110	assert.Nil(t, err)
111
112	// flat code to R1CS
113	a, b, c := circuit.GenerateR1CS()
114	assert.Equal(t, "s0", circuit.PrivateInputs[0])
115	assert.Equal(t, "s1", circuit.PublicInputs[0])
116
117	assert.Equal(t, []string{"one", "s1", "s0", "b0", "s3", "s4", "s5", "out"}, circuit.Signals)
118
119	// expected result
120	b0 := big.NewInt(int64(0))
121	b1 := big.NewInt(int64(1))
122	b5 := big.NewInt(int64(5))
123	aExpected := [][]*big.Int{
124		[]*big.Int{b0, b0, b1, b0, b0, b0, b0, b0},
125		[]*big.Int{b0, b0, b1, b0, b0, b0, b0, b0},
126		[]*big.Int{b0, b0, b1, b0, b1, b0, b0, b0},
127		[]*big.Int{b5, b0, b0, b0, b0, b1, b0, b0},
128		[]*big.Int{b0, b0, b0, b0, b0, b0, b1, b0},
129		[]*big.Int{b0, b1, b0, b0, b0, b0, b0, b0},
130		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
131	}
132	bExpected := [][]*big.Int{
133		[]*big.Int{b0, b0, b1, b0, b0, b0, b0, b0},
134		[]*big.Int{b0, b0, b0, b1, b0, b0, b0, b0},
135		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
136		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
137		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
138		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
139		[]*big.Int{b1, b0, b0, b0, b0, b0, b0, b0},
140	}
141	cExpected := [][]*big.Int{
142		[]*big.Int{b0, b0, b0, b1, b0, b0, b0, b0},
143		[]*big.Int{b0, b0, b0, b0, b1, b0, b0, b0},
144		[]*big.Int{b0, b0, b0, b0, b0, b1, b0, b0},
145		[]*big.Int{b0, b0, b0, b0, b0, b0, b1, b0},
146		[]*big.Int{b0, b1, b0, b0, b0, b0, b0, b0},
147		[]*big.Int{b0, b0, b0, b0, b0, b0, b1, b0},
148		[]*big.Int{b0, b0, b0, b0, b0, b0, b0, b1},
149	}
150
151	assert.Equal(t, aExpected, a)
152	assert.Equal(t, bExpected, b)
153	assert.Equal(t, cExpected, c)
154
155	b3 := big.NewInt(int64(3))
156	privateInputs := []*big.Int{b3}
157	b35 := big.NewInt(int64(35))
158	publicInputs := []*big.Int{b35}
159	// Calculate Witness
160	w, err := circuit.CalculateWitness(privateInputs, publicInputs)
161	assert.Nil(t, err)
162	b9 := big.NewInt(int64(9))
163	b27 := big.NewInt(int64(27))
164	b30 := big.NewInt(int64(30))
165	wExpected := []*big.Int{b1, b35, b3, b9, b27, b30, b35, b1}
166	assert.Equal(t, wExpected, w)
167
168	// circuitJson, _ := json.Marshal(circuit)
169	// fmt.Println("circuit:", string(circuitJson))
170
171	assert.Equal(t, circuit.NPublic, 1)
172	assert.Equal(t, len(circuit.PublicInputs), 1)
173	assert.Equal(t, len(circuit.PrivateInputs), 1)
174}