feature: arbo verifier with generic hash function

arbo/hints.go

@@ -0,0 +1,43 @@
+package arbo
+
+import (
+	"fmt"
+	"math/big"
+
+	"github.com/consensys/gnark/constraint/solver"
+)
+
+func init() {
+	solver.RegisterHint(replaceSiblingHint)
+}
+
+// replaceSiblingHint gnark hint function receives the new sibling to set as
+// first input, the index of the sibling to be replaced as second input, and the
+// rest of the siblings as the rest of the inputs. The function should return
+// the new siblings with the replacement done. The caller should ensure that the
+// the result of the hint is correct checking that every sibling is the same
+// that was passed as input except for the sibling with the index provided that
+// should be replaced with the new sibling.
+func replaceSiblingHint(_ *big.Int, inputs, outputs []*big.Int) error {
+	if len(inputs) != len(outputs)+2 {
+		return fmt.Errorf("invalid number of inputs/outputs")
+	}
+	// get the new sibling and the index to replace
+	newSibling := inputs[0]
+	index := int(inputs[1].Int64())
+	if index >= len(outputs) {
+		return fmt.Errorf("invalid index")
+	}
+	siblings := inputs[2:]
+	if len(siblings) != len(outputs) {
+		return fmt.Errorf("invalid number of siblings")
+	}
+	for i := 0; i < len(outputs); i++ {
+		if i == index {
+			outputs[i] = outputs[i].Set(newSibling)
+		} else {
+			outputs[i] = outputs[i].Set(siblings[i])
+		}
+	}
+	return nil
+}

arbo/utils.go

@@ -0,0 +1,92 @@
+package arbo
+
+import (
+	"fmt"
+	"math/big"
+	"os"
+
+	arbotree "github.com/vocdoni/arbo"
+	"go.vocdoni.io/dvote/db"
+	"go.vocdoni.io/dvote/db/pebbledb"
+	"go.vocdoni.io/dvote/tree/arbo"
+	"go.vocdoni.io/dvote/util"
+)
+
+type censusConfig struct {
+	dir           string
+	validSiblings int
+	totalSiblings int
+	keyLen        int
+	hash          arbotree.HashFunction
+	baseFiled     *big.Int
+}
+
+// generateCensusProofForTest generates a census proof for testing purposes, it
+// receives a configuration and a key-value pair to generate the proof for.
+// It returns the root, key, value, and siblings of the proof. The configuration
+// includes the temp directory to store the database, the number of valid
+// siblings, the total number of siblings, the key length, the hash function to
+// use in the merkle tree, and the base field to use in the finite field.
+func generateCensusProofForTest(conf censusConfig, k, v []byte) (*big.Int, *big.Int, *big.Int, []*big.Int, error) {
+	defer func() {
+		_ = os.RemoveAll(conf.dir)
+	}()
+	database, err := pebbledb.New(db.Options{Path: conf.dir})
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+	tree, err := arbotree.NewTree(arbotree.Config{
+		Database:     database,
+		MaxLevels:    conf.totalSiblings,
+		HashFunction: conf.hash,
+	})
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+
+	k = arbotree.BigToFF(conf.baseFiled, new(big.Int).SetBytes(k)).Bytes()
+	// add the first key-value pair
+	if err = tree.Add(k, v); err != nil {
+		return nil, nil, nil, nil, err
+	}
+	// add random addresses
+	for i := 1; i < conf.validSiblings; i++ {
+		rk := arbotree.BigToFF(conf.baseFiled, new(big.Int).SetBytes(util.RandomBytes(conf.keyLen))).Bytes()
+		rv := new(big.Int).SetBytes(util.RandomBytes(8)).Bytes()
+		if err = tree.Add(rk, rv); err != nil {
+			return nil, nil, nil, nil, err
+		}
+	}
+	// generate the proof
+	_, _, siblings, exist, err := tree.GenProof(k)
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+	if !exist {
+		return nil, nil, nil, nil, fmt.Errorf("error building the merkle tree: key not found")
+	}
+	unpackedSiblings, err := arbo.UnpackSiblings(tree.HashFunction(), siblings)
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+	paddedSiblings := make([]*big.Int, conf.totalSiblings)
+	for i := 0; i < conf.totalSiblings; i++ {
+		if i < len(unpackedSiblings) {
+			paddedSiblings[i] = arbo.BytesLEToBigInt(unpackedSiblings[i])
+		} else {
+			paddedSiblings[i] = big.NewInt(0)
+		}
+	}
+	root, err := tree.Root()
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+	verified, err := arbotree.CheckProof(tree.HashFunction(), k, v, root, siblings)
+	if !verified {
+		return nil, nil, nil, nil, fmt.Errorf("error verifying the proof")
+	}
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+	return arbo.BytesLEToBigInt(root), arbo.BytesLEToBigInt(k), new(big.Int).SetBytes(v), paddedSiblings, nil
+}

arbo/verifier.go

@@ -1,23 +1,24 @@
 package arbo
 
-import (
-	"github.com/consensys/gnark/frontend"
-	"github.com/vocdoni/gnark-crypto-primitives/poseidon"
-)
+import "github.com/consensys/gnark/frontend"
 
-// 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, error) {
-	// l, r = path == 1 ? sibling, current : current, sibling
-	l, r := api.Select(ipath, sibling, leaf), api.Select(ipath, leaf, sibling)
+type Hash func(frontend.API, ...frontend.Variable) (frontend.Variable, error)
+
+// intermediateLeafKey function calculates the intermediate leaf key of the
+// path position provided. The leaf key is calculated by hashing the sibling
+// and the key provided. The position of the sibling and the key is decided by
+// the path position. If the current sibling is not valid, the method will
+// return the key provided.
+func intermediateLeafKey(api frontend.API, hFn Hash, ipath, valid, key, sibling frontend.Variable) (frontend.Variable, error) {
+	// l, r = path == 1 ? sibling, key : key, sibling
+	l, r := api.Select(ipath, sibling, key), api.Select(ipath, key, sibling)
 	// intermediateLeafKey = H(l | r)
-	intermediateLeafKey, err := poseidon.Hash(api, l, r)
+	intermediateLeafKey, err := hFn(api, l, r)
 	if err != nil {
 		return 0, err
 	}
-	// newCurrent = valid == 1 ? current : intermediateLeafKey
-	return api.Select(valid, intermediateLeafKey, leaf), nil
+	// newCurrent = valid == 1 ? intermediateLeafKey : key
+	return api.Select(valid, intermediateLeafKey, key), nil
 }
 
 // strictCmp function compares a and b and returns:
@@ -29,40 +30,110 @@ func strictCmp(api frontend.API, a, b frontend.Variable) frontend.Variable {
 }
 
 // isValid function returns 1 if the the sibling provided is a valid sibling or
-// 0 otherwise. To check if the sibling is valid, its leaf value and it must be
-// different from the previous leaf value and the previous sibling.
+// 0 otherwise. To check if the sibling is valid, its leaf key and it must be
+// different from the previous leaf key and the previous sibling.
 func isValid(api frontend.API, sibling, prevSibling, leaf, prevLeaf frontend.Variable) frontend.Variable {
 	cmp1, cmp2 := strictCmp(api, leaf, prevLeaf), strictCmp(api, sibling, prevSibling)
 	return api.Select(api.Or(cmp1, cmp2), 1, 0)
 }
 
-// CheckProof receives the parameters of a proof of Arbo to recalculate the
-// root with them and compare it with the provided one, verifiying the proof.
-func CheckProof(api frontend.API, key, value, root frontend.Variable, siblings []frontend.Variable) error {
+// replaceFirstPaddedSibling function replaces the first padded sibling with the
+// new sibling provided. The function receives the new sibling, the siblings and
+// returns the new siblings with the replacement done. It first calculates the
+// index of the first padded sibling and then calls the hint function to replace
+// it. The hint function should return the new siblings with the replacement
+// done. The function ensures that the replacement was done correctly.
+func replaceFirstPaddedSibling(api frontend.API, newSibling frontend.Variable, siblings []frontend.Variable) ([]frontend.Variable, error) {
+	// the valid siblins are always the first n siblings that are not zero, so
+	// we need to iterate through the siblings in reverse order to find the
+	// first non-zero sibling and count the number of valid siblings from there,
+	// so the index of the last padded sibling is the number of valid siblings
+	index := frontend.Variable(0)
+	nonZeroFound := frontend.Variable(0)
+	for i := len(siblings) - 1; i >= 0; i-- {
+		isNotZero := strictCmp(api, siblings[i], 0)
+		nonZeroFound = api.Or(nonZeroFound, isNotZero)
+		index = api.Add(index, nonZeroFound)
+	}
+	// call the hint function to replace the sibling with the index to be
+	// replaced, the new sibling and the rest of the siblings
+	newSiblings, err := api.Compiler().NewHint(replaceSiblingHint, len(siblings),
+		append([]frontend.Variable{newSibling, index}, siblings...)...)
+	if err != nil {
+		return nil, err
+	}
+	// check that the hint successfully replaced the first padded sibling
+	newSiblingFound := frontend.Variable(0)
+	for i := 0; i < len(newSiblings); i++ {
+		correctIndex := api.IsZero(api.Sub(index, frontend.Variable(i)))
+		correctSibling := api.IsZero(api.Sub(newSiblings[i], newSibling))
+		newSiblingFound = api.Or(newSiblingFound, api.And(correctIndex, correctSibling))
+	}
+	api.AssertIsEqual(newSiblingFound, 1)
+	return newSiblings, nil
+}
+
+// CheckInclusionProof receives the parameters of an inclusion proof of Arbo to
+// recalculate the root with them and compare it with the provided one,
+// verifiying the proof.
+func CheckInclusionProof(api frontend.API, hFn Hash, key, value, root frontend.Variable,
+	siblings []frontend.Variable,
+) error {
 	// 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
-	//   leafValue = H(key | value | 1)
-	leafValue, err := poseidon.Hash(api, key, value, 1)
+	path := api.ToBinary(key, len(siblings))
+	// calculate the current leaf key to start with it to rebuild the tree
+	//   leafKey = H(key | value | 1)
+	leafKey, err := hFn(api, key, value, 1)
 	if err != nil {
 		return err
 	}
-	// calculate the root and compare it with the provided one
-	prevLeaf := leafValue
-	currentLeaf := leafValue
-	prevSibling := frontend.Variable(0)
+	// calculate the root iterating through the siblings in inverse order,
+	// calculating the intermediate leaf key based on the path and the validity
+	// of the current sibling
+	prevKey := leafKey                  // init prevKey with computed leafKey
+	prevSibling := frontend.Variable(0) // init prevSibling with 0
 	for i := len(siblings) - 1; i >= 0; i-- {
 		// check if the sibling is valid
-		valid := isValid(api, siblings[i], prevSibling, currentLeaf, prevLeaf)
-		prevLeaf = currentLeaf
-		prevSibling = siblings[i]
-		// compute the next leaf value
-		currentLeaf, err = prevLevel(api, currentLeaf, path[i], valid, siblings[i])
+		valid := isValid(api, siblings[i], prevSibling, leafKey, prevKey)
+		prevKey = leafKey         // update prevKey to the lastKey
+		prevSibling = siblings[i] // update prevSibling to the current sibling
+		// compute the intermediate leaf key and update the lastKey
+		leafKey, err = intermediateLeafKey(api, hFn, path[i], valid, leafKey, siblings[i])
 		if err != nil {
 			return err
 		}
 	}
-	api.AssertIsEqual(currentLeaf, root)
+	api.AssertIsEqual(leafKey, root)
 	return nil
 }
+
+// CheckAdditionProof method proves that the addition of a new key-value pair
+// to the tree is valid. It gets the root (old root), the key-value pair of the
+// first sibling of the new key-value pair (the old leaf key), and the rest of
+// the siblings to check. It also gets the root after including the new
+// key-value pair (the new leaf key) and the new key-value pair itself. It
+// includes the old leaf key as a sibling of the new leaf key and verifies the
+// proof of the new root. In this way, it ensures that the addition of the new
+// key-value pair is valid by checking the state of the tree before and after
+// the addition.
+func CheckAdditionProof(api frontend.API, hFn Hash, key, value, root, oldKey, oldValue,
+	oldRoot frontend.Variable, siblings []frontend.Variable,
+) error {
+	// check that the old proof is valid
+	if err := CheckInclusionProof(api, hFn, oldKey, oldValue, oldRoot, siblings); err != nil {
+		return err
+	}
+	// calculate the old leaf key (as new sibling)
+	newLeafKey, err := hFn(api, oldKey, oldValue, 1)
+	if err != nil {
+		return err
+	}
+	// include the old leaf key as a sibling of the new leaf key
+	newSiblings, err := replaceFirstPaddedSibling(api, newLeafKey, siblings)
+	if err != nil {
+		return err
+	}
+	// verify the proof of the new leaf key
+	return CheckInclusionProof(api, hFn, key, value, root, newSiblings)
+}

arbo/verifier_bls12377_test.go

@@ -0,0 +1,79 @@
+package arbo
+
+import (
+	"fmt"
+	"math/big"
+	"testing"
+	"time"
+
+	"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/std/hash/mimc"
+	"github.com/consensys/gnark/test"
+	qt "github.com/frankban/quicktest"
+	arbotree "github.com/vocdoni/arbo"
+	"go.vocdoni.io/dvote/util"
+)
+
+const (
+	v_siblings = 10
+	n_siblings = 160
+	k_len      = n_siblings / 8
+)
+
+type testVerifierBLS12377 struct {
+	Root     frontend.Variable
+	Key      frontend.Variable
+	Value    frontend.Variable
+	Siblings [n_siblings]frontend.Variable
+}
+
+func (circuit *testVerifierBLS12377) Define(api frontend.API) error {
+	// use mimc hash function
+	hash := func(api frontend.API, data ...frontend.Variable) (frontend.Variable, error) {
+		h, err := mimc.NewMiMC(api)
+		if err != nil {
+			return 0, err
+		}
+		h.Write(data...)
+		return h.Sum(), nil
+	}
+	return CheckInclusionProof(api, hash, circuit.Key, circuit.Value, circuit.Root, circuit.Siblings[:])
+}
+
+func TestVerifierBLS12377(t *testing.T) {
+	c := qt.New(t)
+	// profile the circuit compilation
+	p := profile.Start()
+	now := time.Now()
+	_, _ = frontend.Compile(ecc.BLS12_377.ScalarField(), r1cs.NewBuilder, &testVerifierBLS12377{})
+	fmt.Println("elapsed", time.Since(now))
+	p.Stop()
+	fmt.Println("constrains", p.NbConstraints())
+	// generate census proof
+	root, key, value, siblings, err := generateCensusProofForTest(censusConfig{
+		dir:           t.TempDir() + "/bls12377",
+		validSiblings: v_siblings,
+		totalSiblings: n_siblings,
+		keyLen:        k_len,
+		hash:          arbotree.HashFunctionMiMC_BLS12_377,
+		baseFiled:     arbotree.BLS12377BaseField,
+	}, util.RandomBytes(k_len), big.NewInt(10).Bytes())
+	c.Assert(err, qt.IsNil)
+	// init and print inputs
+	fSiblings := [n_siblings]frontend.Variable{}
+	for i := 0; i < n_siblings; i++ {
+		fSiblings[i] = siblings[i]
+	}
+	inputs := testVerifierBLS12377{
+		Root:     root,
+		Key:      key,
+		Value:    value,
+		Siblings: fSiblings,
+	}
+	assert := test.NewAssert(t)
+	assert.SolvingSucceeded(&testVerifierBLS12377{}, &inputs, test.WithCurves(ecc.BLS12_377), test.WithBackends(backend.GROTH16))
+}