github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/crypto/hash/merkle_tree.go (about)

     1  package hash
     2  
     3  import (
     4  	"errors"
     5  
     6  	"github.com/nspcc-dev/neo-go/pkg/util"
     7  )
     8  
     9  // MerkleTree implementation.
    10  type MerkleTree struct {
    11  	root  *MerkleTreeNode
    12  	depth int
    13  }
    14  
    15  // NewMerkleTree returns a new MerkleTree object.
    16  func NewMerkleTree(hashes []util.Uint256) (*MerkleTree, error) {
    17  	if len(hashes) == 0 {
    18  		return nil, errors.New("length of the hashes cannot be zero")
    19  	}
    20  
    21  	nodes := make([]*MerkleTreeNode, len(hashes))
    22  	for i := 0; i < len(hashes); i++ {
    23  		nodes[i] = &MerkleTreeNode{
    24  			hash: hashes[i],
    25  		}
    26  	}
    27  
    28  	return &MerkleTree{
    29  		root:  buildMerkleTree(nodes),
    30  		depth: 1,
    31  	}, nil
    32  }
    33  
    34  // Root returns the computed root hash of the MerkleTree.
    35  func (t *MerkleTree) Root() util.Uint256 {
    36  	return t.root.hash
    37  }
    38  
    39  func buildMerkleTree(leaves []*MerkleTreeNode) *MerkleTreeNode {
    40  	if len(leaves) == 0 {
    41  		panic("length of leaves cannot be zero")
    42  	}
    43  	if len(leaves) == 1 {
    44  		return leaves[0]
    45  	}
    46  
    47  	parents := make([]*MerkleTreeNode, (len(leaves)+1)/2)
    48  	for i := 0; i < len(parents); i++ {
    49  		parents[i] = &MerkleTreeNode{}
    50  		parents[i].leftChild = leaves[i*2]
    51  		leaves[i*2].parent = parents[i]
    52  
    53  		if i*2+1 == len(leaves) {
    54  			parents[i].rightChild = parents[i].leftChild
    55  		} else {
    56  			parents[i].rightChild = leaves[i*2+1]
    57  			leaves[i*2+1].parent = parents[i]
    58  		}
    59  
    60  		b1 := parents[i].leftChild.hash.BytesBE()
    61  		b2 := parents[i].rightChild.hash.BytesBE()
    62  		b1 = append(b1, b2...)
    63  		parents[i].hash = DoubleSha256(b1)
    64  	}
    65  
    66  	return buildMerkleTree(parents)
    67  }
    68  
    69  // CalcMerkleRoot calculates the Merkle root hash value for the given slice of hashes.
    70  // It doesn't create a full MerkleTree structure and it uses the given slice as a
    71  // scratchpad, so it will destroy its contents in the process. But it's much more
    72  // memory efficient if you only need a root hash value. While NewMerkleTree would
    73  // make 3*N allocations for N hashes, this function will only make 4.
    74  func CalcMerkleRoot(hashes []util.Uint256) util.Uint256 {
    75  	if len(hashes) == 0 {
    76  		return util.Uint256{}
    77  	}
    78  	if len(hashes) == 1 {
    79  		return hashes[0]
    80  	}
    81  
    82  	scratch := make([]byte, 64)
    83  	parents := hashes[:(len(hashes)+1)/2]
    84  	for i := 0; i < len(parents); i++ {
    85  		copy(scratch, hashes[i*2].BytesBE())
    86  
    87  		if i*2+1 == len(hashes) {
    88  			copy(scratch[32:], hashes[i*2].BytesBE())
    89  		} else {
    90  			copy(scratch[32:], hashes[i*2+1].BytesBE())
    91  		}
    92  
    93  		parents[i] = DoubleSha256(scratch)
    94  	}
    95  
    96  	return CalcMerkleRoot(parents)
    97  }
    98  
    99  // MerkleTreeNode represents a node in the MerkleTree.
   100  type MerkleTreeNode struct {
   101  	hash       util.Uint256
   102  	parent     *MerkleTreeNode
   103  	leftChild  *MerkleTreeNode
   104  	rightChild *MerkleTreeNode
   105  }
   106  
   107  // IsLeaf returns whether this node is a leaf node or not.
   108  func (n *MerkleTreeNode) IsLeaf() bool {
   109  	return n.leftChild == nil && n.rightChild == nil
   110  }
   111  
   112  // IsRoot returns whether this node is a root node or not.
   113  func (n *MerkleTreeNode) IsRoot() bool {
   114  	return n.parent == nil
   115  }