Remove mmr batch and add more tests

pkg/mmr/mmr.go

@@ -23,9 +23,10 @@ var (
 
 type MMRElement []byte
 
-type MMRNode struct {
-	pos      uint64
-	elements []MMRElement
+type MMRStorage interface {
+	getElement(pos uint64) (*MMRElement, error)
+	append(pos uint64, elements []MMRElement) error
+	commit() error
 }
 
 type MMR struct {
@@ -193,7 +194,7 @@ func (mmr *MMR) bagPeaks(peaks []MMRElement) MMRElement {
 		rightPeak, peaks = peaks[len(peaks)-1], peaks[:len(peaks)-1]
 		leftPeak, peaks = peaks[len(peaks)-1], peaks[:len(peaks)-1]
 
-		mergedPeak := mmr.merge(leftPeak, rightPeak)
+		mergedPeak := mmr.merge(rightPeak, leftPeak)
 		peaks = append(peaks, mergedPeak)
 	}
 

pkg/mmr/utils.go → pkg/mmr/mmr_mem_storage.go

@@ -1,11 +1,6 @@
-// Copyright 2024 ChainSafe Systems (ON)
-// SPDX-License-Identifier: LGPL-3.0-only
-
 package mmr
 
 import (
-	"hash"
-
 	"github.com/tidwall/btree"
 )
 
@@ -37,8 +32,3 @@ func (s *MemStorage) commit() error {
 	// Do nothing since all changes are automatically committed
 	return nil
 }
-
-func NewInMemMMR(hasher hash.Hash) *MMR {
-	storage := NewMemStorage()
-	return NewMMR(0, storage, hasher)
-}

pkg/mmr/mmr_mem_storage_test.go

@@ -0,0 +1,35 @@
+package mmr
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetElement(t *testing.T) {
+	memStorage := NewMemStorage()
+	elements := make(map[uint64]MMRElement)
+
+	for i := uint64(1); i < 100; i++ {
+		value := MMRElement(fmt.Sprintf("value%d", i))
+		elements[i] = value
+		memStorage.append(i, []MMRElement{value})
+	}
+
+	// Check all elements are in the right position
+	for pos, expected := range elements {
+		element, err := memStorage.getElement(pos)
+		assert.NoError(t, err)
+		assert.NotNil(t, element)
+		assert.Equal(t, *element, expected)
+	}
+}
+
+func TestGetNotFoundElement(t *testing.T) {
+	memStorage := NewMemStorage()
+
+	element, err := memStorage.getElement(100)
+	assert.NoError(t, err)
+	assert.Nil(t, element)
+}

pkg/mmr/mmr_test.go

@@ -5,27 +5,35 @@ package mmr
 
 import (
 	"encoding/binary"
+	"hash"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/crypto/blake2b"
 )
 
-func hashNumber(number int) MMRElement {
+func newInMemMMR(hasher hash.Hash) *MMR {
+	return NewMMR(0, NewMemStorage(), hasher)
+}
+
+func hashNumber(number uint32) MMRElement {
 	hasher, _ := blake2b.New256(nil)
-	buf := make([]byte, 8)
-	binary.BigEndian.PutUint64(buf, uint64(number))
-	hasher.Write(buf)
-	return hasher.Sum(nil)
+	var numBytes [4]byte
+	binary.LittleEndian.PutUint32(numBytes[:], number)
+	hasher.Write(numBytes[:])
+
+	var hash [32]byte
+	hasher.Sum(hash[:0])
+	return hash[:]
 }
 
-func TestPush1Elements(t *testing.T) {
+func TestPushOneElement_RootShouldBeSameLeaf(t *testing.T) {
 	hasher, err := blake2b.New256(nil)
 	assert.NoError(t, err)
 
-	inMemMMR := NewInMemMMR(hasher)
+	inMemMMR := newInMemMMR(hasher)
 
-	leaf := hashNumber(42)
+	leaf := hashNumber(0)
 	_, err = inMemMMR.Push(leaf)
 	assert.NoError(t, err)
 
@@ -35,13 +43,14 @@ func TestPush1Elements(t *testing.T) {
 	assert.Equal(t, root, leaf)
 }
 
-func TestPush4Elements(t *testing.T) {
+// Compared with the same MMR using substrate's implementation
+func TestPushManyElementsGetRootOk(t *testing.T) {
 	hasher, err := blake2b.New256(nil)
 	assert.NoError(t, err)
 
-	inMemMMR := NewInMemMMR(hasher)
+	inMemMMR := newInMemMMR(hasher)
 
-	for i := 0; i < 2; i++ {
+	for i := uint32(0); i < 100; i++ {
 		leaf := hashNumber(i)
 		_, err := inMemMMR.Push(leaf)
 		assert.NoError(t, err)
@@ -50,9 +59,9 @@ func TestPush4Elements(t *testing.T) {
 	root, err := inMemMMR.Root()
 	assert.NoError(t, err)
 
-	assert.Equal(t, []byte(root), []byte{
-		0x2a, 0x44, 0xf7, 0xc, 0xa4, 0x6b, 0xee, 0x95, 0xa, 0x4b, 0xd3, 0x52,
-		0x8a, 0x3a, 0x3a, 0x10, 0xc4, 0x3d, 0x19, 0x51, 0x9c, 0xfe, 0x67, 0xc7,
-		0x93, 0x94, 0x3a, 0x12, 0xfc, 0x7, 0xf7, 0xe7,
-	})
+	assert.Equal(t, []byte{
+		0x5, 0x0, 0xd0, 0xeb, 0xdb, 0xca, 0xd3, 0x6a, 0x79, 0xd3, 0x32, 0x5d,
+		0xbd, 0x2a, 0x4b, 0x2b, 0x97, 0x30, 0x1d, 0x8e, 0x48, 0x2a, 0x9b, 0xe2,
+		0x2, 0x1, 0x6e, 0x9f, 0x1c, 0xaa, 0xe1, 0x3f,
+	}, []byte(root))
 }