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		s1 = s5 * 1
 29		s5 = s1 * 1
 30		out = 1 * 1
 31	`
 32	fmt.Print("\nflat code of the circuit:")
 33	fmt.Println(flatCode)
 34
 35	// parse the code
 36	parser := circuitcompiler.NewParser(strings.NewReader(flatCode))
 37	circuit, err := parser.Parse()
 38	assert.Nil(t, err)
 39	fmt.Println("\ncircuit data:", circuit)
 40	circuitJson, _ := json.Marshal(circuit)
 41	fmt.Println("circuit:", string(circuitJson))
 42
 43	b3 := big.NewInt(int64(3))
 44	privateInputs := []*big.Int{b3}
 45	b35 := big.NewInt(int64(35))
 46	publicSignals := []*big.Int{b35}
 47
 48	// wittness
 49	w, err := circuit.CalculateWitness(privateInputs, publicSignals)
 50	assert.Nil(t, err)
 51	fmt.Println("\n", circuit.Signals)
 52	fmt.Println("witness", w)
 53
 54	// flat code to R1CS
 55	fmt.Println("\ngenerating R1CS from flat code")
 56	a, b, c := circuit.GenerateR1CS()
 57	fmt.Println("\nR1CS:")
 58	fmt.Println("a:", a)
 59	fmt.Println("b:", b)
 60	fmt.Println("c:", c)
 61
 62	// R1CS to QAP
 63	// TODO zxQAP is not used and is an old impl, bad calculated. TODO remove
 64	alphas, betas, gammas, zxQAP := Utils.PF.R1CSToQAP(a, b, c)
 65	fmt.Println("qap")
 66	fmt.Println("alphas", len(alphas))
 67	fmt.Println("alphas[1]", alphas[1])
 68	fmt.Println("betas", len(betas))
 69	fmt.Println("gammas", len(gammas))
 70	fmt.Println("zx length", len(zxQAP))
 71	assert.True(t, !bytes.Equal(alphas[1][1].Bytes(), big.NewInt(int64(0)).Bytes()))
 72
 73	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
 74	fmt.Println("ax length", len(ax))
 75	fmt.Println("bx length", len(bx))
 76	fmt.Println("cx length", len(cx))
 77	fmt.Println("px length", len(px))
 78	fmt.Println("px[last]", px[0])
 79
 80	hxQAP := Utils.PF.DivisorPolynomial(px, zxQAP)
 81	fmt.Println("hx length", len(hxQAP))
 82
 83	// hx==px/zx so px==hx*zx
 84	assert.Equal(t, px, Utils.PF.Mul(hxQAP, zxQAP))
 85
 86	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
 87	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
 88	assert.Equal(t, abc, px)
 89	hzQAP := Utils.PF.Mul(hxQAP, zxQAP)
 90	assert.Equal(t, abc, hzQAP)
 91
 92	div, rem := Utils.PF.Div(px, zxQAP)
 93	assert.Equal(t, hxQAP, div)
 94	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(6))
 95
 96	// calculate trusted setup
 97	setup, err := GenerateTrustedSetup(len(w), *circuit, alphas, betas, gammas)
 98	assert.Nil(t, err)
 99	fmt.Println("\nt:", setup.Toxic.T)
100
101	// 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
102	// assert.Equal(t, zxQAP, setup.Pk.Z)
103
104	fmt.Println("hx pk.z", hxQAP)
105	hx := Utils.PF.DivisorPolynomial(px, setup.Pk.Z)
106	fmt.Println("hx pk.z", hx)
107	// assert.Equal(t, hxQAP, hx)
108	div, rem = Utils.PF.Div(px, setup.Pk.Z)
109	assert.Equal(t, hx, div)
110	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(6))
111
112	assert.Equal(t, px, Utils.PF.Mul(hxQAP, zxQAP))
113	// hx==px/zx so px==hx*zx
114	assert.Equal(t, px, Utils.PF.Mul(hx, setup.Pk.Z))
115
116	// check length of polynomials H(x) and Z(x)
117	assert.Equal(t, len(hx), len(px)-len(setup.Pk.Z)+1)
118	assert.Equal(t, len(hxQAP), len(px)-len(zxQAP)+1)
119
120	// fmt.Println("pk.Z", len(setup.Pk.Z))
121	// fmt.Println("zxQAP", len(zxQAP))
122
123	proof, err := GenerateProofs(*circuit, setup, w, px)
124	assert.Nil(t, err)
125
126	// fmt.Println("\n proofs:")
127	// fmt.Println(proof)
128
129	// fmt.Println("public signals:", proof.PublicSignals)
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
143/*
144func TestZkMultiplication(t *testing.T) {
145
146	// compile circuit and get the R1CS
147	flatCode := `
148	func test(a, b):
149		out = a * b
150	`
151
152	// parse the code
153	parser := circuitcompiler.NewParser(strings.NewReader(flatCode))
154	circuit, err := parser.Parse()
155	assert.Nil(t, err)
156
157	b3 := big.NewInt(int64(3))
158	b4 := big.NewInt(int64(4))
159	inputs := []*big.Int{b3, b4}
160	// wittness
161	w, err := circuit.CalculateWitness(inputs)
162	assert.Nil(t, err)
163
164	fmt.Println("circuit")
165	fmt.Println(circuit.NPublic)
166
167	// flat code to R1CS
168	a, b, c := circuit.GenerateR1CS()
169	fmt.Println("\nR1CS:")
170	fmt.Println("a:", a)
171	fmt.Println("b:", b)
172	fmt.Println("c:", c)
173
174	// R1CS to QAP
175	alphas, betas, gammas, zx := Utils.PF.R1CSToQAP(a, b, c)
176	fmt.Println("qap")
177	fmt.Println("alphas", alphas)
178	fmt.Println("betas", betas)
179	fmt.Println("gammas", gammas)
180
181	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
182
183	hx := Utils.PF.DivisorPolynomial(px, zx)
184
185	// hx==px/zx so px==hx*zx
186	assert.Equal(t, px, Utils.PF.Mul(hx, zx))
187
188	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
189	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
190	assert.Equal(t, abc, px)
191	hz := Utils.PF.Mul(hx, zx)
192	assert.Equal(t, abc, hz)
193
194	div, rem := Utils.PF.Div(px, zx)
195	assert.Equal(t, hx, div)
196	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(1))
197
198	// calculate trusted setup
199	setup, err := GenerateTrustedSetup(len(w), *circuit, alphas, betas, gammas, zx)
200	assert.Nil(t, err)
201
202	// piA = g1 * A(t), piB = g2 * B(t), piC = g1 * C(t), piH = g1 * H(t)
203	proof, err := GenerateProofs(*circuit, setup, hx, w)
204	assert.Nil(t, err)
205
206	// assert.True(t, VerifyProof(*circuit, setup, proof, false))
207	b35 := big.NewInt(int64(35))
208	publicSignals := []*big.Int{b35}
209	assert.True(t, VerifyProof(*circuit, setup, proof, publicSignals, true))
210}
211*/
212/*
213func TestZkFromHardcodedR1CS(t *testing.T) {
214	b0 := big.NewInt(int64(0))
215	b1 := big.NewInt(int64(1))
216	b3 := big.NewInt(int64(3))
217	b5 := big.NewInt(int64(5))
218	b9 := big.NewInt(int64(9))
219	b27 := big.NewInt(int64(27))
220	b30 := big.NewInt(int64(30))
221	b35 := big.NewInt(int64(35))
222	a := [][]*big.Int{
223		[]*big.Int{b0, b0, b1, b0, b0, b0},
224		[]*big.Int{b0, b0, b0, b1, b0, b0},
225		[]*big.Int{b0, b0, b1, b0, b1, b0},
226		[]*big.Int{b5, b0, b0, b0, b0, b1},
227	}
228	b := [][]*big.Int{
229		[]*big.Int{b0, b0, b1, b0, b0, b0},
230		[]*big.Int{b0, b0, b1, b0, b0, b0},
231		[]*big.Int{b1, b0, b0, b0, b0, b0},
232		[]*big.Int{b1, b0, b0, b0, b0, b0},
233	}
234	c := [][]*big.Int{
235		[]*big.Int{b0, b0, b0, b1, b0, b0},
236		[]*big.Int{b0, b0, b0, b0, b1, b0},
237		[]*big.Int{b0, b0, b0, b0, b0, b1},
238		[]*big.Int{b0, b1, b0, b0, b0, b0},
239	}
240	alphas, betas, gammas, zx := Utils.PF.R1CSToQAP(a, b, c)
241
242	// wittness = 1, 35, 3, 9, 27, 30
243	w := []*big.Int{b1, b35, b3, b9, b27, b30}
244	circuit := circuitcompiler.Circuit{
245		NVars:    6,
246		NPublic:  1,
247		NSignals: len(w),
248	}
249	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
250
251	hx := Utils.PF.DivisorPolynomial(px, zx)
252
253	// hx==px/zx so px==hx*zx
254	assert.Equal(t, px, Utils.PF.Mul(hx, zx))
255
256	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
257	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
258	assert.Equal(t, abc, px)
259	hz := Utils.PF.Mul(hx, zx)
260	assert.Equal(t, abc, hz)
261
262	div, rem := Utils.PF.Div(px, zx)
263	assert.Equal(t, hx, div)
264	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(4))
265
266	// calculate trusted setup
267	setup, err := GenerateTrustedSetup(len(w), circuit, alphas, betas, gammas, zx)
268	assert.Nil(t, err)
269
270	// piA = g1 * A(t), piB = g2 * B(t), piC = g1 * C(t), piH = g1 * H(t)
271	proof, err := GenerateProofs(circuit, setup, hx, w)
272	assert.Nil(t, err)
273
274	// assert.True(t, VerifyProof(circuit, setup, proof, true))
275	publicSignals := []*big.Int{b35}
276	assert.True(t, VerifyProof(circuit, setup, proof, publicSignals, true))
277}
278
279func TestZkMultiplication(t *testing.T) {
280
281	// compile circuit and get the R1CS
282	flatCode := `
283	func test(a, b):
284		out = a * b
285	`
286
287	// parse the code
288	parser := circuitcompiler.NewParser(strings.NewReader(flatCode))
289	circuit, err := parser.Parse()
290	assert.Nil(t, err)
291
292	b3 := big.NewInt(int64(3))
293	b4 := big.NewInt(int64(4))
294	inputs := []*big.Int{b3, b4}
295	// wittness
296	w, err := circuit.CalculateWitness(inputs)
297	assert.Nil(t, err)
298
299	// flat code to R1CS
300	a, b, c := circuit.GenerateR1CS()
301
302	// R1CS to QAP
303	alphas, betas, gammas, zx := Utils.PF.R1CSToQAP(a, b, c)
304
305	ax, bx, cx, px := Utils.PF.CombinePolynomials(w, alphas, betas, gammas)
306
307	hx := Utils.PF.DivisorPolynomial(px, zx)
308
309	// hx==px/zx so px==hx*zx
310	assert.Equal(t, px, Utils.PF.Mul(hx, zx))
311
312	// p(x) = a(x) * b(x) - c(x) == h(x) * z(x)
313	abc := Utils.PF.Sub(Utils.PF.Mul(ax, bx), cx)
314	assert.Equal(t, abc, px)
315	hz := Utils.PF.Mul(hx, zx)
316	assert.Equal(t, abc, hz)
317
318	div, rem := Utils.PF.Div(px, zx)
319	assert.Equal(t, hx, div)
320	assert.Equal(t, rem, r1csqap.ArrayOfBigZeros(1))
321
322	// calculate trusted setup
323	setup, err := GenerateTrustedSetup(len(w), *circuit, alphas, betas, gammas, zx)
324	assert.Nil(t, err)
325
326	// piA = g1 * A(t), piB = g2 * B(t), piC = g1 * C(t), piH = g1 * H(t)
327	proof, err := GenerateProofs(*circuit, setup, hx, w)
328	assert.Nil(t, err)
329
330	// assert.True(t, VerifyProof(*circuit, setup, proof, false))
331	b35 := big.NewInt(int64(35))
332	publicSignals := []*big.Int{b35}
333	assert.True(t, VerifyProof(*circuit, setup, proof, publicSignals, true))
334}
335*/