blockchain: Add RollingMerkleTree

RollingMerkleTree is a much more memory efficient way of calculating the
merkle root of a tx commitment inside the bitcoin block header.  The
current way of calculating the merkle root allocates 2*N elements. With
the RollingMerkleTree, we are able to reduce the memory allocated to
log2(N).

This results in significant memory savings (99.9% in an average block),
allowing for a faster block verification.
This commit is contained in:
Conner Fromknecht 2019-01-17 02:14:26 -08:00 committed by Calvin Kim
parent 3ba9feeeee
commit 2bb6824067
2 changed files with 310 additions and 0 deletions

View File

@ -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]
}

View File

@ -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)
}
}