github.com/aergoio/aergo@v1.3.1/internal/merkle/merkle.go (about)

     1  package merkle
     2  
     3  import (
     4  	"github.com/minio/sha256-simd"
     5  	"hash"
     6  )
     7  
     8  type MerkleEntry interface {
     9  	GetHash() []byte
    10  }
    11  
    12  var (
    13  	HashSize = 32
    14  	nilHash  = make([]byte, HashSize)
    15  	//logger = log.NewLogger("merkle")
    16  )
    17  
    18  func CalculateMerkleRoot(entries []MerkleEntry) []byte {
    19  	merkles := CalculateMerkleTree(entries)
    20  
    21  	return merkles[len(merkles)-1]
    22  }
    23  
    24  func CalculateMerkleTree(entries []MerkleEntry) [][]byte {
    25  	var merkles [][]byte
    26  	entriesLen := len(entries)
    27  
    28  	if entriesLen == 0 {
    29  		merkles = append(merkles, nilHash)
    30  		return merkles
    31  	}
    32  
    33  	//leaf count for full binary tree = 2 ^ n > entryLen
    34  	getLeafCount := func(num int) int {
    35  		if (num&num - 1) == 0 {
    36  			return num
    37  		}
    38  		x := 1
    39  		for x < num {
    40  			x = x << 1
    41  		}
    42  		return x
    43  	}
    44  
    45  	calcMerkle := func(hasher hash.Hash, lc []byte, rc []byte) []byte {
    46  		hasher.Reset()
    47  		hasher.Write(lc)
    48  		hasher.Write(rc)
    49  		return hasher.Sum(nil)
    50  	}
    51  
    52  	hasher := sha256.New()
    53  
    54  	leafCount := getLeafCount(len(entries))
    55  	totalCount := leafCount*2 - 1
    56  
    57  	//logger.Debug().Int("leafcount", leafCount).Int("totCount", totalCount).Msg("start merkling")
    58  
    59  	merkles = make([][]byte, totalCount)
    60  
    61  	// init leaf hash (0 <= node# < entry len)
    62  	for i, entry := range entries {
    63  		merkles[i] = entry.GetHash()
    64  	}
    65  
    66  	// start from branch height 1 (merkles[leafcount] ~ merkles[totalCount -1])
    67  	var childIdx = 0
    68  	var lc, rc int
    69  	for i := leafCount; i < totalCount; i++ {
    70  		// hash of branch node is zero if all child not exist
    71  		lc = childIdx
    72  		rc = childIdx + 1
    73  		childIdx += 2
    74  
    75  		// If all child is nil, merkle is nil
    76  		if merkles[lc] == nil {
    77  			merkles[i] = nil
    78  			continue
    79  		}
    80  		// If only exist left child, copy left child hash to right child
    81  		if merkles[rc] == nil {
    82  			merkles[rc] = merkles[lc]
    83  		}
    84  
    85  		merkles[i] = calcMerkle(hasher, merkles[lc], merkles[rc])
    86  		//logger.Debug().Int("i", i).Str("m", EncodeB64(merkles[i])).Msg("merkling")
    87  	}
    88  
    89  	return merkles
    90  }