Add more tests, update README.md

README.md

@@ -3,3 +3,35 @@
 MerkleTree compatible with version from [circomlib](https://github.com/iden3/circomlib).
 
 Adaptation of the merkletree from https://github.com/iden3/go-iden3-core/tree/v0.0.8
+
+## Usage
+More detailed examples can be found at the [tests](https://github.com/iden3/go-merkletree/blob/master/merkletree_test.go), and in the [documentation](https://godoc.org/github.com/iden3/go-merkletree).
+
+```go
+import (
+	"fmt"
+	"math/big"
+	"testing"
+
+	"github.com/iden3/go-iden3-core/db"
+	"github.com/stretchr/testify/assert"
+)
+
+[...]
+
+func TestExampleMerkleTree(t *testing.T) {
+	mt, err := NewMerkleTree(db.NewMemoryStorage(), 10)
+	assert.Nil(t, err)
+
+	key := big.NewInt(1)
+	value := big.NewInt(2)
+	err = mt.Add(key, value)
+	assert.Nil(t, err)
+	fmt.Println(mt.Root().String())
+
+	proof, err := mt.GenerateProof(key, nil)
+	assert.Nil(t, err)
+
+	assert.True(t, VerifyProof(mt.Root(), proof, key, value))
+}
+```

go.mod

@@ -1,4 +1,4 @@
-module go-merkletree
+module github.com/iden3/go-merkletree
 
 go 1.14
 

merkletree.go

@@ -106,6 +106,11 @@ func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) {
 	return &mt, nil
 }
 
+// DB returns the MerkleTree.DB()
+func (mt *MerkleTree) DB() db.Storage {
+	return mt.db
+}
+
 // Root returns the MerkleRoot
 func (mt *MerkleTree) Root() *Hash {
 	return mt.rootKey

merkletree_test.go

@@ -1,6 +1,8 @@
 package merkletree
 
 import (
+	"encoding/hex"
+	"fmt"
 	"math/big"
 	"testing"
 
@@ -9,6 +11,21 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+var debug = false
+
+type Fatalable interface {
+	Fatal(args ...interface{})
+}
+
+func newTestingMerkle(f Fatalable, numLevels int) *MerkleTree {
+	mt, err := NewMerkleTree(db.NewMemoryStorage(), numLevels)
+	if err != nil {
+		f.Fatal(err)
+		return nil
+	}
+	return mt
+}
+
 func TestNewTree(t *testing.T) {
 	mt, err := NewMerkleTree(db.NewMemoryStorage(), 10)
 	assert.Nil(t, err)
@@ -34,6 +51,77 @@ func TestNewTree(t *testing.T) {
 	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45)))
 }
 
+func TestAddDifferentOrder(t *testing.T) {
+	mt1 := newTestingMerkle(t, 140)
+	defer mt1.db.Close()
+	for i := 0; i < 16; i++ {
+		k := big.NewInt(int64(i))
+		v := big.NewInt(0)
+		if err := mt1.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	mt2 := newTestingMerkle(t, 140)
+	defer mt2.db.Close()
+	for i := 16 - 1; i >= 0; i-- {
+		k := big.NewInt(int64(i))
+		v := big.NewInt(0)
+		if err := mt2.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	assert.Equal(t, mt1.Root().Hex(), mt2.Root().Hex())
+	assert.Equal(t, "ee3b91f9df12d26c1430b9d6693d4d5062b95b40f3f0a0a34ae560d677b76709", mt1.Root().Hex())
+}
+
+func TestAddRepeatedIndex(t *testing.T) {
+	mt := newTestingMerkle(t, 140)
+	defer mt.db.Close()
+	k := big.NewInt(int64(3))
+	v := big.NewInt(int64(12))
+	if err := mt.Add(k, v); err != nil {
+		t.Fatal(err)
+	}
+	err := mt.Add(k, v)
+	assert.NotNil(t, err)
+	assert.Equal(t, err, ErrEntryIndexAlreadyExists)
+}
+
+func TestGenerateAndVerifyProof128(t *testing.T) {
+	mt, err := NewMerkleTree(db.NewMemoryStorage(), 140)
+	require.Nil(t, err)
+	defer mt.db.Close()
+
+	for i := 0; i < 128; i++ {
+		k := big.NewInt(int64(i))
+		v := big.NewInt(0)
+		if err := mt.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+	proof, err := mt.GenerateProof(big.NewInt(42), nil)
+	assert.Nil(t, err)
+	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(42), big.NewInt(0)))
+}
+
+func TestTreeLimit(t *testing.T) {
+	mt, err := NewMerkleTree(db.NewMemoryStorage(), 5)
+	require.Nil(t, err)
+	defer mt.db.Close()
+
+	for i := 0; i < 16; i++ {
+		err = mt.Add(big.NewInt(int64(i)), big.NewInt(int64(i)))
+		assert.Nil(t, err)
+	}
+
+	// here the tree is full, should not allow to add more data as reaches the maximum number of levels
+	err = mt.Add(big.NewInt(int64(16)), big.NewInt(int64(16)))
+	assert.NotNil(t, err)
+	assert.Equal(t, ErrReachedMaxLevel, err)
+}
+
 func TestSiblingsFromProof(t *testing.T) {
 	mt, err := NewMerkleTree(db.NewMemoryStorage(), 140)
 	require.Nil(t, err)
@@ -61,3 +149,83 @@ func TestSiblingsFromProof(t *testing.T) {
 	assert.Equal(t, "b6c905f21c9928efa19a2c4e55d88d5fd0af493032f5b84620eb872e48ff5d01", siblings[4].Hex())
 	assert.Equal(t, "69873a951f49bbaff71039d672ffe52d73605aed6b40c1ce7ab068ad86a44d1e", siblings[5].Hex())
 }
+
+func TestVerifyProofCases(t *testing.T) {
+	mt := newTestingMerkle(t, 140)
+	defer mt.DB().Close()
+
+	for i := 0; i < 8; i++ {
+		if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Existence proof
+
+	proof, err := mt.GenerateProof(big.NewInt(4), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, proof.Existence, true)
+	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(4), big.NewInt(0)))
+	assert.Equal(t, "000300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f", hex.EncodeToString(proof.Bytes()))
+
+	for i := 8; i < 32; i++ {
+		proof, err = mt.GenerateProof(big.NewInt(int64(i)), nil)
+		assert.Nil(t, err)
+		if debug {
+			fmt.Println(i, proof)
+		}
+	}
+	// Non-existence proof, empty aux
+	proof, err = mt.GenerateProof(big.NewInt(12), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, proof.Existence, false)
+	// assert.True(t, proof.nodeAux == nil)
+	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(12), big.NewInt(0)))
+	assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f04000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
+
+	// Non-existence proof, diff. node aux
+	proof, err = mt.GenerateProof(big.NewInt(10), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, proof.Existence, false)
+	assert.True(t, proof.NodeAux != nil)
+	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(10), big.NewInt(0)))
+	assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b1662163097fcf8f911b271df196e0a75667b8a4f3024ef39f87201ed2b7cda349ba202296b7aeba35dc19ab0d4f65e175536c9952a90b6de18c3205611c3cd4fb408f01602000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
+}
+
+func TestVerifyProofFalse(t *testing.T) {
+	mt := newTestingMerkle(t, 140)
+	defer mt.DB().Close()
+
+	for i := 0; i < 8; i++ {
+		if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Invalid existence proof (node used for verification doesn't
+	// correspond to node in the proof)
+	proof, err := mt.GenerateProof(big.NewInt(int64(4)), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, proof.Existence, true)
+	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(5)), big.NewInt(int64(5))))
+
+	// Invalid non-existence proof (Non-existence proof, diff. node aux)
+	proof, err = mt.GenerateProof(big.NewInt(int64(4)), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, proof.Existence, true)
+	// Now we change the proof from existence to non-existence, and add e's
+	// data as auxiliary node.
+	proof.Existence = false
+	proof.NodeAux = &NodeAux{Key: NewHashFromBigInt(big.NewInt(int64(4))), Value: NewHashFromBigInt(big.NewInt(4))}
+	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(4)), big.NewInt(0)))
+}