1package snark
  2
  3import (
  4	"bytes"
  5	"encoding/json"
  6	"fmt"
  7	"math/big"
  8	"strings"
  9	"testing"
 10	"time"
 11
 12	"github.com/arnaucube/go-snark/circuitcompiler"
 13	"github.com/arnaucube/go-snark/r1csqap"
 14	"github.com/stretchr/testify/assert"
 15)
 16
 17func TestZkFromFlatCircuitCode(t *testing.T) {
 18	// compile circuit and get the R1CS
 19
 20	// circuit function
 21	// y = x^3 + x + 5
 22	flatCode := `
 23	func test(private s0, public s1):
 24		s2 = s0 * s0
 25		s3 = s2 * s0
 26		s4 = s3 + s0
 27		s5 = s4 + 5
 28		equals(s1, s5)
 29		out = 1 * 1
 30	`
 31	fmt.Print("\nflat code of the circuit:")
 32	fmt.Println(flatCode)
 33
 34	// parse the code
 35	parser := circuitcompiler.NewParser(strings.NewReader(flatCode))
 36	circuit, err := parser.Parse()
 37	assert.Nil(t, err)
 38	fmt.Println("\ncircuit data:", circuit)
 39	circuitJson, _ := json.Marshal(circuit)
 40	fmt.Println("circuit:", string(circuitJson))
 41
 42	b3 := big.NewInt(int64(3))
 43	privateInputs := []*big.Int{b3}
 44	b35 := big.NewInt(int64(35))
 45	publicSignals := []*big.Int{b35}
 46
 47	// wittness
 48	w, err := circuit.CalculateWitness(privateInputs, publicSignals)
 49	assert.Nil(t, err)
 50	fmt.Println("\n", circuit.Signals)
 51	fmt.Println("witness", w)
 52
 53	// flat code to R1CS
 54	fmt.Println("\ngenerating R1CS from flat code")
 55	a, b, c := circuit.GenerateR1CS()
 56	fmt.Println("\nR1CS:")
 57	fmt.Println("a:", a)
 58	fmt.Println("b:", b)
 59	fmt.Println("c:", c)
 60
 61	// R1CS to QAP
 62	// TODO zxQAP is not used and is an old impl, bad calculated. TODO remove
 63	alphas, betas, gammas, zxQAP := Utils.PF.R1CSToQAP(a, b, c)
 64	fmt.Println("qap")
 65	fmt.Println("alphas", len(alphas))
 66	fmt.Println("alphas[1]", alphas[1])
 67	fmt.Println("betas", len(betas))
 68	fmt.Println("gammas", len(gammas))
 69	fmt.Println("zx length", len(zxQAP))
 70	assert.True(t, !bytes.Equal(alphas[1][1].Bytes(), big.NewInt(int64(0)).Bytes()))
 71
 72	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
 73	fmt.Println("ax length", len(ax))
 74	fmt.Println("bx length", len(bx))
 75	fmt.Println("cx length", len(cx))
 76	fmt.Println("px length", len(px))
 77	fmt.Println("px[last]", px[0])
 78
 79	hxQAP := Utils.PF.DivisorPolynomial(px, zxQAP)
 80	fmt.Println("hx length", len(hxQAP))
 81
 82	// hx==px/zx so px==hx*zx
 83	assert.Equal(t, px, Utils.PF.Mul(hxQAP, zxQAP))
 84
 85	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
 86	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
 87	assert.Equal(t, abc, px)
 88	hzQAP := Utils.PF.Mul(hxQAP, zxQAP)
 89	assert.Equal(t, abc, hzQAP)
 90
 91	div, rem := Utils.PF.Div(px, zxQAP)
 92	assert.Equal(t, hxQAP, div)
 93	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(6))
 94
 95	// calculate trusted setup
 96	setup, err := GenerateTrustedSetup(len(w), *circuit, alphas, betas, gammas)
 97	assert.Nil(t, err)
 98	fmt.Println("\nt:", setup.Toxic.T)
 99
100	// zx and setup.Pk.Z should be the same (currently not, the correct one is the calculation used inside GenerateTrustedSetup function), the calculation is repeated. TODO avoid repeating calculation
101	// assert.Equal(t, zxQAP, setup.Pk.Z)
102
103	fmt.Println("hx pk.z", hxQAP)
104	hx := Utils.PF.DivisorPolynomial(px, setup.Pk.Z)
105	fmt.Println("hx pk.z", hx)
106	// assert.Equal(t, hxQAP, hx)
107	div, rem = Utils.PF.Div(px, setup.Pk.Z)
108	assert.Equal(t, hx, div)
109	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(6))
110
111	assert.Equal(t, px, Utils.PF.Mul(hxQAP, zxQAP))
112	// hx==px/zx so px==hx*zx
113	assert.Equal(t, px, Utils.PF.Mul(hx, setup.Pk.Z))
114
115	// check length of polynomials H(x) and Z(x)
116	assert.Equal(t, len(hx), len(px)-len(setup.Pk.Z)+1)
117	assert.Equal(t, len(hxQAP), len(px)-len(zxQAP)+1)
118
119	// fmt.Println("pk.Z", len(setup.Pk.Z))
120	// fmt.Println("zxQAP", len(zxQAP))
121
122	proof, err := GenerateProofs(*circuit, setup, w, px)
123	assert.Nil(t, err)
124
125	// fmt.Println("\n proofs:")
126	// fmt.Println(proof)
127
128	// fmt.Println("public signals:", proof.PublicSignals)
129	fmt.Println("\n", circuit.Signals)
130	fmt.Println("\nwitness", w)
131	b35Verif := big.NewInt(int64(35))
132	publicSignalsVerif := []*big.Int{b35Verif}
133	before := time.Now()
134	assert.True(t, VerifyProof(*circuit, setup, proof, publicSignalsVerif, true))
135	fmt.Println("verify proof time elapsed:", time.Since(before))
136
137	// check that with another public input the verification returns false
138	bOtherWrongPublic := big.NewInt(int64(34))
139	wrongPublicSignalsVerif := []*big.Int{bOtherWrongPublic}
140	assert.True(t, !VerifyProof(*circuit, setup, proof, wrongPublicSignalsVerif, true))
141}
142
143func TestZkMultiplication(t *testing.T) {
144	flatCode := `
145	func test(private a, private b, public c):
146		d = a * b
147		equals(c, d)
148		out = 1 * 1
149	`
150	fmt.Print("\nflat code of the circuit:")
151	fmt.Println(flatCode)
152
153	// parse the code
154	parser := circuitcompiler.NewParser(strings.NewReader(flatCode))
155	circuit, err := parser.Parse()
156	assert.Nil(t, err)
157	fmt.Println("\ncircuit data:", circuit)
158	circuitJson, _ := json.Marshal(circuit)
159	fmt.Println("circuit:", string(circuitJson))
160
161	b3 := big.NewInt(int64(3))
162	b4 := big.NewInt(int64(4))
163	privateInputs := []*big.Int{b3, b4}
164	b12 := big.NewInt(int64(12))
165	publicSignals := []*big.Int{b12}
166
167	// wittness
168	w, err := circuit.CalculateWitness(privateInputs, publicSignals)
169	assert.Nil(t, err)
170	fmt.Println("\n", circuit.Signals)
171	fmt.Println("witness", w)
172
173	// flat code to R1CS
174	fmt.Println("\ngenerating R1CS from flat code")
175	a, b, c := circuit.GenerateR1CS()
176	fmt.Println("\nR1CS:")
177	fmt.Println("a:", a)
178	fmt.Println("b:", b)
179	fmt.Println("c:", c)
180
181	// R1CS to QAP
182	// TODO zxQAP is not used and is an old impl, bad calculated. TODO remove
183	alphas, betas, gammas, zxQAP := Utils.PF.R1CSToQAP(a, b, c)
184	fmt.Println("qap")
185	fmt.Println("alphas", len(alphas))
186	fmt.Println("alphas[1]", alphas[1])
187	fmt.Println("betas", len(betas))
188	fmt.Println("gammas", len(gammas))
189	fmt.Println("zx length", len(zxQAP))
190	assert.True(t, !bytes.Equal(alphas[1][1].Bytes(), big.NewInt(int64(0)).Bytes()))
191
192	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
193	fmt.Println("ax length", len(ax))
194	fmt.Println("bx length", len(bx))
195	fmt.Println("cx length", len(cx))
196	fmt.Println("px length", len(px))
197	fmt.Println("px[last]", px[0])
198
199	hxQAP := Utils.PF.DivisorPolynomial(px, zxQAP)
200	fmt.Println("hx length", len(hxQAP))
201
202	// hx==px/zx so px==hx*zx
203	assert.Equal(t, px, Utils.PF.Mul(hxQAP, zxQAP))
204
205	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
206	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
207	assert.Equal(t, abc, px)
208	hzQAP := Utils.PF.Mul(hxQAP, zxQAP)
209	assert.Equal(t, abc, hzQAP)
210
211	div, rem := Utils.PF.Div(px, zxQAP)
212	assert.Equal(t, hxQAP, div)
213	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(4))
214
215	// calculate trusted setup
216	setup, err := GenerateTrustedSetup(len(w), *circuit, alphas, betas, gammas)
217	assert.Nil(t, err)
218	fmt.Println("\nt:", setup.Toxic.T)
219
220	// zx and setup.Pk.Z should be the same (currently not, the correct one is the calculation used inside GenerateTrustedSetup function), the calculation is repeated. TODO avoid repeating calculation
221	// assert.Equal(t, zxQAP, setup.Pk.Z)
222
223	fmt.Println("hx pk.z", hxQAP)
224	hx := Utils.PF.DivisorPolynomial(px, setup.Pk.Z)
225	fmt.Println("hx pk.z", hx)
226	// assert.Equal(t, hxQAP, hx)
227	div, rem = Utils.PF.Div(px, setup.Pk.Z)
228	assert.Equal(t, hx, div)
229	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(4))
230
231	assert.Equal(t, px, Utils.PF.Mul(hxQAP, zxQAP))
232	// hx==px/zx so px==hx*zx
233	assert.Equal(t, px, Utils.PF.Mul(hx, setup.Pk.Z))
234
235	// check length of polynomials H(x) and Z(x)
236	assert.Equal(t, len(hx), len(px)-len(setup.Pk.Z)+1)
237	assert.Equal(t, len(hxQAP), len(px)-len(zxQAP)+1)
238
239	// fmt.Println("pk.Z", len(setup.Pk.Z))
240	// fmt.Println("zxQAP", len(zxQAP))
241
242	proof, err := GenerateProofs(*circuit, setup, w, px)
243	assert.Nil(t, err)
244
245	// fmt.Println("\n proofs:")
246	// fmt.Println(proof)
247
248	// fmt.Println("public signals:", proof.PublicSignals)
249	fmt.Println("\n", circuit.Signals)
250	fmt.Println("\nwitness", w)
251	b12Verif := big.NewInt(int64(12))
252	publicSignalsVerif := []*big.Int{b12Verif}
253	before := time.Now()
254	assert.True(t, VerifyProof(*circuit, setup, proof, publicSignalsVerif, true))
255	fmt.Println("verify proof time elapsed:", time.Since(before))
256
257	// check that with another public input the verification returns false
258	bOtherWrongPublic := big.NewInt(int64(11))
259	wrongPublicSignalsVerif := []*big.Int{bOtherWrongPublic}
260	assert.True(t, !VerifyProof(*circuit, setup, proof, wrongPublicSignalsVerif, true))
261}
262
263/*
264func TestZkFromHardcodedR1CS(t *testing.T) {
265	b0 := big.NewInt(int64(0))
266	b1 := big.NewInt(int64(1))
267	b3 := big.NewInt(int64(3))
268	b5 := big.NewInt(int64(5))
269	b9 := big.NewInt(int64(9))
270	b27 := big.NewInt(int64(27))
271	b30 := big.NewInt(int64(30))
272	b35 := big.NewInt(int64(35))
273	a := [][]*big.Int{
274		[]*big.Int{b0, b0, b1, b0, b0, b0},
275		[]*big.Int{b0, b0, b0, b1, b0, b0},
276		[]*big.Int{b0, b0, b1, b0, b1, b0},
277		[]*big.Int{b5, b0, b0, b0, b0, b1},
278	}
279	b := [][]*big.Int{
280		[]*big.Int{b0, b0, b1, b0, b0, b0},
281		[]*big.Int{b0, b0, b1, b0, b0, b0},
282		[]*big.Int{b1, b0, b0, b0, b0, b0},
283		[]*big.Int{b1, b0, b0, b0, b0, b0},
284	}
285	c := [][]*big.Int{
286		[]*big.Int{b0, b0, b0, b1, b0, b0},
287		[]*big.Int{b0, b0, b0, b0, b1, b0},
288		[]*big.Int{b0, b0, b0, b0, b0, b1},
289		[]*big.Int{b0, b1, b0, b0, b0, b0},
290	}
291	alphas, betas, gammas, zx := Utils.PF.R1CSToQAP(a, b, c)
292
293	// wittness = 1, 35, 3, 9, 27, 30
294	w := []*big.Int{b1, b35, b3, b9, b27, b30}
295	circuit := circuitcompiler.Circuit{
296		NVars:    6,
297		NPublic:  1,
298		NSignals: len(w),
299	}
300	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
301
302	hx := Utils.PF.DivisorPolynomial(px, zx)
303
304	// hx==px/zx so px==hx*zx
305	assert.Equal(t, px, Utils.PF.Mul(hx, zx))
306
307	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
308	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
309	assert.Equal(t, abc, px)
310	hz := Utils.PF.Mul(hx, zx)
311	assert.Equal(t, abc, hz)
312
313	div, rem := Utils.PF.Div(px, zx)
314	assert.Equal(t, hx, div)
315	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(4))
316
317	// calculate trusted setup
318	setup, err := GenerateTrustedSetup(len(w), circuit, alphas, betas, gammas, zx)
319	assert.Nil(t, err)
320
321	// piA = g1 * A(t), piB = g2 * B(t), piC = g1 * C(t), piH = g1 * H(t)
322	proof, err := GenerateProofs(circuit, setup, hx, w)
323	assert.Nil(t, err)
324
325	// assert.True(t, VerifyProof(circuit, setup, proof, true))
326	publicSignals := []*big.Int{b35}
327	assert.True(t, VerifyProof(circuit, setup, proof, publicSignals, true))
328}
329
330func TestZkMultiplication(t *testing.T) {
331
332	// compile circuit and get the R1CS
333	flatCode := `
334	func test(a, b):
335		out = a * b
336	`
337
338	// parse the code
339	parser := circuitcompiler.NewParser(strings.NewReader(flatCode))
340	circuit, err := parser.Parse()
341	assert.Nil(t, err)
342
343	b3 := big.NewInt(int64(3))
344	b4 := big.NewInt(int64(4))
345	inputs := []*big.Int{b3, b4}
346	// wittness
347	w, err := circuit.CalculateWitness(inputs)
348	assert.Nil(t, err)
349
350	// flat code to R1CS
351	a, b, c := circuit.GenerateR1CS()
352
353	// R1CS to QAP
354	alphas, betas, gammas, zx := Utils.PF.R1CSToQAP(a, b, c)
355
356	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
357
358	hx := Utils.PF.DivisorPolynomial(px, zx)
359
360	// hx==px/zx so px==hx*zx
361	assert.Equal(t, px, Utils.PF.Mul(hx, zx))
362
363	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
364	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
365	assert.Equal(t, abc, px)
366	hz := Utils.PF.Mul(hx, zx)
367	assert.Equal(t, abc, hz)
368
369	div, rem := Utils.PF.Div(px, zx)
370	assert.Equal(t, hx, div)
371	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(1))
372
373	// calculate trusted setup
374	setup, err := GenerateTrustedSetup(len(w), *circuit, alphas, betas, gammas, zx)
375	assert.Nil(t, err)
376
377	// piA = g1 * A(t), piB = g2 * B(t), piC = g1 * C(t), piH = g1 * H(t)
378	proof, err := GenerateProofs(*circuit, setup, hx, w)
379	assert.Nil(t, err)
380
381	// assert.True(t, VerifyProof(*circuit, setup, proof, false))
382	b35 := big.NewInt(int64(35))
383	publicSignals := []*big.Int{b35}
384	assert.True(t, VerifyProof(*circuit, setup, proof, publicSignals, true))
385}
386*/