lbcd/claimtrie/merkletrie/ramtrie.go

140 lines
3 KiB
Go
Raw Normal View History

package merkletrie
import (
"bytes"
"errors"
"runtime"
"sync"
"github.com/lbryio/lbcd/chaincfg/chainhash"
"github.com/lbryio/lbcd/claimtrie/node"
)
type MerkleTrie interface {
SetRoot(h *chainhash.Hash) error
Update(name []byte, h *chainhash.Hash, restoreChildren bool)
MerkleHash() *chainhash.Hash
MerkleHashAllClaims() *chainhash.Hash
Flush() error
}
type RamTrie struct {
collapsedTrie
bufs *sync.Pool
}
func NewRamTrie() *RamTrie {
return &RamTrie{
bufs: &sync.Pool{
New: func() interface{} {
return new(bytes.Buffer)
},
},
collapsedTrie: collapsedTrie{Root: &collapsedVertex{merkleHash: EmptyTrieHash}},
}
}
var ErrFullRebuildRequired = errors.New("a full rebuild is required")
func (rt *RamTrie) SetRoot(h *chainhash.Hash) error {
if rt.Root.merkleHash.IsEqual(h) {
runtime.GC()
return nil
}
// should technically clear the old trie first, but this is abused for partial rebuilds so don't
return ErrFullRebuildRequired
}
func (rt *RamTrie) Update(name []byte, h *chainhash.Hash, _ bool) {
if h == nil {
rt.Erase(name)
} else {
_, n := rt.InsertOrFind(name)
n.claimHash = h
}
}
func (rt *RamTrie) MerkleHash() *chainhash.Hash {
if h := rt.merkleHash(rt.Root); h == nil {
return EmptyTrieHash
}
return rt.Root.merkleHash
}
func (rt *RamTrie) merkleHash(v *collapsedVertex) *chainhash.Hash {
if v.merkleHash != nil {
return v.merkleHash
}
b := rt.bufs.Get().(*bytes.Buffer)
defer rt.bufs.Put(b)
b.Reset()
for _, ch := range v.children {
h := rt.merkleHash(ch) // h is a pointer; don't destroy its data
b.WriteByte(ch.key[0]) // nolint : errchk
b.Write(rt.completeHash(h, ch.key)) // nolint : errchk
}
if v.claimHash != nil {
b.Write(v.claimHash[:])
}
if b.Len() > 0 {
h := chainhash.DoubleHashH(b.Bytes())
v.merkleHash = &h
}
return v.merkleHash
}
func (rt *RamTrie) completeHash(h *chainhash.Hash, childKey KeyType) []byte {
var data [chainhash.HashSize + 1]byte
copy(data[1:], h[:])
for i := len(childKey) - 1; i > 0; i-- {
data[0] = childKey[i]
copy(data[1:], chainhash.DoubleHashB(data[:]))
}
return data[1:]
}
func (rt *RamTrie) MerkleHashAllClaims() *chainhash.Hash {
if h := rt.merkleHashAllClaims(rt.Root); h == nil {
return EmptyTrieHash
}
return rt.Root.merkleHash
}
func (rt *RamTrie) merkleHashAllClaims(v *collapsedVertex) *chainhash.Hash {
if v.merkleHash != nil {
return v.merkleHash
}
childHashes := make([]*chainhash.Hash, 0, len(v.children))
for _, ch := range v.children {
h := rt.merkleHashAllClaims(ch)
childHashes = append(childHashes, h)
}
claimHash := NoClaimsHash
if v.claimHash != nil {
claimHash = v.claimHash
} else if len(childHashes) == 0 {
return nil
}
childHash := NoChildrenHash
if len(childHashes) > 0 {
// this shouldn't be referencing node; where else can we put this merkle root func?
childHash = node.ComputeMerkleRoot(childHashes)
}
v.merkleHash = node.HashMerkleBranches(childHash, claimHash)
return v.merkleHash
}
func (rt *RamTrie) Flush() error {
return nil
}