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,20 +139,50 @@ 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++
 	}
 
 	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]
+}