mirror of
https://github.com/nspcc-dev/neo-go.git
synced 2024-11-23 03:38:35 +00:00
e833d333fe
This is not a problem in practice, as all keys are prefixed by a contract ID. However in theory it can lead to a different state root after new portion of changes thus this fix.
506 lines
12 KiB
Go
506 lines
12 KiB
Go
package mpt
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
|
|
"github.com/nspcc-dev/neo-go/pkg/core/storage"
|
|
"github.com/nspcc-dev/neo-go/pkg/io"
|
|
"github.com/nspcc-dev/neo-go/pkg/util"
|
|
)
|
|
|
|
// Trie is an MPT trie storing all key-value pairs.
|
|
type Trie struct {
|
|
Store *storage.MemCachedStore
|
|
|
|
root Node
|
|
refcountEnabled bool
|
|
refcount map[util.Uint256]*cachedNode
|
|
}
|
|
|
|
type cachedNode struct {
|
|
bytes []byte
|
|
initial int32
|
|
refcount int32
|
|
}
|
|
|
|
// ErrNotFound is returned when requested trie item is missing.
|
|
var ErrNotFound = errors.New("item not found")
|
|
|
|
// NewTrie returns new MPT trie. It accepts a MemCachedStore to decouple storage errors from logic errors
|
|
// so that all storage errors are processed during `store.Persist()` at the caller.
|
|
// This also has the benefit, that every `Put` can be considered an atomic operation.
|
|
func NewTrie(root Node, enableRefCount bool, store *storage.MemCachedStore) *Trie {
|
|
if root == nil {
|
|
root = new(HashNode)
|
|
}
|
|
|
|
return &Trie{
|
|
Store: store,
|
|
root: root,
|
|
|
|
refcountEnabled: enableRefCount,
|
|
refcount: make(map[util.Uint256]*cachedNode),
|
|
}
|
|
}
|
|
|
|
// Get returns value for the provided key in t.
|
|
func (t *Trie) Get(key []byte) ([]byte, error) {
|
|
path := toNibbles(key)
|
|
r, bs, err := t.getWithPath(t.root, path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
t.root = r
|
|
return bs, nil
|
|
}
|
|
|
|
// getWithPath returns value the provided path in a subtrie rooting in curr.
|
|
// It also returns a current node with all hash nodes along the path
|
|
// replaced to their "unhashed" counterparts.
|
|
func (t *Trie) getWithPath(curr Node, path []byte) (Node, []byte, error) {
|
|
switch n := curr.(type) {
|
|
case *LeafNode:
|
|
if len(path) == 0 {
|
|
return curr, copySlice(n.value), nil
|
|
}
|
|
case *BranchNode:
|
|
i, path := splitPath(path)
|
|
r, bs, err := t.getWithPath(n.Children[i], path)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
n.Children[i] = r
|
|
return n, bs, nil
|
|
case *HashNode:
|
|
if !n.IsEmpty() {
|
|
if r, err := t.getFromStore(n.hash); err == nil {
|
|
return t.getWithPath(r, path)
|
|
}
|
|
}
|
|
case *ExtensionNode:
|
|
if bytes.HasPrefix(path, n.key) {
|
|
r, bs, err := t.getWithPath(n.next, path[len(n.key):])
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
n.next = r
|
|
return curr, bs, err
|
|
}
|
|
default:
|
|
panic("invalid MPT node type")
|
|
}
|
|
return curr, nil, ErrNotFound
|
|
}
|
|
|
|
// Put puts key-value pair in t.
|
|
func (t *Trie) Put(key, value []byte) error {
|
|
if len(key) == 0 {
|
|
return errors.New("key is empty")
|
|
} else if len(key) > MaxKeyLength {
|
|
return errors.New("key is too big")
|
|
} else if len(value) > MaxValueLength {
|
|
return errors.New("value is too big")
|
|
}
|
|
if len(value) == 0 {
|
|
return t.Delete(key)
|
|
}
|
|
path := toNibbles(key)
|
|
n := NewLeafNode(value)
|
|
r, err := t.putIntoNode(t.root, path, n)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
t.root = r
|
|
return nil
|
|
}
|
|
|
|
// putIntoLeaf puts val to trie if current node is a Leaf.
|
|
// It returns Node if curr needs to be replaced and error if any.
|
|
func (t *Trie) putIntoLeaf(curr *LeafNode, path []byte, val Node) (Node, error) {
|
|
v := val.(*LeafNode)
|
|
if len(path) == 0 {
|
|
t.removeRef(curr.Hash(), curr.bytes)
|
|
t.addRef(val.Hash(), val.Bytes())
|
|
return v, nil
|
|
}
|
|
|
|
b := NewBranchNode()
|
|
b.Children[path[0]] = t.newSubTrie(path[1:], v, true)
|
|
b.Children[lastChild] = curr
|
|
t.addRef(b.Hash(), b.bytes)
|
|
return b, nil
|
|
}
|
|
|
|
// putIntoBranch puts val to trie if current node is a Branch.
|
|
// It returns Node if curr needs to be replaced and error if any.
|
|
func (t *Trie) putIntoBranch(curr *BranchNode, path []byte, val Node) (Node, error) {
|
|
i, path := splitPath(path)
|
|
t.removeRef(curr.Hash(), curr.bytes)
|
|
r, err := t.putIntoNode(curr.Children[i], path, val)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
curr.Children[i] = r
|
|
curr.invalidateCache()
|
|
t.addRef(curr.Hash(), curr.bytes)
|
|
return curr, nil
|
|
}
|
|
|
|
// putIntoExtension puts val to trie if current node is an Extension.
|
|
// It returns Node if curr needs to be replaced and error if any.
|
|
func (t *Trie) putIntoExtension(curr *ExtensionNode, path []byte, val Node) (Node, error) {
|
|
t.removeRef(curr.Hash(), curr.bytes)
|
|
if bytes.HasPrefix(path, curr.key) {
|
|
r, err := t.putIntoNode(curr.next, path[len(curr.key):], val)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
curr.next = r
|
|
curr.invalidateCache()
|
|
t.addRef(curr.Hash(), curr.bytes)
|
|
return curr, nil
|
|
}
|
|
|
|
pref := lcp(curr.key, path)
|
|
lp := len(pref)
|
|
keyTail := curr.key[lp:]
|
|
pathTail := path[lp:]
|
|
|
|
s1 := t.newSubTrie(keyTail[1:], curr.next, false)
|
|
b := NewBranchNode()
|
|
b.Children[keyTail[0]] = s1
|
|
|
|
i, pathTail := splitPath(pathTail)
|
|
s2 := t.newSubTrie(pathTail, val, true)
|
|
b.Children[i] = s2
|
|
|
|
t.addRef(b.Hash(), b.bytes)
|
|
if lp > 0 {
|
|
e := NewExtensionNode(copySlice(pref), b)
|
|
t.addRef(e.Hash(), e.bytes)
|
|
return e, nil
|
|
}
|
|
return b, nil
|
|
}
|
|
|
|
// putIntoHash puts val to trie if current node is a HashNode.
|
|
// It returns Node if curr needs to be replaced and error if any.
|
|
func (t *Trie) putIntoHash(curr *HashNode, path []byte, val Node) (Node, error) {
|
|
if curr.IsEmpty() {
|
|
hn := t.newSubTrie(path, val, true)
|
|
return hn, nil
|
|
}
|
|
|
|
result, err := t.getFromStore(curr.hash)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return t.putIntoNode(result, path, val)
|
|
}
|
|
|
|
// newSubTrie create new trie containing node at provided path.
|
|
func (t *Trie) newSubTrie(path []byte, val Node, newVal bool) Node {
|
|
if newVal {
|
|
t.addRef(val.Hash(), val.Bytes())
|
|
}
|
|
if len(path) == 0 {
|
|
return val
|
|
}
|
|
e := NewExtensionNode(path, val)
|
|
t.addRef(e.Hash(), e.bytes)
|
|
return e
|
|
}
|
|
|
|
// putIntoNode puts val with provided path inside curr and returns updated node.
|
|
// Reference counters are updated for both curr and returned value.
|
|
func (t *Trie) putIntoNode(curr Node, path []byte, val Node) (Node, error) {
|
|
switch n := curr.(type) {
|
|
case *LeafNode:
|
|
return t.putIntoLeaf(n, path, val)
|
|
case *BranchNode:
|
|
return t.putIntoBranch(n, path, val)
|
|
case *ExtensionNode:
|
|
return t.putIntoExtension(n, path, val)
|
|
case *HashNode:
|
|
return t.putIntoHash(n, path, val)
|
|
default:
|
|
panic("invalid MPT node type")
|
|
}
|
|
}
|
|
|
|
// Delete removes key from trie.
|
|
// It returns no error on missing key.
|
|
func (t *Trie) Delete(key []byte) error {
|
|
path := toNibbles(key)
|
|
r, err := t.deleteFromNode(t.root, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
t.root = r
|
|
return nil
|
|
}
|
|
|
|
func (t *Trie) deleteFromBranch(b *BranchNode, path []byte) (Node, error) {
|
|
i, path := splitPath(path)
|
|
h := b.Hash()
|
|
bs := b.bytes
|
|
r, err := t.deleteFromNode(b.Children[i], path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
t.removeRef(h, bs)
|
|
b.Children[i] = r
|
|
b.invalidateCache()
|
|
var count, index int
|
|
for i := range b.Children {
|
|
h, ok := b.Children[i].(*HashNode)
|
|
if !ok || !h.IsEmpty() {
|
|
index = i
|
|
count++
|
|
}
|
|
}
|
|
// count is >= 1 because branch node had at least 2 children before deletion.
|
|
if count > 1 {
|
|
t.addRef(b.Hash(), b.bytes)
|
|
return b, nil
|
|
}
|
|
c := b.Children[index]
|
|
if index == lastChild {
|
|
return c, nil
|
|
}
|
|
if h, ok := c.(*HashNode); ok {
|
|
c, err = t.getFromStore(h.Hash())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if e, ok := c.(*ExtensionNode); ok {
|
|
t.removeRef(e.Hash(), e.bytes)
|
|
e.key = append([]byte{byte(index)}, e.key...)
|
|
e.invalidateCache()
|
|
t.addRef(e.Hash(), e.bytes)
|
|
return e, nil
|
|
}
|
|
|
|
e := NewExtensionNode([]byte{byte(index)}, c)
|
|
t.addRef(e.Hash(), e.bytes)
|
|
return e, nil
|
|
}
|
|
|
|
func (t *Trie) deleteFromExtension(n *ExtensionNode, path []byte) (Node, error) {
|
|
if !bytes.HasPrefix(path, n.key) {
|
|
return n, nil
|
|
}
|
|
h := n.Hash()
|
|
bs := n.bytes
|
|
r, err := t.deleteFromNode(n.next, path[len(n.key):])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
t.removeRef(h, bs)
|
|
switch nxt := r.(type) {
|
|
case *ExtensionNode:
|
|
t.removeRef(nxt.Hash(), nxt.bytes)
|
|
n.key = append(n.key, nxt.key...)
|
|
n.next = nxt.next
|
|
case *HashNode:
|
|
if nxt.IsEmpty() {
|
|
return nxt, nil
|
|
}
|
|
default:
|
|
n.next = r
|
|
}
|
|
n.invalidateCache()
|
|
t.addRef(n.Hash(), n.bytes)
|
|
return n, nil
|
|
}
|
|
|
|
// deleteFromNode removes value with provided path from curr and returns an updated node.
|
|
// Reference counters are updated for both curr and returned value.
|
|
func (t *Trie) deleteFromNode(curr Node, path []byte) (Node, error) {
|
|
switch n := curr.(type) {
|
|
case *LeafNode:
|
|
if len(path) == 0 {
|
|
t.removeRef(curr.Hash(), curr.Bytes())
|
|
return new(HashNode), nil
|
|
}
|
|
return curr, nil
|
|
case *BranchNode:
|
|
return t.deleteFromBranch(n, path)
|
|
case *ExtensionNode:
|
|
return t.deleteFromExtension(n, path)
|
|
case *HashNode:
|
|
if n.IsEmpty() {
|
|
return n, nil
|
|
}
|
|
newNode, err := t.getFromStore(n.Hash())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return t.deleteFromNode(newNode, path)
|
|
default:
|
|
panic("invalid MPT node type")
|
|
}
|
|
}
|
|
|
|
// StateRoot returns root hash of t.
|
|
func (t *Trie) StateRoot() util.Uint256 {
|
|
if hn, ok := t.root.(*HashNode); ok && hn.IsEmpty() {
|
|
return util.Uint256{}
|
|
}
|
|
return t.root.Hash()
|
|
}
|
|
|
|
func makeStorageKey(mptKey []byte) []byte {
|
|
return append([]byte{byte(storage.DataMPT)}, mptKey...)
|
|
}
|
|
|
|
// Flush puts every node in the trie except Hash ones to the storage.
|
|
// Because we care only about block-level changes, there is no need to put every
|
|
// new node to storage. Normally, flush should be called with every StateRoot persist, i.e.
|
|
// after every block.
|
|
func (t *Trie) Flush() {
|
|
for h, node := range t.refcount {
|
|
if node.refcount != 0 {
|
|
if node.bytes == nil {
|
|
panic("item not in trie")
|
|
}
|
|
if t.refcountEnabled {
|
|
node.initial = t.updateRefCount(h)
|
|
if node.initial == 0 {
|
|
delete(t.refcount, h)
|
|
}
|
|
} else if node.refcount > 0 {
|
|
_ = t.Store.Put(makeStorageKey(h.BytesBE()), node.bytes)
|
|
}
|
|
node.refcount = 0
|
|
} else {
|
|
delete(t.refcount, h)
|
|
}
|
|
}
|
|
}
|
|
|
|
// updateRefCount should be called only when refcounting is enabled.
|
|
func (t *Trie) updateRefCount(h util.Uint256) int32 {
|
|
if !t.refcountEnabled {
|
|
panic("`updateRefCount` is called, but GC is disabled")
|
|
}
|
|
var data []byte
|
|
key := makeStorageKey(h.BytesBE())
|
|
node := t.refcount[h]
|
|
cnt := node.initial
|
|
if cnt == 0 {
|
|
// A newly created item which may be in store.
|
|
var err error
|
|
data, err = t.Store.Get(key)
|
|
if err == nil {
|
|
cnt = int32(binary.LittleEndian.Uint32(data[len(data)-4:]))
|
|
}
|
|
}
|
|
if len(data) == 0 {
|
|
data = append(node.bytes, 0, 0, 0, 0)
|
|
}
|
|
cnt += node.refcount
|
|
switch {
|
|
case cnt < 0:
|
|
// BUG: negative reference count
|
|
panic(fmt.Sprintf("negative reference count: %s new %d, upd %d", h.StringBE(), cnt, t.refcount[h]))
|
|
case cnt == 0:
|
|
_ = t.Store.Delete(key)
|
|
default:
|
|
binary.LittleEndian.PutUint32(data[len(data)-4:], uint32(cnt))
|
|
_ = t.Store.Put(key, data)
|
|
}
|
|
return cnt
|
|
}
|
|
|
|
func (t *Trie) addRef(h util.Uint256, bs []byte) {
|
|
node := t.refcount[h]
|
|
if node == nil {
|
|
t.refcount[h] = &cachedNode{
|
|
refcount: 1,
|
|
bytes: bs,
|
|
}
|
|
return
|
|
}
|
|
node.refcount++
|
|
if node.bytes == nil {
|
|
node.bytes = bs
|
|
}
|
|
}
|
|
|
|
func (t *Trie) removeRef(h util.Uint256, bs []byte) {
|
|
node := t.refcount[h]
|
|
if node == nil {
|
|
t.refcount[h] = &cachedNode{
|
|
refcount: -1,
|
|
bytes: bs,
|
|
}
|
|
return
|
|
}
|
|
node.refcount--
|
|
if node.bytes == nil {
|
|
node.bytes = bs
|
|
}
|
|
}
|
|
|
|
func (t *Trie) getFromStore(h util.Uint256) (Node, error) {
|
|
data, err := t.Store.Get(makeStorageKey(h.BytesBE()))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var n NodeObject
|
|
r := io.NewBinReaderFromBuf(data)
|
|
n.DecodeBinary(r)
|
|
if r.Err != nil {
|
|
return nil, r.Err
|
|
}
|
|
|
|
if t.refcountEnabled {
|
|
data = data[:len(data)-4]
|
|
node := t.refcount[h]
|
|
if node != nil {
|
|
node.bytes = data
|
|
node.initial = int32(r.ReadU32LE())
|
|
}
|
|
}
|
|
n.Node.(flushedNode).setCache(data, h)
|
|
return n.Node, nil
|
|
}
|
|
|
|
// Collapse compresses all nodes at depth n to the hash nodes.
|
|
// Note: this function does not perform any kind of storage flushing so
|
|
// `Flush()` should be called explicitly before invoking function.
|
|
func (t *Trie) Collapse(depth int) {
|
|
if depth < 0 {
|
|
panic("negative depth")
|
|
}
|
|
t.root = collapse(depth, t.root)
|
|
t.refcount = make(map[util.Uint256]*cachedNode)
|
|
}
|
|
|
|
func collapse(depth int, node Node) Node {
|
|
if _, ok := node.(*HashNode); ok {
|
|
return node
|
|
} else if depth == 0 {
|
|
return NewHashNode(node.Hash())
|
|
}
|
|
|
|
switch n := node.(type) {
|
|
case *BranchNode:
|
|
for i := range n.Children {
|
|
n.Children[i] = collapse(depth-1, n.Children[i])
|
|
}
|
|
case *ExtensionNode:
|
|
n.next = collapse(depth-1, n.next)
|
|
case *LeafNode:
|
|
case *HashNode:
|
|
default:
|
|
panic("invalid MPT node type")
|
|
}
|
|
return node
|
|
}
|