From 866386a98225b6be9345cbeaa6c9ac91b1d9eb38 Mon Sep 17 00:00:00 2001
From: Gui Iribarren <gui@altermundi.net>
Date: Mon, 9 Dec 2024 16:19:44 +0100
Subject: [PATCH] finish working PoC

---
 circuits/statetransition/circuit.go | 30 +++--------------------------
 circuits/test/state.go              | 26 +++++++++++++++++++++----
 2 files changed, 25 insertions(+), 31 deletions(-)

diff --git a/circuits/statetransition/circuit.go b/circuits/statetransition/circuit.go
index d6867c0..b319278 100644
--- a/circuits/statetransition/circuit.go
+++ b/circuits/statetransition/circuit.go
@@ -118,6 +118,9 @@ func (circuit Circuit) verifyBallots(api frontend.API) {
 	var ballotCount, overwrittenCount frontend.Variable = 0, 0
 
 	for _, b := range circuit.Ballot {
+		// TODO: check that Hash(NewCiphertext) matches b.NewValue
+		// and Hash(OldCiphertext) matches b.OldValue
+
 		ballotSum.Add(api, ballotSum,
 			elgamal.NewCiphertext().Select(api, b.IsInsertOrUpdate(api), &b.NewCiphertext, zero))
 
@@ -135,30 +138,3 @@ func (circuit Circuit) verifyBallots(api frontend.API) {
 	api.AssertIsEqual(circuit.NumNewVotes, ballotCount)
 	api.AssertIsEqual(circuit.NumOverwrites, overwrittenCount)
 }
-
-// // verifyMockBallots counts the ballots using a simple api.Add
-// //
-// // (not intended for production, just for testing purposes)
-// func (circuit Circuit) verifyMockBallots(api frontend.API) {
-// 	var ballotSum, overwrittenSum, ballotCount, overwrittenCount frontend.Variable = 0, 0, 0, 0
-
-// 	for _, b := range circuit.Ballot {
-// 		ballotSum = api.Add(ballotSum, api.Select(api.Or(isUpdate(api, b), isInsert(api, b)),
-// 			b.NewValue, 0))
-// 		overwrittenSum = api.Add(overwrittenSum, api.Select(isUpdate(api, b),
-// 			b.OldValue, 0))
-// 		ballotCount = api.Add(ballotCount, api.Select(api.Or(isUpdate(api, b), isInsert(api, b)),
-// 			1, 0))
-// 		overwrittenCount = api.Add(overwrittenCount, api.Select(isUpdate(api, b),
-// 			1, 0))
-// 	}
-
-// 	api.AssertIsEqual(
-// 		api.Add(circuit.ResultsAdd.OldValue, ballotSum),
-// 		circuit.ResultsAdd.NewValue)
-// 	api.AssertIsEqual(
-// 		api.Add(circuit.ResultsSub.OldValue, overwrittenSum),
-// 		circuit.ResultsSub.NewValue)
-// 	api.AssertIsEqual(circuit.NumNewVotes, ballotCount)
-// 	api.AssertIsEqual(circuit.NumOverwrites, overwrittenCount)
-// }
diff --git a/circuits/test/state.go b/circuits/test/state.go
index 204fb5b..49deae5 100644
--- a/circuits/test/state.go
+++ b/circuits/test/state.go
@@ -12,9 +12,25 @@ import (
 	"go.vocdoni.io/dvote/tree/arbo"
 )
 
-// absolute hack, need to reimplement this, store the ciphertexts <-> hash relationship somewhere else
+// absolute hack, need to reimplement this:
+// store the serialized ciphertext in arbo tree, (something simple, concat the bytes)
+// and inside the circuit check the hash(serializedCiphertext) = Leaf
 var memDB map[string]*encrypt.ElGamalCiphertext
 
+func init() {
+	memDB = make(map[string]*encrypt.ElGamalCiphertext)
+}
+
+func oldVote(nullifier []byte) *encrypt.ElGamalCiphertext {
+	return memDB[string(nullifier)]
+}
+
+func storeVote(nullifier []byte, vote *encrypt.ElGamalCiphertext) {
+	memDB[string(nullifier)] = vote
+}
+
+// end absolute hack
+
 var hashFunc = arbo.HashFunctionPoseidon
 
 var (
@@ -126,8 +142,8 @@ func (o *State) AddVote(v Vote) error {
 
 	// if nullifier exists, it's a vote overwrite, need to count the overwritten vote
 	// so it's later added to circuit.ResultsSub
-	if _, v, err := o.tree.Get(v.nullifier); err == nil {
-		o.overwriteSum.Add(o.overwriteSum, memDB[string(v)])
+	if _, _, err := o.tree.Get(v.nullifier); err == nil {
+		o.overwriteSum.Add(o.overwriteSum, oldVote(v.nullifier))
 		o.overwriteCount++
 	}
 
@@ -151,10 +167,10 @@ func (o *State) EndBatch() error {
 	// add Ballots
 	for i := range o.Witnesses.Ballot {
 		if i < len(o.votes) {
+			o.Witnesses.Ballot[i].OldCiphertext = o.votes[i].elgamalBallot.ToGnark() // mock
 			o.Witnesses.Ballot[i].MerkleTransition, err = statetransition.MerkleTransitionFromAddOrUpdate(o.tree,
 				o.votes[i].nullifier, arbo.BigIntToBytesLE(32, &o.votes[i].ballot))
 			o.Witnesses.Ballot[i].NewCiphertext = o.votes[i].elgamalBallot.ToGnark()
-			o.Witnesses.Ballot[i].OldCiphertext = o.votes[i].elgamalBallot.ToGnark() // mock
 		} else {
 			o.Witnesses.Ballot[i], err = statetransition.MerkleTransitionElGamalFromNoop(o.tree)
 		}
@@ -255,6 +271,8 @@ func NewVote(nullifier, amount uint64) Vote {
 
 	v.elgamalBallot = NewEncryptedBallot(amount)
 
+	storeVote(v.nullifier, v.elgamalBallot)
+
 	v.address = arbo.BigIntToBytesLE(statetransition.MaxKeyLen,
 		big.NewInt(int64(nullifier)+int64(KeyAddressesOffset))) // mock
 	v.commitment.SetUint64(amount + 256) // mock