Caplin: Introduced cached merkle tree

cl/cltypes/solid/hash_list.go

@@ -19,7 +19,6 @@ package solid
 import (
 	"encoding/json"
 
-	"github.com/erigontech/erigon-lib/common"
 	libcommon "github.com/erigontech/erigon-lib/common"
 	"github.com/erigontech/erigon-lib/common/length"
 	"github.com/erigontech/erigon-lib/types/clonable"
@@ -32,7 +31,7 @@ type hashList struct {
 	u    []byte
 	l, c int
 
-	hashBuf
+	*merkle_tree.MerkleTree
 }
 
 func NewHashList(c int) HashListSSZ {
@@ -68,6 +67,9 @@ func (arr *hashList) UnmarshalJSON(buf []byte) error {
 
 func (h *hashList) Append(val libcommon.Hash) {
 	offset := h.l * length.Hash
+	if h.MerkleTree != nil {
+		h.MerkleTree.AppendLeaf()
+	}
 	if offset == len(h.u) {
 		h.u = append(h.u, val[:]...)
 		h.l++
@@ -87,6 +89,7 @@ func (h *hashList) Length() int {
 
 func (h *hashList) Clear() {
 	h.l = 0
+	h.MerkleTree = nil
 }
 
 func (h *hashList) Clone() clonable.Clonable {
@@ -100,6 +103,12 @@ func (h *hashList) CopyTo(t IterableSSZ[libcommon.Hash]) {
 	if len(h.u) > len(tu.u) {
 		tu.u = make([]byte, len(h.u))
 	}
+	if h.MerkleTree != nil {
+		if tu.MerkleTree == nil {
+			tu.MerkleTree = &merkle_tree.MerkleTree{}
+		}
+		h.MerkleTree.CopyInto(tu.MerkleTree)
+	}
 	copy(tu.u, h.u)
 }
 
@@ -111,6 +120,7 @@ func (h *hashList) DecodeSSZ(buf []byte, _ int) error {
 	if len(buf)%length.Hash > 0 {
 		return ssz.ErrBadDynamicLength
 	}
+	h.MerkleTree = nil
 	h.u = libcommon.Copy(buf)
 	h.l = len(h.u) / length.Hash
 	return nil
@@ -136,43 +146,30 @@ func (h *hashList) Set(index int, newValue libcommon.Hash) {
 	if index >= h.l {
 		panic("too big bruh")
 	}
+	if h.MerkleTree != nil {
+		h.MerkleTree.MarkLeafAsDirty(index)
+	}
 	copy(h.u[index*length.Hash:], newValue[:])
 }
 
 func (h *hashList) hashVectorSSZ() ([32]byte, error) {
-	depth := GetDepth(uint64(h.c))
-	offset := length.Hash * h.l
-	elements := common.Copy(h.u[:offset])
-	for i := uint8(0); i < depth; i++ {
-		// Sequential
-		if len(elements)%64 != 0 {
-			elements = append(elements, merkle_tree.ZeroHashes[i][:]...)
-		}
-		outputLen := len(elements) / 2
-		h.makeBuf(outputLen)
-		if err := merkle_tree.HashByteSlice(h.buf, elements); err != nil {
-			return [32]byte{}, err
-		}
-		elements = h.buf
+	if h.MerkleTree == nil {
+		cap := uint64(h.c)
+		h.MerkleTree = &merkle_tree.MerkleTree{}
+		h.MerkleTree.Initialize(h.l, merkle_tree.OptimalMaxTreeCacheDepth, func(idx int, out []byte) {
+			copy(out, h.u[idx*length.Hash:(idx+1)*length.Hash])
+		}, /*limit=*/ &cap)
 	}
-
-	return common.BytesToHash(elements[:32]), nil
+	return h.MerkleTree.ComputeRoot(), nil
 }
 
 func (h *hashList) HashSSZ() ([32]byte, error) {
-	depth := GetDepth(uint64(h.c))
-	baseRoot := [32]byte{}
-	var err error
-	if h.l == 0 {
-		copy(baseRoot[:], merkle_tree.ZeroHashes[depth][:])
-	} else {
-		baseRoot, err = h.hashVectorSSZ()
-		if err != nil {
-			return [32]byte{}, err
-		}
-	}
 	lengthRoot := merkle_tree.Uint64Root(uint64(h.l))
-	return utils.Sha256(baseRoot[:], lengthRoot[:]), nil
+	coreRoot, err := h.hashVectorSSZ()
+	if err != nil {
+		return [32]byte{}, err
+	}
+	return utils.Sha256(coreRoot[:], lengthRoot[:]), nil
 }
 
 func (h *hashList) Range(fn func(int, libcommon.Hash, int) bool) {

cl/cltypes/solid/uint64_raw_list.go

@@ -126,7 +126,7 @@ func (arr *RawUint64List) SetReusableHashBuffer(buf []byte) {
 }
 
 func (arr *RawUint64List) hashBufLength() int {
-	return (((len(arr.u) * 4) + 3) / 4) * length.Hash
+	return ((len(arr.u) + 3) / 4) * length.Hash
 }
 
 func (arr *RawUint64List) HashSSZ() ([32]byte, error) {

cl/cltypes/solid/uint64_vector.go

@@ -19,7 +19,6 @@ package solid
 import (
 	"encoding/json"
 
-	"github.com/erigontech/erigon-lib/common/length"
 	"github.com/erigontech/erigon-lib/types/clonable"
 )
 
@@ -29,10 +28,9 @@ type uint64VectorSSZ struct {
 
 func NewUint64VectorSSZ(size int) Uint64VectorSSZ {
 	o := &byteBasedUint64Slice{
-		c:               size,
-		l:               size,
-		u:               make([]byte, size*8),
-		treeCacheBuffer: make([]byte, getTreeCacheSize((size+3)/4, treeCacheDepthUint64Slice)*length.Hash),
+		c: size,
+		l: size,
+		u: make([]byte, size*8),
 	}
 	return &uint64VectorSSZ{
 		u: o,

cl/cltypes/solid/uint64slice_byte.go

@@ -17,20 +17,16 @@
 package solid
 
 import (
-	"bytes"
 	"encoding/binary"
 	"encoding/json"
 	"strconv"
 
-	"github.com/erigontech/erigon-lib/common"
 	"github.com/erigontech/erigon-lib/common/length"
 	"github.com/erigontech/erigon-lib/types/ssz"
 	"github.com/erigontech/erigon/cl/merkle_tree"
 	"github.com/erigontech/erigon/cl/utils"
 )
 
-const treeCacheDepthUint64Slice = 0
-
 func convertDepthToChunkSize(d int) int {
 	return (1 << d) // just power of 2
 }
@@ -46,16 +42,15 @@ func getTreeCacheSize(listLen int, cacheDepth int) int {
 // memory usage, especially when dealing with large slices.
 type byteBasedUint64Slice struct {
 	// The bytes that back the slice
-	u               []byte
-	treeCacheBuffer []byte
+	u []byte
 
 	// Length of the slice
 	l int
 
 	// Capacity of the slice
 	c int
 
-	hashBuf
+	*merkle_tree.MerkleTree
 }
 
 // NewUint64Slice creates a new instance of byteBasedUint64Slice with a specified capacity limit.
@@ -70,25 +65,27 @@ func NewUint64Slice(limit int) *byteBasedUint64Slice {
 func (arr *byteBasedUint64Slice) Clear() {
 	arr.l = 0
 	clear(arr.u)
-	clear(arr.treeCacheBuffer)
+	arr.MerkleTree = nil
 }
 
 // CopyTo copies the slice to a target slice.
 func (arr *byteBasedUint64Slice) CopyTo(target *byteBasedUint64Slice) {
 	target.Clear()
-
+	// TODO: implement CopyTo for MPT
 	target.c = arr.c
 	target.l = arr.l
 	if len(target.u) < len(arr.u) {
 		target.u = make([]byte, len(arr.u))
 	}
-	if len(target.treeCacheBuffer) < len(arr.treeCacheBuffer) {
-		target.treeCacheBuffer = make([]byte, len(arr.treeCacheBuffer))
+	if arr.MerkleTree != nil {
+		if target.MerkleTree == nil {
+			target.MerkleTree = &merkle_tree.MerkleTree{}
+		}
+		arr.MerkleTree.CopyInto(target.MerkleTree)
 	}
-	target.treeCacheBuffer = target.treeCacheBuffer[:len(arr.treeCacheBuffer)]
+
 	target.u = target.u[:len(arr.u)]
 	copy(target.u, arr.u)
-	copy(target.treeCacheBuffer, arr.treeCacheBuffer)
 }
 
 func (arr *byteBasedUint64Slice) MarshalJSON() ([]byte, error) {
@@ -110,6 +107,7 @@ func (arr *byteBasedUint64Slice) UnmarshalJSON(buf []byte) error {
 	for _, elem := range list {
 		arr.Append(elem)
 	}
+	arr.MerkleTree = nil
 	return nil
 }
 
@@ -129,27 +127,22 @@ func (arr *byteBasedUint64Slice) Pop() uint64 {
 	val := binary.LittleEndian.Uint64(arr.u[offset : offset+8])
 	binary.LittleEndian.PutUint64(arr.u[offset:offset+8], 0)
 	arr.l = arr.l - 1
-	arr.treeCacheBuffer = arr.treeCacheBuffer[:getTreeCacheSize((arr.l+3)/4, treeCacheDepthUint64Slice)*length.Hash]
+	arr.MerkleTree = nil
 	return val
 }
 
 // Append adds a new element to the end of the slice.
 func (arr *byteBasedUint64Slice) Append(v uint64) {
 	if len(arr.u) <= arr.l*8 {
 		arr.u = append(arr.u, make([]byte, 32)...)
+		if arr.MerkleTree != nil {
+			arr.MerkleTree.AppendLeaf()
+		}
 	}
-
 	offset := arr.l * 8
 	binary.LittleEndian.PutUint64(arr.u[offset:offset+8], v)
-	arr.l = arr.l + 1
-	treeBufferExpectCache := getTreeCacheSize((arr.l+3)/4, treeCacheDepthUint64Slice) * length.Hash
-	if len(arr.treeCacheBuffer) < treeBufferExpectCache {
-		arr.treeCacheBuffer = append(arr.treeCacheBuffer, make([]byte, treeBufferExpectCache-len(arr.treeCacheBuffer))...)
-	}
-	ihIdx := (((arr.l - 1) / 4) / convertDepthToChunkSize(treeCacheDepthUint64Slice)) * length.Hash
-	for i := ihIdx; i < ihIdx+length.Hash; i++ {
-		arr.treeCacheBuffer[i] = 0
-	}
+
+	arr.l++
 }
 
 // Get returns the element at the given index.
@@ -163,11 +156,10 @@ func (arr *byteBasedUint64Slice) Get(index int) uint64 {
 
 // Set replaces the element at the given index with a new value.
 func (arr *byteBasedUint64Slice) Set(index int, v uint64) {
-	offset := index * 8
-	ihIdx := ((index / 4) / convertDepthToChunkSize(treeCacheDepthUint64Slice)) * length.Hash
-	for i := ihIdx; i < ihIdx+length.Hash; i++ {
-		arr.treeCacheBuffer[i] = 0
+	if arr.MerkleTree != nil {
+		arr.MerkleTree.MarkLeafAsDirty(index / 4)
 	}
+	offset := index * 8
 	binary.LittleEndian.PutUint64(arr.u[offset:offset+8], v)
 }
 
@@ -183,66 +175,30 @@ func (arr *byteBasedUint64Slice) Cap() int {
 
 // HashListSSZ computes the SSZ hash of the slice as a list. It returns the hash and any error encountered.
 func (arr *byteBasedUint64Slice) HashListSSZ() ([32]byte, error) {
-	depth := GetDepth((uint64(arr.c)*8 + 31) / 32)
-	baseRoot := [32]byte{}
-	var err error
-	if arr.l == 0 {
-		copy(baseRoot[:], merkle_tree.ZeroHashes[depth][:])
-	} else {
-		baseRoot, err = arr.HashVectorSSZ()
-		if err != nil {
-			return [32]byte{}, err
-		}
+	if arr.MerkleTree == nil {
+		arr.MerkleTree = &merkle_tree.MerkleTree{}
+		cap := uint64((arr.c*8 + length.Hash - 1) / length.Hash)
+
+		arr.MerkleTree.Initialize((arr.l+3)/4, merkle_tree.OptimalMaxTreeCacheDepth, func(idx int, out []byte) {
+			copy(out, arr.u[idx*length.Hash:])
+		}, &cap)
 	}
+
+	coreRoot := arr.ComputeRoot()
 	lengthRoot := merkle_tree.Uint64Root(uint64(arr.l))
-	return utils.Sha256(baseRoot[:], lengthRoot[:]), nil
+	return utils.Sha256(coreRoot[:], lengthRoot[:]), nil
 }
 
 // HashVectorSSZ computes the SSZ hash of the slice as a vector. It returns the hash and any error encountered.
 func (arr *byteBasedUint64Slice) HashVectorSSZ() ([32]byte, error) {
-	chunkSize := convertDepthToChunkSize(treeCacheDepthUint64Slice) * length.Hash
-	depth := GetDepth((uint64(arr.c)*8 + length.Hash - 1) / length.Hash)
-	emptyHashBytes := make([]byte, length.Hash)
-
-	layerBuffer := make([]byte, chunkSize)
-	maxTo := length.Hash*((arr.l-1)/4) + length.Hash
-
-	offset := 0
-	for i := 0; i < maxTo; i += chunkSize {
-		offset = (i / chunkSize) * length.Hash
-		from := i
-		to := min(from+chunkSize, maxTo)
-
-		if !bytes.Equal(arr.treeCacheBuffer[offset:offset+length.Hash], emptyHashBytes) {
-			continue
-		}
-		layerBuffer = layerBuffer[:to-from]
-		copy(layerBuffer, arr.u[from:to])
-		if err := computeFlatRootsToBuffer(uint8(min(treeCacheDepthUint64Slice, uint64(depth))), layerBuffer, arr.treeCacheBuffer[offset:]); err != nil {
-			return [32]byte{}, err
-		}
-	}
-	if treeCacheDepthUint64Slice >= depth {
-		return common.BytesToHash(arr.treeCacheBuffer[:32]), nil
-	}
-
-	arr.makeBuf(offset + length.Hash)
-	copy(arr.buf, arr.treeCacheBuffer[:offset+length.Hash])
-	elements := arr.buf
-	for i := uint8(treeCacheDepthUint64Slice); i < depth; i++ {
-		layerLen := len(elements)
-		if layerLen%64 == 32 {
-			elements = append(elements, merkle_tree.ZeroHashes[i][:]...)
-		}
-		outputLen := len(elements) / 2
-		arr.makeBuf(outputLen)
-		if err := merkle_tree.HashByteSlice(arr.buf, elements); err != nil {
-			return [32]byte{}, err
-		}
-		elements = arr.buf
+	if arr.MerkleTree == nil {
+		arr.MerkleTree = &merkle_tree.MerkleTree{}
+		arr.MerkleTree.Initialize((arr.l+3)/4, merkle_tree.OptimalMaxTreeCacheDepth, func(idx int, out []byte) {
+			copy(out, arr.u[idx*length.Hash:])
+		}, nil)
 	}
 
-	return common.BytesToHash(elements[:32]), nil
+	return arr.ComputeRoot(), nil
 }
 
 // EncodeSSZ encodes the slice in SSZ format. It appends the encoded data to the provided buffer and returns the result.
@@ -259,7 +215,7 @@ func (arr *byteBasedUint64Slice) DecodeSSZ(buf []byte, _ int) error {
 	bufferLength := length.Hash*((arr.l-1)/4) + length.Hash
 	arr.u = make([]byte, bufferLength)
 	copy(arr.u, buf)
-	arr.treeCacheBuffer = make([]byte, getTreeCacheSize((arr.l+3)/4, treeCacheDepthUint64Slice)*length.Hash)
+	arr.MerkleTree = nil
 	return nil
 }
 

cl/cltypes/solid/uint64slice_byte_test.go

@@ -38,6 +38,7 @@ func TestUint64SliceBasic(t *testing.T) {
 
 	out, err := slice.HashListSSZ()
 	require.NoError(t, err)
+
 	require.EqualValues(t, common.HexToHash("eb8cec5eaec74a32e8b9b56cc42f7627cef722f81081ead786c97a4df1c8be5d"), out)
 
 }