Merge pull request #1979 from kcalvinalvin/merkle-calc-fast

blockchain, integration, mining, main: Rolling merkle root calculation

blockchain/fullblocktests/generate.go

@@ -309,8 +309,7 @@ func calcMerkleRoot(txns []*wire.MsgTx) chainhash.Hash {
 	for _, tx := range txns {
 		utilTxns = append(utilTxns, btcutil.NewTx(tx))
 	}
-	merkles := blockchain.BuildMerkleTreeStore(utilTxns, false)
-	return *merkles[len(merkles)-1]
+	return blockchain.CalcMerkleRoot(utilTxns, false)
 }
 
 // solveBlock attempts to find a nonce which makes the passed block header hash

blockchain/merkle.go

@@ -58,14 +58,13 @@ func nextPowerOfTwo(n int) int {
 // HashMerkleBranches takes two hashes, treated as the left and right tree
 // nodes, and returns the hash of their concatenation.  This is a helper
 // function used to aid in the generation of a merkle tree.
-func HashMerkleBranches(left *chainhash.Hash, right *chainhash.Hash) *chainhash.Hash {
+func HashMerkleBranches(left, right *chainhash.Hash) chainhash.Hash {
 	// Concatenate the left and right nodes.
 	var hash [chainhash.HashSize * 2]byte
 	copy(hash[:chainhash.HashSize], left[:])
 	copy(hash[chainhash.HashSize:], right[:])
 
-	newHash := chainhash.DoubleHashH(hash[:])
-	return &newHash
+	return chainhash.DoubleHashH(hash[:])
 }
 
 // BuildMerkleTreeStore creates a merkle tree from a slice of transactions,
@@ -140,13 +139,13 @@ func BuildMerkleTreeStore(transactions []*btcutil.Tx, witness bool) []*chainhash
 		// hashing the concatenation of the left child with itself.
 		case merkles[i+1] == nil:
 			newHash := HashMerkleBranches(merkles[i], merkles[i])
-			merkles[offset] = newHash
+			merkles[offset] = &newHash
 
 		// The normal case sets the parent node to the double sha256
 		// of the concatentation of the left and right children.
 		default:
 			newHash := HashMerkleBranches(merkles[i], merkles[i+1])
-			merkles[offset] = newHash
+			merkles[offset] = &newHash
 		}
 		offset++
 	}
@@ -154,6 +153,36 @@ func BuildMerkleTreeStore(transactions []*btcutil.Tx, witness bool) []*chainhash
 	return merkles
 }
 
+// CalcMerkleRoot computes the merkle root over a set of hashed leaves. The
+// interior nodes are computed opportunistically as the leaves are added to the
+// abstract tree to reduce the total number of allocations. Throughout the
+// computation, this computation only requires storing O(log n) interior
+// nodes.
+//
+// This method differs from BuildMerkleTreeStore in that the interior nodes are
+// discarded instead of being returned along with the root. CalcMerkleRoot is
+// slightly faster than BuildMerkleTreeStore and requires significantly less
+// memory and fewer allocations.
+//
+// A merkle tree is a tree in which every non-leaf node is the hash of its
+// children nodes. A diagram depicting how this works for bitcoin transactions
+// where h(x) is a double sha256 follows:
+//
+//	         root = h1234 = h(h12 + h34)
+//	        /                           \
+//	  h12 = h(h1 + h2)            h34 = h(h3 + h4)
+//	   /            \              /            \
+//	h1 = h(tx1)  h2 = h(tx2)    h3 = h(tx3)  h4 = h(tx4)
+//
+// The additional bool parameter indicates if we are generating the merkle tree
+// using witness transaction id's rather than regular transaction id's. This
+// also presents an additional case wherein the wtxid of the coinbase transaction
+// is the zeroHash.
+func CalcMerkleRoot(transactions []*btcutil.Tx, witness bool) chainhash.Hash {
+	s := newRollingMerkleTreeStore(uint64(len(transactions)))
+	return s.calcMerkleRoot(transactions, witness)
+}
+
 // ExtractWitnessCommitment attempts to locate, and return the witness
 // commitment for a block. The witness commitment is of the form:
 // SHA256(witness root || witness nonce). The function additionally returns a
@@ -246,8 +275,7 @@ func ValidateWitnessCommitment(blk *btcutil.Block) error {
 	// the extracted witnessCommitment is equal to:
 	// SHA256(witnessMerkleRoot || witnessNonce). Where witnessNonce is the
 	// coinbase transaction's only witness item.
-	witnessMerkleTree := BuildMerkleTreeStore(blk.Transactions(), true)
-	witnessMerkleRoot := witnessMerkleTree[len(witnessMerkleTree)-1]
+	witnessMerkleRoot := CalcMerkleRoot(blk.Transactions(), true)
 
 	var witnessPreimage [chainhash.HashSize * 2]byte
 	copy(witnessPreimage[:], witnessMerkleRoot[:])

blockchain/merkle_test.go

@@ -5,19 +5,105 @@
 package blockchain
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/btcsuite/btcd/btcutil"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/stretchr/testify/require"
 )
 
 // TestMerkle tests the BuildMerkleTreeStore API.
 func TestMerkle(t *testing.T) {
 	block := btcutil.NewBlock(&Block100000)
-	merkles := BuildMerkleTreeStore(block.Transactions(), false)
-	calculatedMerkleRoot := merkles[len(merkles)-1]
+	calcMerkleRoot := CalcMerkleRoot(block.Transactions(), false)
+	merkleStoreTree := BuildMerkleTreeStore(block.Transactions(), false)
+	merkleStoreRoot := merkleStoreTree[len(merkleStoreTree)-1]
+
+	require.Equal(t, *merkleStoreRoot, calcMerkleRoot)
+
 	wantMerkle := &Block100000.Header.MerkleRoot
-	if !wantMerkle.IsEqual(calculatedMerkleRoot) {
+	if !wantMerkle.IsEqual(&calcMerkleRoot) {
 		t.Errorf("BuildMerkleTreeStore: merkle root mismatch - "+
-			"got %v, want %v", calculatedMerkleRoot, wantMerkle)
+			"got %v, want %v", calcMerkleRoot, wantMerkle)
+	}
+}
+
+func makeHashes(size int) []*chainhash.Hash {
+	var hashes = make([]*chainhash.Hash, size)
+	for i := range hashes {
+		hashes[i] = new(chainhash.Hash)
+	}
+	return hashes
+}
+
+func makeTxs(size int) []*btcutil.Tx {
+	var txs = make([]*btcutil.Tx, size)
+	for i := range txs {
+		tx := btcutil.NewTx(wire.NewMsgTx(2))
+		tx.Hash()
+		txs[i] = tx
+	}
+	return txs
+}
+
+// BenchmarkRollingMerkle benches the RollingMerkleTree while varying the number
+// of leaves pushed to the tree.
+func BenchmarkRollingMerkle(b *testing.B) {
+	sizes := []int{
+		1000,
+		2000,
+		4000,
+		8000,
+		16000,
+		32000,
+	}
+
+	for _, size := range sizes {
+		txs := makeTxs(size)
+		name := fmt.Sprintf("%d", size)
+		b.Run(name, func(b *testing.B) {
+			benchmarkRollingMerkle(b, txs)
+		})
+	}
+}
+
+// BenchmarkMerkle benches the BuildMerkleTreeStore while varying the number
+// of leaves pushed to the tree.
+func BenchmarkMerkle(b *testing.B) {
+	sizes := []int{
+		1000,
+		2000,
+		4000,
+		8000,
+		16000,
+		32000,
+	}
+
+	for _, size := range sizes {
+		txs := makeTxs(size)
+		name := fmt.Sprintf("%d", size)
+		b.Run(name, func(b *testing.B) {
+			benchmarkMerkle(b, txs)
+		})
+	}
+}
+
+func benchmarkRollingMerkle(b *testing.B, txs []*btcutil.Tx) {
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		CalcMerkleRoot(txs, false)
+	}
+}
+
+func benchmarkMerkle(b *testing.B, txs []*btcutil.Tx) {
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		BuildMerkleTreeStore(txs, false)
 	}
 }

blockchain/rolling_merkle.go

@@ -0,0 +1,136 @@
+package blockchain
+
+import (
+	"math/bits"
+
+	"github.com/btcsuite/btcd/btcutil"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+)
+
+// rollingMerkleTreeStore calculates the merkle root by only allocating O(logN)
+// memory where N is the total amount of leaves being included in the tree.
+type rollingMerkleTreeStore struct {
+	// roots are where the temporary merkle roots get stored while the
+	// merkle root is being calculated.
+	roots []chainhash.Hash
+
+	// numLeaves is the total leaves the store has processed.  numLeaves
+	// is required for the root calculation algorithm to work.
+	numLeaves uint64
+}
+
+// newRollingMerkleTreeStore returns a rollingMerkleTreeStore with the roots
+// allocated based on the passed in size.
+//
+// NOTE: If more elements are added in than the passed in size, there will be
+// additional allocations which in turn hurts performance.
+func newRollingMerkleTreeStore(size uint64) rollingMerkleTreeStore {
+	var alloc int
+	if size != 0 {
+		alloc = bits.Len64(size - 1)
+	}
+	return rollingMerkleTreeStore{roots: make([]chainhash.Hash, 0, alloc)}
+}
+
+// add adds a single hash to the merkle tree store.  Refer to algorithm 1 "AddOne" in
+// the utreexo paper (https://eprint.iacr.org/2019/611.pdf) for the exact algorithm.
+func (s *rollingMerkleTreeStore) add(add chainhash.Hash) {
+	// We can tell where the roots are by looking at the binary representation
+	// of the numLeaves.  Wherever there's a 1, there's a root.
+	//
+	// numLeaves of 8 will be '1000' in binary, so there will be one root at
+	// row 3. numLeaves of 3 will be '11' in binary, so there's two roots.  One at
+	// row 0 and one at row 1.  Row 0 is the leaf row.
+	//
+	// In this loop below, we're looking for these roots by checking if there's
+	// a '1', starting from the LSB.  If there is a '1', we'll hash the root being
+	// added with that root until we hit a '0'.
+	newRoot := add
+	for h := uint8(0); (s.numLeaves>>h)&1 == 1; h++ {
+		// Pop off the last root.
+		var root chainhash.Hash
+		root, s.roots = s.roots[len(s.roots)-1], s.roots[:len(s.roots)-1]
+
+		// Calculate the hash of the new root and append it.
+		newRoot = HashMerkleBranches(&root, &newRoot)
+	}
+	s.roots = append(s.roots, newRoot)
+	s.numLeaves++
+}
+
+// calcMerkleRoot returns the merkle root for the passed in transactions.
+func (s *rollingMerkleTreeStore) calcMerkleRoot(adds []*btcutil.Tx, witness bool) chainhash.Hash {
+	for i := range adds {
+		// If we're computing a witness merkle root, instead of the
+		// regular txid, we use the modified wtxid which includes a
+		// transaction's witness data within the digest.  Additionally,
+		// the coinbase's wtxid is all zeroes.
+		switch {
+		case witness && i == 0:
+			var zeroHash chainhash.Hash
+			s.add(zeroHash)
+		case witness:
+			s.add(*adds[i].WitnessHash())
+		default:
+			s.add(*adds[i].Hash())
+		}
+	}
+
+	// If we only have one leaf, then the hash of that tx is the merkle root.
+	if s.numLeaves == 1 {
+		return s.roots[0]
+	}
+
+	// Add on the last tx again if there's an odd number of txs.
+	if len(adds) > 0 && len(adds)%2 != 0 {
+		switch {
+		case witness:
+			s.add(*adds[len(adds)-1].WitnessHash())
+		default:
+			s.add(*adds[len(adds)-1].Hash())
+		}
+	}
+
+	// If we still have more than 1 root after adding on the last tx again,
+	// we need to do the same for the upper rows.
+	//
+	// For exmaple, the below tree has 6 leaves.  For row 1, you'll need to
+	// hash 'F' with itself to create 'C' so you have something to hash with
+	// 'B'.  For bigger trees we may need to do the same in rows 2 or 3 as
+	// well.
+	//
+	// row :3         A
+	//              /   \
+	// row :2     B       C
+	//           / \     / \
+	// row :1   D   E   F   F
+	//         / \ / \ / \
+	// row :0  1 2 3 4 5 6
+	for len(s.roots) > 1 {
+		// If we have to keep adding the last node in the set, bitshift
+		// the num leaves right by 1.  This effectively moves the row up
+		// for calculation.  We do this until we reach a row where there's
+		// an odd number of leaves.
+		//
+		// row :3         A
+		//              /   \
+		// row :2     B       C        D
+		//           / \     / \     /   \
+		// row :1   E   F   G   H   I     J
+		//         / \ / \ / \ / \ / \   / \
+		// row :0  1 2 3 4 5 6 7 8 9 10 11 12
+		//
+		// In the above tree, 12 leaves were added and there's an odd amount
+		// of leaves at row 2.  Because of this, we'll bitshift right twice.
+		currentLeaves := s.numLeaves
+		for h := uint8(0); (currentLeaves>>h)&1 == 0; h++ {
+			s.numLeaves >>= 1
+		}
+
+		// Add the last root again so that it'll get hashed with itself.
+		h := s.roots[len(s.roots)-1]
+		s.add(h)
+	}
+
+	return s.roots[0]
+}

blockchain/rolling_merkle_test.go

@@ -0,0 +1,174 @@
+package blockchain
+
+import (
+	"testing"
+
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRollingMerkleAdd(t *testing.T) {
+	tests := []struct {
+		leaves            []chainhash.Hash
+		expectedRoots     []chainhash.Hash
+		expectedNumLeaves uint64
+	}{
+		// 00  (00 is also a root)
+		{
+			leaves: []chainhash.Hash{
+				{0x00},
+			},
+			expectedRoots: []chainhash.Hash{
+				{0x00},
+			},
+			expectedNumLeaves: 1,
+		},
+
+		// root
+		// |---\
+		// 00  01
+		{
+			leaves: []chainhash.Hash{
+				{0x00},
+				{0x01},
+			},
+			expectedRoots: []chainhash.Hash{
+				func() chainhash.Hash {
+					hash, err := chainhash.NewHashFromStr(
+						"c2bf026e62af95cd" +
+							"7b785e2cd5a5f1ec" +
+							"01fafda85886a8eb" +
+							"d34482c0b05dc2c2")
+					require.NoError(t, err)
+					return *hash
+				}(),
+			},
+			expectedNumLeaves: 2,
+		},
+
+		// root
+		// |---\
+		// 00  01  02
+		{
+			leaves: []chainhash.Hash{
+				{0x00},
+				{0x01},
+				{0x02},
+			},
+			expectedRoots: []chainhash.Hash{
+				func() chainhash.Hash {
+					hash, err := chainhash.NewHashFromStr(
+						"c2bf026e62af95cd" +
+							"7b785e2cd5a5f1ec" +
+							"01fafda85886a8eb" +
+							"d34482c0b05dc2c2")
+					require.NoError(t, err)
+					return *hash
+				}(),
+				{0x02},
+			},
+			expectedNumLeaves: 3,
+		},
+
+		// root
+		// |-------\
+		// br      br
+		// |---\   |---\
+		// 00  01  02  03
+		{
+			leaves: []chainhash.Hash{
+				{0x00},
+				{0x01},
+				{0x02},
+				{0x03},
+			},
+			expectedRoots: []chainhash.Hash{
+				func() chainhash.Hash {
+					hash, err := chainhash.NewHashFromStr(
+						"270714425ea73eb8" +
+							"5942f0f705788f25" +
+							"1fefa3f533410a3f" +
+							"338de46e641082c4")
+					require.NoError(t, err)
+					return *hash
+				}(),
+			},
+			expectedNumLeaves: 4,
+		},
+
+		// root
+		// |-------\
+		// br      br
+		// |---\   |---\
+		// 00  01  02  03  04
+		{
+			leaves: []chainhash.Hash{
+				{0x00},
+				{0x01},
+				{0x02},
+				{0x03},
+				{0x04},
+			},
+			expectedRoots: []chainhash.Hash{
+				func() chainhash.Hash {
+					hash, err := chainhash.NewHashFromStr(
+						"270714425ea73eb8" +
+							"5942f0f705788f25" +
+							"1fefa3f533410a3f" +
+							"338de46e641082c4")
+					require.NoError(t, err)
+					return *hash
+				}(),
+				{0x04},
+			},
+			expectedNumLeaves: 5,
+		},
+
+		// root
+		// |-------\
+		// br      br      root
+		// |---\   |---\   |---\
+		// 00  01  02  03  04  05
+		{
+			leaves: []chainhash.Hash{
+				{0x00},
+				{0x01},
+				{0x02},
+				{0x03},
+				{0x04},
+				{0x05},
+			},
+			expectedRoots: []chainhash.Hash{
+				func() chainhash.Hash {
+					hash, err := chainhash.NewHashFromStr(
+						"270714425ea73eb8" +
+							"5942f0f705788f25" +
+							"1fefa3f533410a3f" +
+							"338de46e641082c4")
+					require.NoError(t, err)
+					return *hash
+				}(),
+				func() chainhash.Hash {
+					hash, err := chainhash.NewHashFromStr(
+						"e5c2407ba454ffeb" +
+							"28cf0c50c5c293a8" +
+							"4c9a75788f8a8f35" +
+							"ccb974e606280377")
+					require.NoError(t, err)
+					return *hash
+				}(),
+			},
+			expectedNumLeaves: 6,
+		},
+	}
+
+	for _, test := range tests {
+		s := newRollingMerkleTreeStore(uint64(len(test.leaves)))
+		for _, leaf := range test.leaves {
+			s.add(leaf)
+		}
+
+		require.Equal(t, s.roots, test.expectedRoots)
+		require.Equal(t, s.numLeaves, test.expectedNumLeaves)
+	}
+}

blockchain/validate.go

@@ -529,12 +529,11 @@ func checkBlockSanity(block *btcutil.Block, powLimit *big.Int, timeSource Median
 	// checks.  Bitcoind builds the tree here and checks the merkle root
 	// after the following checks, but there is no reason not to check the
 	// merkle root matches here.
-	merkles := BuildMerkleTreeStore(block.Transactions(), false)
-	calculatedMerkleRoot := merkles[len(merkles)-1]
-	if !header.MerkleRoot.IsEqual(calculatedMerkleRoot) {
+	calcMerkleRoot := CalcMerkleRoot(block.Transactions(), false)
+	if !header.MerkleRoot.IsEqual(&calcMerkleRoot) {
 		str := fmt.Sprintf("block merkle root is invalid - block "+
 			"header indicates %v, but calculated value is %v",
-			header.MerkleRoot, calculatedMerkleRoot)
+			header.MerkleRoot, calcMerkleRoot)
 		return ruleError(ErrBadMerkleRoot, str)
 	}
 

btcutil/bloom/merkleblock.go

@@ -41,7 +41,8 @@ func (m *merkleBlock) calcHash(height, pos uint32) *chainhash.Hash {
 	} else {
 		right = left
 	}
-	return blockchain.HashMerkleBranches(left, right)
+	res := blockchain.HashMerkleBranches(left, right)
+	return &res
 }
 
 // traverseAndBuild builds a partial merkle tree using a recursive depth-first

integration/rpctest/blockgen.go

@@ -197,12 +197,12 @@ func CreateBlock(prevBlock *btcutil.Block, inclusionTxs []*btcutil.Tx,
 		_ = mining.AddWitnessCommitment(coinbaseTx, blockTxns)
 	}
 
-	merkles := blockchain.BuildMerkleTreeStore(blockTxns, false)
+	merkleRoot := blockchain.CalcMerkleRoot(blockTxns, false)
 	var block wire.MsgBlock
 	block.Header = wire.BlockHeader{
 		Version:    blockVersion,
 		PrevBlock:  *prevHash,
-		MerkleRoot: *merkles[len(merkles)-1],
+		MerkleRoot: merkleRoot,
 		Timestamp:  ts,
 		Bits:       net.PowLimitBits,
 	}

mining/mining.go

@@ -823,12 +823,11 @@ mempoolLoop:
 	}
 
 	// Create a new block ready to be solved.
-	merkles := blockchain.BuildMerkleTreeStore(blockTxns, false)
 	var msgBlock wire.MsgBlock
 	msgBlock.Header = wire.BlockHeader{
 		Version:    nextBlockVersion,
 		PrevBlock:  best.Hash,
-		MerkleRoot: *merkles[len(merkles)-1],
+		MerkleRoot: blockchain.CalcMerkleRoot(blockTxns, false),
 		Timestamp:  ts,
 		Bits:       reqDifficulty,
 	}
@@ -875,9 +874,7 @@ func AddWitnessCommitment(coinbaseTx *btcutil.Tx,
 	// Next, obtain the merkle root of a tree which consists of the
 	// wtxid of all transactions in the block. The coinbase
 	// transaction will have a special wtxid of all zeroes.
-	witnessMerkleTree := blockchain.BuildMerkleTreeStore(blockTxns,
-		true)
-	witnessMerkleRoot := witnessMerkleTree[len(witnessMerkleTree)-1]
+	witnessMerkleRoot := blockchain.CalcMerkleRoot(blockTxns, true)
 
 	// The preimage to the witness commitment is:
 	// witnessRoot || coinbaseWitness
@@ -953,8 +950,8 @@ func (g *BlkTmplGenerator) UpdateExtraNonce(msgBlock *wire.MsgBlock, blockHeight
 
 	// Recalculate the merkle root with the updated extra nonce.
 	block := btcutil.NewBlock(msgBlock)
-	merkles := blockchain.BuildMerkleTreeStore(block.Transactions(), false)
-	msgBlock.Header.MerkleRoot = *merkles[len(merkles)-1]
+	merkleRoot := blockchain.CalcMerkleRoot(block.Transactions(), false)
+	msgBlock.Header.MerkleRoot = merkleRoot
 	return nil
 }
 

rpcserver.go

@@ -1665,8 +1665,8 @@ func (state *gbtWorkState) updateBlockTemplate(s *rpcServer, useCoinbaseValue bo
 
 			// Update the merkle root.
 			block := btcutil.NewBlock(template.Block)
-			merkles := blockchain.BuildMerkleTreeStore(block.Transactions(), false)
-			template.Block.Header.MerkleRoot = *merkles[len(merkles)-1]
+			merkleRoot := blockchain.CalcMerkleRoot(block.Transactions(), false)
+			template.Block.Header.MerkleRoot = merkleRoot
 		}
 
 		// Set locals for convenience.