Merge pull request #1403 from nspcc-dev/merkle-memory-optimization

hash: introduce memory-optimized merkle root hash calculation routine

pkg/consensus/consensus.go

@@ -562,12 +562,7 @@ func (s *service) newBlockFromContext(ctx *dbft.Context) block.Block {
 	hashes := make([]util.Uint256, len(ctx.TransactionHashes)+1)
 	hashes[0] = block.Block.ConsensusData.Hash()
 	copy(hashes[1:], ctx.TransactionHashes)
-	mt, err := hash.NewMerkleTree(hashes)
-	if err != nil {
-		s.log.Fatal("can't calculate merkle root for the new block")
-		return nil
-	}
-	block.Block.MerkleRoot = mt.Root()
+	block.Block.MerkleRoot = hash.CalcMerkleRoot(hashes)
 
 	return block
 }

pkg/core/block/block.go

@@ -47,25 +47,19 @@ func (b *Block) Header() *Header {
 }
 
 // computeMerkleTree computes Merkle tree based on actual block's data.
-func (b *Block) computeMerkleTree() (*hash.MerkleTree, error) {
+func (b *Block) computeMerkleTree() util.Uint256 {
 	hashes := make([]util.Uint256, len(b.Transactions)+1)
 	hashes[0] = b.ConsensusData.Hash()
 	for i, tx := range b.Transactions {
 		hashes[i+1] = tx.Hash()
 	}
 
-	return hash.NewMerkleTree(hashes)
+	return hash.CalcMerkleRoot(hashes)
 }
 
 // RebuildMerkleRoot rebuilds the merkleroot of the block.
-func (b *Block) RebuildMerkleRoot() error {
-	merkle, err := b.computeMerkleTree()
-	if err != nil {
-		return err
-	}
-
-	b.MerkleRoot = merkle.Root()
-	return nil
+func (b *Block) RebuildMerkleRoot() {
+	b.MerkleRoot = b.computeMerkleTree()
 }
 
 // Verify verifies the integrity of the block.
@@ -81,11 +75,8 @@ func (b *Block) Verify() error {
 		}
 	}
 
-	merkle, err := b.computeMerkleTree()
-	if err != nil {
-		return err
-	}
-	if !b.MerkleRoot.Equals(merkle.Root()) {
+	merkle := b.computeMerkleTree()
+	if !b.MerkleRoot.Equals(merkle) {
 		return errors.New("MerkleRoot mismatch")
 	}
 	return nil

pkg/core/block/block_test.go

@@ -111,11 +111,11 @@ func TestHashBlockEqualsHashHeader(t *testing.T) {
 func TestBlockVerify(t *testing.T) {
 	block := newDumbBlock()
 	assert.NotNil(t, block.Verify())
-	assert.Nil(t, block.RebuildMerkleRoot())
+	block.RebuildMerkleRoot()
 	assert.Nil(t, block.Verify())
 
 	block.Transactions = []*transaction.Transaction{}
-	assert.NoError(t, block.RebuildMerkleRoot())
+	block.RebuildMerkleRoot()
 	assert.Nil(t, block.Verify())
 }
 

pkg/core/helper_test.go

@@ -72,11 +72,7 @@ func newBlock(cfg config.ProtocolConfiguration, index uint32, prev util.Uint256,
 		},
 		Transactions: txs,
 	}
-	err := b.RebuildMerkleRoot()
-	if err != nil {
-		panic(err)
-	}
-
+	b.RebuildMerkleRoot()
 	b.Script.InvocationScript = testchain.Sign(b.GetSignedPart())
 	return b
 }

pkg/core/util.go

@@ -66,10 +66,7 @@ func createGenesisBlock(cfg config.ProtocolConfiguration) (*block.Block, error)
 			Nonce:        2083236893,
 		},
 	}
-
-	if err = b.RebuildMerkleRoot(); err != nil {
-		return nil, err
-	}
+	b.RebuildMerkleRoot()
 
 	return b, nil
 }

pkg/crypto/hash/merkle_bench_test.go

@@ -0,0 +1,37 @@
+package hash
+
+import (
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/stretchr/testify/require"
+)
+
+func BenchmarkMerkle(t *testing.B) {
+	var err error
+	var hashes = make([]util.Uint256, 100000)
+	var h = make([]byte, 32)
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	for i := range hashes {
+		r.Read(h)
+		hashes[i], err = util.Uint256DecodeBytesBE(h)
+		require.NoError(t, err)
+	}
+
+	t.Run("NewMerkleTree", func(t *testing.B) {
+		t.ResetTimer()
+		for n := 0; n < t.N; n++ {
+			tr, err := NewMerkleTree(hashes)
+			require.NoError(t, err)
+			_ = tr.Root()
+		}
+	})
+	t.Run("CalcMerkleRoot", func(t *testing.B) {
+		t.ResetTimer()
+		for n := 0; n < t.N; n++ {
+			_ = CalcMerkleRoot(hashes)
+		}
+	})
+}

pkg/crypto/hash/merkle_tree.go

@@ -66,6 +66,38 @@ func buildMerkleTree(leaves []*MerkleTreeNode) *MerkleTreeNode {
 	return buildMerkleTree(parents)
 }
 
+// CalcMerkleRoot calculcates Merkle root hash value for a given slice of hashes.
+// It doesn't create a full MerkleTree structure and it uses given slice as a
+// scratchpad, so it will destroy its contents in the process. But it's much more
+// memory efficient if you only need root hash value, while NewMerkleTree would
+// make 3*N allocations for N hashes, this function will only make 4. It also is
+// an error to call this function for zero-length hashes slice, the function will
+// panic.
+func CalcMerkleRoot(hashes []util.Uint256) util.Uint256 {
+	if len(hashes) == 0 {
+		panic("length of the hashes cannot be zero")
+	}
+	if len(hashes) == 1 {
+		return hashes[0]
+	}
+
+	scratch := make([]byte, 64)
+	parents := hashes[:(len(hashes)+1)/2]
+	for i := 0; i < len(parents); i++ {
+		copy(scratch, hashes[i*2].BytesBE())
+
+		if i*2+1 == len(hashes) {
+			copy(scratch[32:], hashes[i*2].BytesBE())
+		} else {
+			copy(scratch[32:], hashes[i*2+1].BytesBE())
+		}
+
+		parents[i] = DoubleSha256(scratch)
+	}
+
+	return CalcMerkleRoot(parents)
+}
+
 // MerkleTreeNode represents a node in the MerkleTree.
 type MerkleTreeNode struct {
 	hash       util.Uint256

pkg/crypto/hash/merkle_tree_test.go

@@ -18,6 +18,8 @@ func testComputeMerkleTree(t *testing.T, hexHashes []string, result string) {
 
 	merkle, err := NewMerkleTree(hashes)
 	require.NoError(t, err)
+	optimized := CalcMerkleRoot(hashes)
+	assert.Equal(t, result, optimized.StringLE())
 	assert.Equal(t, result, merkle.Root().StringLE())
 	assert.Equal(t, true, merkle.root.IsRoot())
 	assert.Equal(t, false, merkle.root.IsLeaf())

pkg/internal/testchain/address.go

@@ -177,7 +177,7 @@ func NewBlock(t *testing.T, bc blockchainer.Blockchainer, offset uint32, primary
 		},
 		Transactions: txs,
 	}
-	_ = b.RebuildMerkleRoot()
+	b.RebuildMerkleRoot()
 
 	b.Script.InvocationScript = Sign(b.GetSignedPart())
 	return b