Add DumpLeafs & ImportDumpedLeafs for MerkleTree backup functionality

merkletree.go

@@ -76,6 +76,22 @@ func (h *Hash) BigInt() *big.Int {
 	return new(big.Int).SetBytes(common.SwapEndianness(h[:]))
 }
 
+// Bytes returns the [32]byte representation of the *Hash
+func (h *Hash) Bytes() [32]byte {
+	bi := new(big.Int).SetBytes(common.SwapEndianness(h[:])).Bytes()
+	b := [32]byte{}
+	copy(b[:], bi[:])
+	return b
+}
+
+// NewBigIntFromBytes returns a *big.Int from a byte array, swapping the endianness in the process. This is the intended method to get a *big.Int from a byte array that previously has ben generated by the Hash.Bytes() method.
+func NewBigIntFromBytes(b []byte) (*big.Int, error) {
+	if len(b) != 32 {
+		return nil, fmt.Errorf("Expected 32 bytes, found %d bytes", len(b))
+	}
+	return new(big.Int).SetBytes(common.SwapEndianness(b[:32])), nil
+}
+
 // NewHashFromBigInt returns a *Hash representation of the given *big.Int
 func NewHashFromBigInt(b *big.Int) *Hash {
 	r := &Hash{}
@@ -279,8 +295,16 @@ func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) {
 	return k, nil
 }
 
-// Delete removes the specified Key from the MerkleTree, and updates the pad from the delted key to the Root with the new values.
-// This method removes the key from the MerkleTree, but does not remove the old nodes from the key-value database; this means that if the tree is accessed by an old Root where the key was not deleted yet, the key will still exist. If is desired to remove the key-values from the database that are not under the current Root, an option could be to dump all the claims and import them in a new MerkleTree in a new database, but this will loose all the Root history of the MerkleTree
+// Delete removes the specified Key from the MerkleTree and updates the path
+// from the deleted key to the Root with the new values.  This method removes
+// the key from the MerkleTree, but does not remove the old nodes from the
+// key-value database; this means that if the tree is accessed by an old Root
+// where the key was not deleted yet, the key will still exist. If is desired
+// to remove the key-values from the database that are not under the current
+// Root, an option could be to dump all the leafs (using mt.DumpLeafs) and
+// import them in a new MerkleTree in a new database (using
+// mt.ImportDumpedLeafs), but this will loose all the Root history of the
+// MerkleTree
 func (mt *MerkleTree) Delete(k *big.Int) error {
 	// verify that the MerkleTree is writable
 	if !mt.writable {
@@ -630,7 +654,7 @@ func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
 }
 
 // RootFromProof calculates the root that would correspond to a tree whose
-// siblings are the ones in the proof with the claim hashing to hIndex and
+// siblings are the ones in the proof with the leaf hashing to hIndex and
 // hValue.
 func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
 	kHash := NewHashFromBigInt(k)
@@ -708,8 +732,8 @@ func (mt *MerkleTree) walk(key *Hash, f func(*Node)) error {
 // Walk iterates over all the branches of a MerkleTree with the given rootKey
 // if rootKey is nil, it will get the current RootKey of the current state of the MerkleTree.
 // For each node, it calls the f function given in the parameters.
-// See some examples of the Walk function usage in the merkletree_test.go
-// test functions: TestMTWalk, TestMTWalkGraphViz, TestMTWalkDumpClaims
+// See some examples of the Walk function usage in the merkletree.go and
+// merkletree_test.go
 func (mt *MerkleTree) Walk(rootKey *Hash, f func(*Node)) error {
 	if rootKey == nil {
 		rootKey = mt.Root()
@@ -773,3 +797,38 @@ func (mt *MerkleTree) PrintGraphViz(rootKey *Hash) error {
 	fmt.Println(w)
 	return nil
 }
+
+// DumpLeafs returns all the Leafs that exist under the given Root. If no Root
+// is given (nil), it uses the current Root of the MerkleTree.
+func (mt *MerkleTree) DumpLeafs(rootKey *Hash) ([]byte, error) {
+	var b []byte
+	err := mt.Walk(rootKey, func(n *Node) {
+		if n.Type == NodeTypeLeaf {
+			l := n.Entry[0].Bytes()
+			r := n.Entry[1].Bytes()
+			b = append(b, append(l[:], r[:]...)...)
+		}
+	})
+	return b, err
+}
+
+// ImportDumpedLeafs parses and adds to the MerkleTree the dumped list of leafs
+// from the DumpLeafs function.
+func (mt *MerkleTree) ImportDumpedLeafs(b []byte) error {
+	for i := 0; i < len(b); i += 64 {
+		lr := b[i : i+64]
+		lB, err := NewBigIntFromBytes(lr[:32])
+		if err != nil {
+			return err
+		}
+		rB, err := NewBigIntFromBytes(lr[32:])
+		if err != nil {
+			return err
+		}
+		err = mt.Add(lB, rB)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

merkletree_test.go

@@ -449,3 +449,27 @@ func TestDeleteNonExistingKeys(t *testing.T) {
 	err = mt.Delete(big.NewInt(33))
 	assert.Equal(t, ErrKeyNotFound, err)
 }
+
+func TestDumpLeafsImportLeafs(t *testing.T) {
+	mt, err := NewMerkleTree(db.NewMemoryStorage(), 140)
+	require.Nil(t, err)
+	defer mt.db.Close()
+
+	for i := 0; i < 10; i++ {
+		k := big.NewInt(int64(i))
+		v := big.NewInt(0)
+		err = mt.Add(k, v)
+		require.Nil(t, err)
+	}
+
+	d, err := mt.DumpLeafs(nil)
+	assert.Nil(t, err)
+
+	mt2, err := NewMerkleTree(db.NewMemoryStorage(), 140)
+	require.Nil(t, err)
+	defer mt2.db.Close()
+	err = mt2.ImportDumpedLeafs(d)
+	assert.Nil(t, err)
+
+	assert.Equal(t, mt.Root(), mt2.Root())
+}

node.go

@@ -11,7 +11,7 @@ type NodeType byte
 const (
 	// NodeTypeMiddle indicates the type of middle Node that has children.
 	NodeTypeMiddle NodeType = 0
-	// NodeTypeLeaf indicates the type of a leaf Node that contains a claim.
+	// NodeTypeLeaf indicates the type of a leaf Node that contains a key & value.
 	NodeTypeLeaf NodeType = 1
 	// NodeTypeEmpty indicates the type of an empty Node.
 	NodeTypeEmpty NodeType = 2

utils.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"math/big"
 
-	cryptoConstants "github.com/iden3/go-iden3-crypto/constants"
 	"github.com/iden3/go-iden3-crypto/poseidon"
 )
 
@@ -24,9 +23,6 @@ func HashElems(elems ...*big.Int) (*Hash, error) {
 
 	poseidonHash, err := poseidon.PoseidonHash(bi)
 	if err != nil {
-		fmt.Println("ERR HashElems PoseidonHash")
-		fmt.Println("e", bi[0])
-		fmt.Println("q", cryptoConstants.Q)
 		return nil, err
 	}
 	return NewHashFromBigInt(poseidonHash), nil