From 2b2fe7be70c9593f86e438998009ccf4100efacc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lucas=20Men=C3=A9ndez?= <hi@lucasmenendez.me>
Date: Tue, 5 Nov 2024 18:09:56 +0100
Subject: [PATCH] code simplified

---
 arbo/verifier.go      | 64 ++++++++++++-------------------------------
 arbo/verifier_test.go |  8 ++++++
 2 files changed, 26 insertions(+), 46 deletions(-)

diff --git a/arbo/verifier.go b/arbo/verifier.go
index 8f9ffee..afad81d 100644
--- a/arbo/verifier.go
+++ b/arbo/verifier.go
@@ -6,49 +6,16 @@ import (
 	"github.com/vocdoni/gnark-crypto-primitives/poseidon"
 )
 
-// leafValueKey returns the encoded childless leaf value for the key-value pair
-// provided, hashing it with the predefined hashing function 'H':
-//
-//	newLeafValue = H(key | value | 1)
-func leafValueKey(api frontend.API, key, value frontend.Variable) frontend.Variable {
-	return poseidon.Hash(api, key, value, 1)
-}
-
-// intermediateLeafKey returns the encoded intermediate leaf value for the
-// key-value pair provided, hashing it with the predefined hashing function 'H':
-//
-//	intermediateLeafValue = H(l | r)
-func intermediateLeafKey(api frontend.API, l, r frontend.Variable) frontend.Variable {
-	return poseidon.Hash(api, l, r)
-}
-
-// switcher returns the l and r parameters swiped if sel is 1, else returns it
-// in the same position:
-//
-//	switcher(1, l, r) == r, l
-//	switcher(0, l, r) == l, r
-func switcher(api frontend.API, sel, l, r frontend.Variable) (outL, outR frontend.Variable) {
-	// aux <== (R-L)*sel;
-	aux := api.Mul(api.Sub(r, l), sel)
-	// outL <==  aux + L;
-	outL = api.Add(aux, l)
-	// outR <== -aux + R;
-	outR = api.Sub(r, aux)
-	return
-}
-
-// calcRoot iterates over valid siblings recursively calculating the root hash
-// with the intermediate leaf keys. To do that, it requires the key path in the
-// tree and the map of valid siblings.
-func calcRoot(api frontend.API, current frontend.Variable, path, valid, siblings []frontend.Variable) frontend.Variable {
-	i := len(siblings) - 1
-	l, r := switcher(api, path[i], current, siblings[i])
-	hash := intermediateLeafKey(api, l, r)
-	newCurrent, _ := switcher(api, valid[i], current, hash)
-	if i == 0 {
-		return newCurrent
-	}
-	return calcRoot(api, newCurrent, path[:i], valid[:i], siblings[:i])
+// prevLevel function calculates the previous level of the merkle tree given the
+// current leaf, the current path bit of the leaf, the validity of the sibling
+// and the sibling itself.
+func prevLevel(api frontend.API, leaf, ipath, valid, sibling frontend.Variable) frontend.Variable {
+	// l, r = path == 1 ? sibling, current : current, sibling
+	l, r := api.Select(ipath, sibling, leaf), api.Select(ipath, leaf, sibling)
+	// intermediateLeafKey = H(l | r)
+	intermediateLeafKey := poseidon.Hash(api, l, r)
+	// newCurrent = valid == 1 ? current : intermediateLeafKey
+	return api.Select(valid, intermediateLeafKey, leaf)
 }
 
 // validSiblings function creates a binary map with the slots where a valid
@@ -70,12 +37,17 @@ func CheckProof(api frontend.API, key, value, root, nsiblings frontend.Variable,
 	api.AssertIsLessOrEqual(nsiblings, len(siblings))
 	// get a map with the valid siblings
 	valid := validSiblings(api, siblings, nsiblings)
-	// calculta the path from the provided key to decide which leaf is the
+	// calculate the path from the provided key to decide which leaf is the
 	// correct one in every level of the tree
 	path := api.ToBinary(key, api.Compiler().FieldBitLen())
 	// calculate the value leaf to start with it to rebuild the tree
-	firstLevel := leafValueKey(api, key, value)
+	//   leafValue = H(key | value | 1)
+	leafValue := poseidon.Hash(api, key, value, 1)
 	// calculate the root and compare it with the provided one
-	api.AssertIsEqual(calcRoot(api, firstLevel, path, valid, siblings), root)
+	currentLevel := leafValue
+	for i := len(siblings) - 1; i >= 0; i-- {
+		currentLevel = prevLevel(api, currentLevel, path[i], valid[i], siblings[i])
+	}
+	api.AssertIsEqual(currentLevel, root)
 	return nil
 }
diff --git a/arbo/verifier_test.go b/arbo/verifier_test.go
index decc766..24018ff 100644
--- a/arbo/verifier_test.go
+++ b/arbo/verifier_test.go
@@ -1,12 +1,15 @@
 package arbo
 
 import (
+	"fmt"
 	"math/big"
 	"testing"
 
 	"github.com/consensys/gnark-crypto/ecc"
 	"github.com/consensys/gnark/backend"
 	"github.com/consensys/gnark/frontend"
+	"github.com/consensys/gnark/frontend/cs/r1cs"
+	"github.com/consensys/gnark/profile"
 	"github.com/consensys/gnark/test"
 	qt "github.com/frankban/quicktest"
 	"go.vocdoni.io/dvote/db"
@@ -80,6 +83,11 @@ func successInputs(t *testing.T, n int) testVerifierCircuit {
 }
 
 func TestVerifier(t *testing.T) {
+	p := profile.Start()
+	_, _ = frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &testVerifierCircuit{})
+	p.Stop()
+	fmt.Println("constrains", p.NbConstraints())
+
 	assert := test.NewAssert(t)
 
 	inputs := successInputs(t, 10)