aergoio/aergo

package merkle

import (
    "hash"

    "github.com/minio/sha256-simd"
)

type MerkleEntry interface {
    GetHash() []byte
}

var (
    HashSize = 32
    nilHash  = make([]byte, HashSize)
    //logger = log.NewLogger("merkle")
)

func CalculateMerkleRoot(entries []MerkleEntry) []byte {
    merkles := CalculateMerkleTree(entries)

    return merkles[len(merkles)-1]
}

func CalculateMerkleTree(entries []MerkleEntry) [][]byte {
    var merkles [][]byte
    entriesLen := len(entries)

    if entriesLen == 0 {
        merkles = append(merkles, nilHash)
        return merkles
    }

    //leaf count for full binary tree = 2 ^ n > entryLen
    getLeafCount := func(num int) int {
        if (num&num - 1) == 0 {
            return num
        }
        x := 1
        for x < num {
            x = x << 1
        }
        return x
    }

    calcMerkle := func(hasher hash.Hash, lc []byte, rc []byte) []byte {
        hasher.Reset()
        hasher.Write(lc)
        hasher.Write(rc)
        return hasher.Sum(nil)
    }

    hasher := sha256.New()

    leafCount := getLeafCount(len(entries))
    totalCount := leafCount*2 - 1

    //logger.Debug().Int("leafcount", leafCount).Int("totCount", totalCount).Msg("start merkling")

    merkles = make([][]byte, totalCount)

    // init leaf hash (0 <= node# < entry len)
    for i, entry := range entries {
        merkles[i] = entry.GetHash()
    }

    // start from branch height 1 (merkles[leafcount] ~ merkles[totalCount -1])
    var childIdx = 0
    var lc, rc int
    for i := leafCount; i < totalCount; i++ {
        // hash of branch node is zero if all child not exist
        lc = childIdx
        rc = childIdx + 1
        childIdx += 2

        // If all child is nil, merkle is nil
        if merkles[lc] == nil {
            merkles[i] = nil
            continue
        }
        // If only exist left child, copy left child hash to right child
        if merkles[rc] == nil {
            merkles[rc] = merkles[lc]
        }

        merkles[i] = calcMerkle(hasher, merkles[lc], merkles[rc])
        //logger.Debug().Int("i", i).Str("m", EncodeB64(merkles[i])).Msg("merkling")
    }

    return merkles
}