Merge pull request #5 from dk-open/perf/mmr-initialisation-improvements

Multiple performance improvements

Author: dk-open

README.md

@@ -76,7 +76,6 @@ go get -u github.com/discretemind/mmr
 ```go
 package main
 
-
 import (
 	"context"
 	"fmt"
@@ -91,24 +90,30 @@ func main() {
 	ctx := context.Background()
 	memoryIndexes := store.MemoryIndexSource[uint64, types.Hash256]()
 	m := merkle.NewMountainRange[uint64, types.Hash256](hasher.Sha3_256, memoryIndexes)
-
+	var transactions []types.Hash256
 	for i := 0; i < 10; i++ {
 		h := hasher.Sha3_256([]byte(fmt.Sprintf("test data %d", i)))
+		transactions = append(transactions, h)
 		fmt.Printf("Adding at %d item %x\n", i, h)
-		if err := m.Add(ctx, h); err != nil {
-			log.Fatal(err)
-		}
+
+	}
+	if err := m.Add(ctx, transactions...); err != nil {
+		log.Fatal(err)
 	}
 	fmt.Println()
-	root := m.Root()
+	root, err := m.Root(ctx)
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	fmt.Printf("Root: %x\n", root)
 	item3, err := m.Get(ctx, 3)
 	if err != nil {
 		log.Fatal(err)
 	}
 	fmt.Printf("Item at index 3: %x\n", item3)
 	fmt.Println("Create a proof for item 4")
-	prooft, err := m.CreateProof(ctx, 4)
+	prooft, err := m.ProofByIndex(ctx, 4)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -118,7 +123,7 @@ func main() {
 		log.Fatal("Proof is not invalid")
 	}
 
-	fmt.Println("Proof is valid")
+	fmt.Println("Proof is Valid")
 }
 
 ```

merkle/index/index.leaf.go

@@ -49,7 +49,6 @@ func (l *leafIndex[TI]) GetSibling() Index[TI] {
 
 // RightUp moves the current node to its parent if it's a right child in the tree hierarchy.
 func (l *leafIndex[TI]) RightUp() Index[TI] {
-	//	value := x.Index()
 	if l.IsRight() {
 		return NodeIndex(l.value)
 	}

merkle/index/index.node_test.go

@@ -0,0 +1,106 @@
+package index_test
+
+import (
+	"github.com/dk-open/go-mmr/merkle/index"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func getNodeIndex(value uint32) index.Index[uint32] {
+	return index.NodeIndex[uint32](value)
+}
+
+func TestNodeIndex_Creation(t *testing.T) {
+	// Test creating a node index
+	value := uint32(4)
+	node := getNodeIndex(value)
+	assert.Equal(t, value, node.Index(), "Node index should be set correctly")
+	assert.Equal(t, 2, node.GetHeight(), "Height should be calculated correctly")
+}
+
+func TestNodeIndex_LeftBranch(t *testing.T) {
+	// Test left branch calculation
+	node := getNodeIndex(uint32(6))
+	leftBranch := node.LeftBranch()
+	assert.NotNil(t, leftBranch, "Left branch should exist")
+	assert.Equal(t, uint32(2), leftBranch.Index(), "Left branch should be calculated correctly")
+}
+
+func TestNodeIndex_GetSibling(t *testing.T) {
+	// Test sibling calculation for both left and right nodes
+	node := getNodeIndex(uint32(3)) // Right sibling
+	sibling := node.GetSibling()
+	assert.Equal(t, uint32(1), sibling.Index(), "Right sibling should be calculated correctly")
+
+	node = getNodeIndex(uint32(2)) // Left sibling
+	sibling = node.GetSibling()
+	assert.Equal(t, uint32(6), sibling.Index(), "Left sibling should be calculated correctly")
+}
+
+func TestNodeIndex_RightUp(t *testing.T) {
+	// Test moving to the parent node from a right child
+	node := getNodeIndex(uint32(6)) // Right child of parent
+	parent := node.RightUp()
+	assert.NotNil(t, parent, "Parent should exist")
+	assert.Equal(t, uint32(4), parent.Index(), "Parent node should be calculated correctly")
+
+	// Test when already at the top
+	node = getNodeIndex(uint32(4)) // Top node
+	parent = node.RightUp()
+	assert.Nil(t, parent, "There should be no parent when already at the top")
+}
+
+func TestNodeIndex_Up(t *testing.T) {
+	// Test moving up in the hierarchy
+	node := getNodeIndex(uint32(3)) // Right child
+	parent := node.Up()
+	assert.Equal(t, uint32(2), parent.Index(), "Parent node should be calculated correctly")
+
+	node = getNodeIndex(uint32(2)) // Left child
+	parent = node.Up()
+	assert.Equal(t, uint32(4), parent.Index(), "Parent node should be calculated correctly")
+}
+
+func TestNodeIndex_IsRight(t *testing.T) {
+	// Test checking if a node is a right child
+	node := getNodeIndex(uint32(3)) // Right child
+	assert.True(t, node.IsRight(), "Node should be identified as a right child")
+
+	node = getNodeIndex(uint32(2)) // Left child
+	assert.False(t, node.IsRight(), "Node should be identified as a left child")
+}
+
+func TestNodeIndex_Top(t *testing.T) {
+	// Test finding the top ancestor
+	node := getNodeIndex(uint32(6)) // Start at node 6
+	top := node.Top()
+	assert.Equal(t, uint32(4), top.Index(), "Top ancestor should be calculated correctly")
+}
+
+func TestNodeIndex_Children(t *testing.T) {
+	// Test retrieving the children of a node
+	node := getNodeIndex(uint32(6)) // Node with height > 0
+	children := node.Children()
+	assert.Len(t, children, 2, "Node should have two children")
+	assert.Equal(t, uint32(5), children[0].Index(), "Left child should be calculated correctly")
+	assert.Equal(t, uint32(7), children[1].Index(), "Right child should be calculated correctly")
+
+	node = getNodeIndex(uint32(1)) // Node with height 0
+	children = node.Children()
+	assert.Len(t, children, 2, "Leaf node should have two children")
+	assert.Equal(t, uint32(0), children[0].Index(), "Left child should be calculated correctly for leaf")
+	assert.Equal(t, uint32(1), children[1].Index(), "Right child should be calculated correctly for leaf")
+}
+
+func TestNodeIndex_IsLeaf(t *testing.T) {
+	// Test checking if a node is a leaf
+	node := getNodeIndex(uint32(1)) // Node at height 0
+	assert.False(t, node.IsLeaf(), "Node should not be identified as a leaf")
+}
+
+func TestNodeIndex_Key(t *testing.T) {
+	// Test generating the key for the node
+	node := getNodeIndex(uint32(5))
+	key := node.Key()
+	assert.Equal(t, "node_5", key, "Node key should be generated correctly")
+}

merkle/index/index_test.go

@@ -98,5 +98,4 @@ func TestPeaks(t *testing.T) {
 	assert.Equal(t, 13, peaks14[1].Index())
 	assert.Equal(t, 10, peaks14[2].Index())
 	assert.Equal(t, 4, peaks14[3].Index())
-
 }

merkle/index/node.leaf_test.go

@@ -0,0 +1,113 @@
+package index_test
+
+import (
+	"github.com/dk-open/go-mmr/merkle/index"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func getLeafIndex(value uint32) index.Index[uint32] {
+	return index.LeafIndex[uint32](value)
+}
+
+func TestLeafIndex_Creation(t *testing.T) {
+	// Test creating a leaf index
+	value := uint32(5)
+	leaf := getLeafIndex(value)
+	assert.Equal(t, value, leaf.Index(), "Leaf index should be set correctly")
+}
+
+func TestLeafIndex_IsLeaf(t *testing.T) {
+	// Test if the node is correctly identified as a leaf
+	leaf := getLeafIndex(uint32(3))
+	assert.True(t, leaf.IsLeaf(), "Leaf index should return true for IsLeaf")
+}
+
+func TestLeafIndex_GetHeight(t *testing.T) {
+	// Test that the height of a leaf node is always 0
+	leaf := getLeafIndex(uint32(3))
+	assert.Equal(t, 0, leaf.GetHeight(), "Leaf height should always be 0")
+}
+
+func TestLeafIndex_LeftBranch(t *testing.T) {
+	// Test left branch calculation
+	leaf := getLeafIndex(uint32(2)) // Left child
+	leftBranch := leaf.LeftBranch()
+	assert.NotNil(t, leftBranch, "Left branch should exist for non-zero, non-right node")
+	assert.Equal(t, uint32(1), leftBranch.Index(), "Left branch should be correctly calculated")
+
+	// Test no left branch for node 0
+	leaf = getLeafIndex(uint32(0)) // Node 0 has no left branch
+	leftBranch = leaf.LeftBranch()
+	assert.Nil(t, leftBranch, "Left branch should not exist for node 0")
+}
+
+func TestLeafIndex_GetSibling(t *testing.T) {
+	// Test sibling calculation for both left and right nodes
+	leaf := getLeafIndex(uint32(3)) // Right sibling
+	sibling := leaf.GetSibling()
+	assert.Equal(t, uint32(2), sibling.Index(), "Right sibling should be calculated correctly")
+
+	leaf = getLeafIndex(uint32(2)) // Left sibling
+	sibling = leaf.GetSibling()
+	assert.Equal(t, uint32(3), sibling.Index(), "Left sibling should be calculated correctly")
+}
+
+func TestLeafIndex_RightUp(t *testing.T) {
+	// Test moving to the parent node from a right child
+	leaf := getLeafIndex(uint32(3)) // Right child of parent
+	parent := leaf.RightUp()
+	assert.NotNil(t, parent, "Parent should exist")
+	assert.Equal(t, uint32(3), parent.Index(), "Parent node should be calculated correctly")
+
+	// Test when there is no parent
+	leaf = getLeafIndex(uint32(0)) // Node 0 has no parent
+	parent = leaf.RightUp()
+	assert.Nil(t, parent, "There should be no parent for node 0")
+}
+
+func TestLeafIndex_Up(t *testing.T) {
+	// Test moving up in the hierarchy
+	leaf := getLeafIndex(uint32(3)) // Right child
+	parent := leaf.Up()
+	assert.Equal(t, uint32(3), parent.Index(), "Parent node should be calculated correctly")
+
+	leaf = getLeafIndex(uint32(2)) // Left child
+	parent = leaf.Up()
+	assert.Equal(t, uint32(3), parent.Index(), "Parent node should be calculated correctly")
+}
+
+func TestLeafIndex_IsRight(t *testing.T) {
+	// Test checking if a leaf is a right child
+	leaf := getLeafIndex(uint32(3)) // Right child
+	assert.True(t, leaf.IsRight(), "Leaf should be identified as a right child")
+
+	leaf = getLeafIndex(uint32(2)) // Left child
+	assert.False(t, leaf.IsRight(), "Leaf should be identified as a left child")
+}
+
+func TestLeafIndex_Top(t *testing.T) {
+	// Test finding the top ancestor
+	leaf := getLeafIndex(uint32(3)) // Start at leaf 3
+	top := leaf.Top()
+	assert.Equal(t, uint32(2), top.Index(), "Top ancestor should be calculated correctly")
+
+	leaf = getLeafIndex(uint32(2)) // Start at leaf 2
+	top = leaf.Top()
+	assert.Equal(t, uint32(2), top.Index(), "Top ancestor should return the leaf itself when it is the top")
+}
+
+func TestLeafIndex_Children(t *testing.T) {
+	// Test that leaf nodes have no children
+	leaf := getLeafIndex(uint32(3))
+	children := leaf.Children()
+	assert.Nil(t, children, "Leaf nodes should have no children")
+}
+
+func TestLeafIndex_Key(t *testing.T) {
+	// Test generating the key for the leaf
+	leaf := getLeafIndex(uint32(5))
+	key := leaf.Key()
+	assert.Equal(t, "leaf_5", key, "Leaf key should be generated correctly")
+}

merkle/merkle.go

@@ -10,16 +10,17 @@ import (
 )
 
 type IMountainRange[TIndex index.Value, THash types.HashType] interface {
-	Add(ctx context.Context, value THash) error
+	Add(ctx context.Context, values ...THash) error
 	Get(ctx context.Context, index TIndex) (THash, error)
-	CreateProof(ctx context.Context, index TIndex) (*Proof[TIndex, THash], error)
-	Root() IRoot[TIndex, THash]
+	ProofByIndex(ctx context.Context, index TIndex) (*Proof[TIndex, THash], error)
+	Proof(ctx context.Context, item THash) (*Proof[TIndex, THash], error)
+	Root(ctx context.Context) (IRoot[TIndex, THash], error)
 	Size() TIndex
 }
 
 type mmr[TIndex index.Value, THash types.HashType] struct {
 	sync.RWMutex
-	root    THash
+	//root    THash
 	size    TIndex
 	hf      types.Hasher[THash]
 	indexes store.IIndexSource[TIndex, THash]
@@ -40,11 +41,15 @@ func (m *mmr[TIndex, THash]) Get(ctx context.Context, index TIndex) (res THash,
 	return
 }
 
-func (m *mmr[TIndex, THash]) Add(ctx context.Context, value THash) error {
+func (m *mmr[TIndex, THash]) Add(ctx context.Context, value ...THash) error {
 	m.Lock()
-	err := m.appendMerkle(ctx, value)
-	m.Unlock()
-	return err
+	defer m.Unlock()
+	for _, v := range value {
+		if err := m.appendMerkle(ctx, v); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 // getProofIndexes collects the indexes needed to create a proof for the given item.
@@ -65,7 +70,15 @@ func (m *mmr[TIndex, THash]) getProofIndexes(item index.Index[TIndex], maxIndex
 	return res
 }
 
-func (m *mmr[TIndex, THash]) CreateProof(ctx context.Context, i TIndex) (*Proof[TIndex, THash], error) {
+func (m *mmr[TIndex, THash]) Proof(ctx context.Context, item THash) (*Proof[TIndex, THash], error) {
+	leafIndex, err := m.indexes.LeafIndex(ctx, item)
+	if err != nil {
+		return nil, err
+	}
+	return m.ProofByIndex(ctx, leafIndex)
+}
+
+func (m *mmr[TIndex, THash]) ProofByIndex(ctx context.Context, i TIndex) (*Proof[TIndex, THash], error) {
 	m.RLock()
 	defer m.RUnlock()
 
@@ -132,6 +145,24 @@ func (m *mmr[TIndex, THash]) Size() TIndex {
 	return m.size
 }
 
-func (m *mmr[TIndex, THash]) Root() IRoot[TIndex, THash] {
-	return newRoot[TIndex, THash](m.root, m.hf)
+func (m *mmr[TIndex, THash]) Root(ctx context.Context) (IRoot[TIndex, THash], error) {
+
+	m.RLock()
+	defer m.RUnlock()
+	peaks := index.GetPeaks(index.LeafIndex(m.size - 1))
+	hashes := make([][]byte, len(peaks))
+	for i, p := range peaks {
+		h, hErr := m.indexes.Get(ctx, p.IsLeaf(), p.Index())
+		if hErr != nil {
+			return nil, hErr
+		}
+		data, hErr := types.HashBytes[THash](h)
+		if hErr != nil {
+			return nil, hErr
+		}
+		hashes[i] = data
+	}
+	r := m.hf(hashes...)
+
+	return newRoot[TIndex, THash](r, m.hf), nil
 }

merkle/merkle.internal.go

@@ -63,20 +63,6 @@ func (m *mmr[TIndex, THash]) appendMerkle(ctx context.Context, value THash) (err
 		return err
 	}
 
-	peaks := index.GetPeaks(leafIndex)
-	hashes := make([][]byte, len(peaks))
-	for i, p := range peaks {
-		h, hErr := m.indexes.Get(ctx, p.IsLeaf(), p.Index())
-		if hErr != nil {
-			return hErr
-		}
-		data, hErr := types.HashBytes[THash](h)
-		if hErr != nil {
-			return hErr
-		}
-		hashes[i] = data
-	}
-	m.root = m.hf(hashes...)
 	m.size = m.size + 1
 	return nil
 }