From d52e79668b0fd839c4bad5bb9c54b9e62760b9a5 Mon Sep 17 00:00:00 2001 From: Roman Khimov Date: Tue, 15 Sep 2020 18:38:15 +0300 Subject: [PATCH] hash: introduce memory-optimized merkle root hash calculation routine NewMerkleTree is a memory hog, we can do better than that: BenchmarkMerkle/NewMerkleTree-8 13 88434670 ns/op 20828207 B/op 300035 allocs/op BenchmarkMerkle/CalcMerkleRoot-8 15 69264150 ns/op 0 B/op 0 allocs/op --- pkg/consensus/consensus.go | 7 +----- pkg/core/block/block.go | 21 +++++----------- pkg/core/block/block_test.go | 4 +-- pkg/core/helper_test.go | 6 +---- pkg/core/util.go | 5 +--- pkg/crypto/hash/merkle_bench_test.go | 37 ++++++++++++++++++++++++++++ pkg/crypto/hash/merkle_tree.go | 32 ++++++++++++++++++++++++ pkg/crypto/hash/merkle_tree_test.go | 2 ++ pkg/internal/testchain/address.go | 2 +- 9 files changed, 83 insertions(+), 33 deletions(-) create mode 100644 pkg/crypto/hash/merkle_bench_test.go diff --git a/pkg/consensus/consensus.go b/pkg/consensus/consensus.go index 9b29cbde7..ca2afd04d 100644 --- a/pkg/consensus/consensus.go +++ b/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 } diff --git a/pkg/core/block/block.go b/pkg/core/block/block.go index 28cbcc12a..9c53ca4bd 100644 --- a/pkg/core/block/block.go +++ b/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 diff --git a/pkg/core/block/block_test.go b/pkg/core/block/block_test.go index 9795d465c..3e9c37317 100644 --- a/pkg/core/block/block_test.go +++ b/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()) } diff --git a/pkg/core/helper_test.go b/pkg/core/helper_test.go index bb0032d14..d327eb2d5 100644 --- a/pkg/core/helper_test.go +++ b/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 } diff --git a/pkg/core/util.go b/pkg/core/util.go index dae367fff..5b6cbcfa0 100644 --- a/pkg/core/util.go +++ 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 } diff --git a/pkg/crypto/hash/merkle_bench_test.go b/pkg/crypto/hash/merkle_bench_test.go new file mode 100644 index 000000000..1cf442fba --- /dev/null +++ b/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) + } + }) +} diff --git a/pkg/crypto/hash/merkle_tree.go b/pkg/crypto/hash/merkle_tree.go index b12d244bc..0d37a964a 100644 --- a/pkg/crypto/hash/merkle_tree.go +++ b/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 diff --git a/pkg/crypto/hash/merkle_tree_test.go b/pkg/crypto/hash/merkle_tree_test.go index d6494f2d1..6a29a5d8d 100644 --- a/pkg/crypto/hash/merkle_tree_test.go +++ b/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()) diff --git a/pkg/internal/testchain/address.go b/pkg/internal/testchain/address.go index 0ee8620a6..9249bf2a4 100644 --- a/pkg/internal/testchain/address.go +++ b/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