github.com/iotexproject/iotex-core@v1.14.1-rc1/crypto/merkle.go (about)

     1  // Copyright (c) 2019 IoTeX Foundation
     2  // This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
     3  // or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
     4  // This source code is governed by Apache License 2.0 that can be found in the LICENSE file.
     5  
     6  package crypto
     7  
     8  import (
     9  	"github.com/iotexproject/go-pkgs/hash"
    10  )
    11  
    12  // Merkle tree struct
    13  type Merkle struct {
    14  	root hash.Hash256
    15  	leaf []hash.Hash256
    16  	size int
    17  }
    18  
    19  // NewMerkleTree creates a merkle tree given hashed leaves
    20  func NewMerkleTree(leaves []hash.Hash256) *Merkle {
    21  	size := len(leaves)
    22  	if size == 0 {
    23  		return nil
    24  	}
    25  
    26  	mk := &Merkle{
    27  		leaf: make([]hash.Hash256, (size+1)>>1<<1),
    28  		size: size,
    29  	}
    30  
    31  	copy(mk.leaf, leaves)
    32  
    33  	if size == 1 {
    34  		mk.root = mk.leaf[0]
    35  		return mk
    36  	}
    37  
    38  	// copy the last hash if original size is odd number
    39  	if size != len(mk.leaf) {
    40  		mk.leaf[size] = mk.leaf[size-1]
    41  		mk.size = len(mk.leaf)
    42  	}
    43  
    44  	return mk
    45  }
    46  
    47  // HashTree calculates the root hash of a merkle tree
    48  func (mk *Merkle) HashTree() hash.Hash256 {
    49  	if mk.root != hash.ZeroHash256 {
    50  		return mk.root
    51  	}
    52  
    53  	length := mk.size >> 1
    54  	merkle := make([]hash.Hash256, length)
    55  
    56  	// first round, compute hash from original leaf
    57  	for i := 0; i < length; i++ {
    58  		h := mk.leaf[i<<1][:]
    59  		h = append(h, mk.leaf[i<<1+1][:]...)
    60  		merkle[i] = hash.Hash256b(h)
    61  	}
    62  
    63  	for length > 1 {
    64  		if length&1 != 0 {
    65  			merkle = append(merkle, merkle[length-1])
    66  			length++
    67  		}
    68  
    69  		length >>= 1
    70  		for i := 0; i < length; i++ {
    71  			h := merkle[i<<1][:]
    72  			h = append(h, merkle[i<<1+1][:]...)
    73  			merkle[i] = hash.Hash256b(h)
    74  		}
    75  		merkle = merkle[0:length]
    76  	}
    77  
    78  	mk.root = merkle[0]
    79  	return mk.root
    80  }