github.com/mit-dci/lit@v0.0.0-20221102210550-8c3d3b49f2ce/btcutil/bloom/merkleblock.go (about)

     1  // Copyright (c) 2013-2016 The btcsuite developers
     2  // Use of this source code is governed by an ISC
     3  // license that can be found in the LICENSE file.
     4  
     5  package bloom
     6  
     7  import (
     8  	"github.com/mit-dci/lit/btcutil"
     9  	"github.com/mit-dci/lit/btcutil/blockchain"
    10  	"github.com/mit-dci/lit/btcutil/chaincfg/chainhash"
    11  	"github.com/mit-dci/lit/wire"
    12  )
    13  
    14  // merkleBlock is used to house intermediate information needed to generate a
    15  // wire.MsgMerkleBlock according to a filter.
    16  type merkleBlock struct {
    17  	numTx       uint32
    18  	allHashes   []*chainhash.Hash
    19  	finalHashes []*chainhash.Hash
    20  	matchedBits []byte
    21  	bits        []byte
    22  }
    23  
    24  // calcTreeWidth calculates and returns the the number of nodes (width) or a
    25  // merkle tree at the given depth-first height.
    26  func (m *merkleBlock) calcTreeWidth(height uint32) uint32 {
    27  	return (m.numTx + (1 << height) - 1) >> height
    28  }
    29  
    30  // calcHash returns the hash for a sub-tree given a depth-first height and
    31  // node position.
    32  func (m *merkleBlock) calcHash(height, pos uint32) *chainhash.Hash {
    33  	if height == 0 {
    34  		return m.allHashes[pos]
    35  	}
    36  
    37  	var right *chainhash.Hash
    38  	left := m.calcHash(height-1, pos*2)
    39  	if pos*2+1 < m.calcTreeWidth(height-1) {
    40  		right = m.calcHash(height-1, pos*2+1)
    41  	} else {
    42  		right = left
    43  	}
    44  	return blockchain.HashMerkleBranches(left, right)
    45  }
    46  
    47  // traverseAndBuild builds a partial merkle tree using a recursive depth-first
    48  // approach.  As it calculates the hashes, it also saves whether or not each
    49  // node is a parent node and a list of final hashes to be included in the
    50  // merkle block.
    51  func (m *merkleBlock) traverseAndBuild(height, pos uint32) {
    52  	// Determine whether this node is a parent of a matched node.
    53  	var isParent byte
    54  	for i := pos << height; i < (pos+1)<<height && i < m.numTx; i++ {
    55  		isParent |= m.matchedBits[i]
    56  	}
    57  	m.bits = append(m.bits, isParent)
    58  
    59  	// When the node is a leaf node or not a parent of a matched node,
    60  	// append the hash to the list that will be part of the final merkle
    61  	// block.
    62  	if height == 0 || isParent == 0x00 {
    63  		m.finalHashes = append(m.finalHashes, m.calcHash(height, pos))
    64  		return
    65  	}
    66  
    67  	// At this point, the node is an internal node and it is the parent of
    68  	// of an included leaf node.
    69  
    70  	// Descend into the left child and process its sub-tree.
    71  	m.traverseAndBuild(height-1, pos*2)
    72  
    73  	// Descend into the right child and process its sub-tree if
    74  	// there is one.
    75  	if pos*2+1 < m.calcTreeWidth(height-1) {
    76  		m.traverseAndBuild(height-1, pos*2+1)
    77  	}
    78  }
    79  
    80  // NewMerkleBlock returns a new *wire.MsgMerkleBlock and an array of the matched
    81  // transaction index numbers based on the passed block and filter.
    82  func NewMerkleBlock(block *btcutil.Block, filter *Filter) (*wire.MsgMerkleBlock, []uint32) {
    83  	numTx := uint32(len(block.Transactions()))
    84  	mBlock := merkleBlock{
    85  		numTx:       numTx,
    86  		allHashes:   make([]*chainhash.Hash, 0, numTx),
    87  		matchedBits: make([]byte, 0, numTx),
    88  	}
    89  
    90  	// Find and keep track of any transactions that match the filter.
    91  	var matchedIndices []uint32
    92  	for txIndex, tx := range block.Transactions() {
    93  		if filter.MatchTxAndUpdate(tx) {
    94  			mBlock.matchedBits = append(mBlock.matchedBits, 0x01)
    95  			matchedIndices = append(matchedIndices, uint32(txIndex))
    96  		} else {
    97  			mBlock.matchedBits = append(mBlock.matchedBits, 0x00)
    98  		}
    99  		mBlock.allHashes = append(mBlock.allHashes, tx.Hash())
   100  	}
   101  
   102  	// Calculate the number of merkle branches (height) in the tree.
   103  	height := uint32(0)
   104  	for mBlock.calcTreeWidth(height) > 1 {
   105  		height++
   106  	}
   107  
   108  	// Build the depth-first partial merkle tree.
   109  	mBlock.traverseAndBuild(height, 0)
   110  
   111  	// Create and return the merkle block.
   112  	msgMerkleBlock := wire.MsgMerkleBlock{
   113  		Header:       block.MsgBlock().Header,
   114  		Transactions: uint32(mBlock.numTx),
   115  		Hashes:       make([]*chainhash.Hash, 0, len(mBlock.finalHashes)),
   116  		Flags:        make([]byte, (len(mBlock.bits)+7)/8),
   117  	}
   118  	for _, hash := range mBlock.finalHashes {
   119  		msgMerkleBlock.AddTxHash(hash)
   120  	}
   121  	for i := uint32(0); i < uint32(len(mBlock.bits)); i++ {
   122  		msgMerkleBlock.Flags[i/8] |= mBlock.bits[i] << (i % 8)
   123  	}
   124  	return &msgMerkleBlock, matchedIndices
   125  }