github.com/lbryio/lbcd@v0.22.119/claimtrie/merkletrie/ramtrie.go (about)

     1  package merkletrie
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"runtime"
     7  	"sync"
     8  
     9  	"github.com/lbryio/lbcd/chaincfg/chainhash"
    10  	"github.com/lbryio/lbcd/claimtrie/node"
    11  )
    12  
    13  type MerkleTrie interface {
    14  	SetRoot(h *chainhash.Hash) error
    15  	Update(name []byte, h *chainhash.Hash, restoreChildren bool)
    16  	MerkleHash() *chainhash.Hash
    17  	MerkleHashAllClaims() *chainhash.Hash
    18  	Flush() error
    19  }
    20  
    21  type RamTrie struct {
    22  	collapsedTrie
    23  	bufs *sync.Pool
    24  }
    25  
    26  func NewRamTrie() *RamTrie {
    27  	return &RamTrie{
    28  		bufs: &sync.Pool{
    29  			New: func() interface{} {
    30  				return new(bytes.Buffer)
    31  			},
    32  		},
    33  		collapsedTrie: collapsedTrie{Root: &collapsedVertex{merkleHash: EmptyTrieHash}},
    34  	}
    35  }
    36  
    37  var ErrFullRebuildRequired = errors.New("a full rebuild is required")
    38  
    39  func (rt *RamTrie) SetRoot(h *chainhash.Hash) error {
    40  	if rt.Root.merkleHash.IsEqual(h) {
    41  		runtime.GC()
    42  		return nil
    43  	}
    44  
    45  	// should technically clear the old trie first, but this is abused for partial rebuilds so don't
    46  	return ErrFullRebuildRequired
    47  }
    48  
    49  func (rt *RamTrie) Update(name []byte, h *chainhash.Hash, _ bool) {
    50  	if h == nil {
    51  		rt.Erase(name)
    52  	} else {
    53  		_, n := rt.InsertOrFind(name)
    54  		n.claimHash = h
    55  	}
    56  }
    57  
    58  func (rt *RamTrie) MerkleHash() *chainhash.Hash {
    59  	if h := rt.merkleHash(rt.Root); h == nil {
    60  		return EmptyTrieHash
    61  	}
    62  	return rt.Root.merkleHash
    63  }
    64  
    65  func (rt *RamTrie) merkleHash(v *collapsedVertex) *chainhash.Hash {
    66  	if v.merkleHash != nil {
    67  		return v.merkleHash
    68  	}
    69  
    70  	b := rt.bufs.Get().(*bytes.Buffer)
    71  	defer rt.bufs.Put(b)
    72  	b.Reset()
    73  
    74  	for _, ch := range v.children {
    75  		h := rt.merkleHash(ch)              // h is a pointer; don't destroy its data
    76  		b.WriteByte(ch.key[0])              // nolint : errchk
    77  		b.Write(rt.completeHash(h, ch.key)) // nolint : errchk
    78  	}
    79  
    80  	if v.claimHash != nil {
    81  		b.Write(v.claimHash[:])
    82  	}
    83  
    84  	if b.Len() > 0 {
    85  		h := chainhash.DoubleHashH(b.Bytes())
    86  		v.merkleHash = &h
    87  	}
    88  
    89  	return v.merkleHash
    90  }
    91  
    92  func (rt *RamTrie) completeHash(h *chainhash.Hash, childKey KeyType) []byte {
    93  	var data [chainhash.HashSize + 1]byte
    94  	copy(data[1:], h[:])
    95  	for i := len(childKey) - 1; i > 0; i-- {
    96  		data[0] = childKey[i]
    97  		copy(data[1:], chainhash.DoubleHashB(data[:]))
    98  	}
    99  	return data[1:]
   100  }
   101  
   102  func (rt *RamTrie) MerkleHashAllClaims() *chainhash.Hash {
   103  	if h := rt.merkleHashAllClaims(rt.Root); h == nil {
   104  		return EmptyTrieHash
   105  	}
   106  	return rt.Root.merkleHash
   107  }
   108  
   109  func (rt *RamTrie) merkleHashAllClaims(v *collapsedVertex) *chainhash.Hash {
   110  	if v.merkleHash != nil {
   111  		return v.merkleHash
   112  	}
   113  
   114  	childHashes := make([]*chainhash.Hash, 0, len(v.children))
   115  	for _, ch := range v.children {
   116  		h := rt.merkleHashAllClaims(ch)
   117  		childHashes = append(childHashes, h)
   118  	}
   119  
   120  	claimHash := NoClaimsHash
   121  	if v.claimHash != nil {
   122  		claimHash = v.claimHash
   123  	} else if len(childHashes) == 0 {
   124  		return nil
   125  	}
   126  
   127  	childHash := NoChildrenHash
   128  	if len(childHashes) > 0 {
   129  		// this shouldn't be referencing node; where else can we put this merkle root func?
   130  		childHash = node.ComputeMerkleRoot(childHashes)
   131  	}
   132  
   133  	v.merkleHash = node.HashMerkleBranches(childHash, claimHash)
   134  	return v.merkleHash
   135  }
   136  
   137  func (rt *RamTrie) Flush() error {
   138  	return nil
   139  }