protocol/bc/merkle.go

   1 package bc
   2
   3 import (
   4         "io"
   5         "math"
   6
   7         "github.com/bytom/crypto/sha3pool"
   8 )
   9
  10 var (
  11         leafPrefix     = []byte{0x00}
  12         interiorPrefix = []byte{0x01}
  13 )
  14
  15 type merkleNode interface {
  16         WriteTo(io.Writer) (int64, error)
  17 }
  18
  19 func merkleRoot(nodes []merkleNode) (root Hash, err error) {
  20         switch {
  21         case len(nodes) == 0:
  22                 return EmptyStringHash, nil
  23
  24         case len(nodes) == 1:
  25                 h := sha3pool.Get256()
  26                 defer sha3pool.Put256(h)
  27
  28                 h.Write(leafPrefix)
  29                 nodes[0].WriteTo(h)
  30                 root.ReadFrom(h)
  31                 return root, nil
  32
  33         default:
  34                 k := prevPowerOfTwo(len(nodes))
  35                 left, err := merkleRoot(nodes[:k])
  36                 if err != nil {
  37                         return root, err
  38                 }
  39
  40                 right, err := merkleRoot(nodes[k:])
  41                 if err != nil {
  42                         return root, err
  43                 }
  44
  45                 h := sha3pool.Get256()
  46                 defer sha3pool.Put256(h)
  47                 h.Write(interiorPrefix)
  48                 left.WriteTo(h)
  49                 right.WriteTo(h)
  50                 root.ReadFrom(h)
  51                 return root, nil
  52         }
  53 }
  54
  55 // TxStatusMerkleRoot creates a merkle tree from a slice of TxVerifyResult
  56 func TxStatusMerkleRoot(tvrs []*TxVerifyResult) (root Hash, err error) {
  57         nodes := []merkleNode{}
  58         for _, tvr := range tvrs {
  59                 nodes = append(nodes, tvr)
  60         }
  61         return merkleRoot(nodes)
  62 }
  63
  64 // TxMerkleRoot creates a merkle tree from a slice of transactions
  65 // and returns the root hash of the tree.
  66 func TxMerkleRoot(transactions []*Tx) (root Hash, err error) {
  67         nodes := []merkleNode{}
  68         for _, tx := range transactions {
  69                 nodes = append(nodes, tx.ID)
  70         }
  71         return merkleRoot(nodes)
  72 }
  73
  74 // prevPowerOfTwo returns the largest power of two that is smaller than a given number.
  75 // In other words, for some input n, the prevPowerOfTwo k is a power of two such that
  76 // k < n <= 2k. This is a helper function used during the calculation of a merkle tree.
  77 func prevPowerOfTwo(n int) int {
  78         // If the number is a power of two, divide it by 2 and return.
  79         if n&(n-1) == 0 {
  80                 return n / 2
  81         }
  82
  83         // Otherwise, find the previous PoT.
  84         exponent := uint(math.Log2(float64(n)))
  85         return 1 << exponent // 2^exponent
  86 }