From c1dc927f942a69cd88de9e7a5462fdb1b1d01854 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Mon, 22 Apr 2019 18:28:59 -0700
Subject: [PATCH] blockchain: Add CalcMerkleRoot

---
 blockchain/merkle.go      | 57 ++++++++++++++++++++++++++++++++++++++-
 blockchain/merkle_test.go | 16 ++++++++---
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/blockchain/merkle.go b/blockchain/merkle.go
index d7e567b283..e5dbf0949a 100644
--- a/blockchain/merkle.go
+++ b/blockchain/merkle.go
@@ -86,7 +86,7 @@ func HashMerkleBranches(left *chainhash.Hash, right *chainhash.Hash) *chainhash.
 //
 // The above stored as a linear array is as follows:
 //
-// 	[h1 h2 h3 h4 h12 h34 root]
+//	[h1 h2 h3 h4 h12 h34 root]
 //
 // As the above shows, the merkle root is always the last element in the array.
 //
@@ -154,6 +154,61 @@ 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 slower than BuildMerkleTreeStore, but 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 {
+	// Iteratively push the leaves onto the rolling merkle tree.
+	merkle := NewStump(len(transactions))
+	for i, tx := range transactions {
+		// 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
+			merkle.add(zeroHash)
+		case witness:
+			wSha := tx.MsgTx().WitnessHash()
+			merkle.add(wSha)
+		default:
+			hash := tx.Hash()
+			merkle.add(*hash)
+		}
+	}
+
+	if len(transactions) != 1 && len(transactions)%2 != 0 {
+		tx := transactions[len(transactions)-1]
+		hash := tx.Hash()
+		merkle.add(*hash)
+	}
+
+	return merkle.Roots[0]
+}
+
 // 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
diff --git a/blockchain/merkle_test.go b/blockchain/merkle_test.go
index ed4f688ed4..14beb8fc31 100644
--- a/blockchain/merkle_test.go
+++ b/blockchain/merkle_test.go
@@ -16,12 +16,20 @@ import (
 // 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]
+
+	if calcMerkleRoot != *merkleStoreRoot {
+		t.Fatalf("BuildMerkleTreeStore root does not match "+
+			"CalcMerkleRoot, store: %v calc: %v", 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)
 	}
 }