Initial MPT implementation (2.x)

cli/server/dump.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/nspcc-dev/neo-go/pkg/core/mpt"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 )
@@ -33,35 +34,7 @@ func toNeoStorageKey(key []byte) []byte {
 	if len(key) < util.Uint160Size {
 		panic("invalid key in storage")
 	}
-
-	var nkey []byte
-	for i := util.Uint160Size - 1; i >= 0; i-- {
-		nkey = append(nkey, key[i])
-	}
-
-	key = key[util.Uint160Size:]
-
-	index := 0
-	remain := len(key)
-	for remain >= 16 {
-		nkey = append(nkey, key[index:index+16]...)
-		nkey = append(nkey, 0)
-		index += 16
-		remain -= 16
-	}
-
-	if remain > 0 {
-		nkey = append(nkey, key[index:]...)
-	}
-
-	padding := 16 - remain
-	for i := 0; i < padding; i++ {
-		nkey = append(nkey, 0)
-	}
-
-	nkey = append(nkey, byte(padding))
-
-	return nkey
+	return mpt.ToNeoStorageKey(key)
 }
 
 // batchToMap converts batch to a map so that JSON is compatible

pkg/core/mpt/branch.go

@@ -0,0 +1,98 @@
+package mpt
+
+import (
+	"encoding/json"
+	"errors"
+
+	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+)
+
+const (
+	// childrenCount represents a number of children of a branch node.
+	childrenCount = 17
+	// lastChild is the index of the last child.
+	lastChild = childrenCount - 1
+)
+
+// BranchNode represents MPT's branch node.
+type BranchNode struct {
+	hash  util.Uint256
+	valid bool
+
+	Children [childrenCount]Node
+}
+
+var _ Node = (*BranchNode)(nil)
+
+// NewBranchNode returns new branch node.
+func NewBranchNode() *BranchNode {
+	b := new(BranchNode)
+	for i := 0; i < childrenCount; i++ {
+		b.Children[i] = new(HashNode)
+	}
+	return b
+}
+
+// Type implements Node interface.
+func (b *BranchNode) Type() NodeType { return BranchT }
+
+// Hash implements Node interface.
+func (b *BranchNode) Hash() util.Uint256 {
+	if !b.valid {
+		b.hash = hash.DoubleSha256(toBytes(b))
+		b.valid = true
+	}
+	return b.hash
+}
+
+// invalidateHash invalidates node hash.
+func (b *BranchNode) invalidateHash() {
+	b.valid = false
+}
+
+// EncodeBinary implements io.Serializable.
+func (b *BranchNode) EncodeBinary(w *io.BinWriter) {
+	for i := 0; i < childrenCount; i++ {
+		if hn, ok := b.Children[i].(*HashNode); ok {
+			hn.EncodeBinary(w)
+			continue
+		}
+		n := NewHashNode(b.Children[i].Hash())
+		n.EncodeBinary(w)
+	}
+}
+
+// DecodeBinary implements io.Serializable.
+func (b *BranchNode) DecodeBinary(r *io.BinReader) {
+	for i := 0; i < childrenCount; i++ {
+		b.Children[i] = new(HashNode)
+		b.Children[i].DecodeBinary(r)
+	}
+}
+
+// MarshalJSON implements json.Marshaler.
+func (b *BranchNode) MarshalJSON() ([]byte, error) {
+	return json.Marshal(b.Children)
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (b *BranchNode) UnmarshalJSON(data []byte) error {
+	var obj NodeObject
+	if err := obj.UnmarshalJSON(data); err != nil {
+		return err
+	} else if u, ok := obj.Node.(*BranchNode); ok {
+		*b = *u
+		return nil
+	}
+	return errors.New("expected branch node")
+}
+
+// splitPath splits path for a branch node.
+func splitPath(path []byte) (byte, []byte) {
+	if len(path) != 0 {
+		return path[0], path[1:]
+	}
+	return lastChild, path
+}

pkg/core/mpt/doc.go

@@ -0,0 +1,45 @@
+/*
+Package mpt implements MPT (Merkle-Patricia Tree).
+
+MPT stores key-value pairs and is a trie over 16-symbol alphabet. https://en.wikipedia.org/wiki/Trie
+Trie is a tree where values are stored in leafs and keys are paths from root to the leaf node.
+MPT consists of 4 type of nodes:
+- Leaf node contains only value.
+- Extension node contains both key and value.
+- Branch node contains 2 or more children.
+- Hash node is a compressed node and contains only actual node's hash.
+  The actual node must be retrieved from storage or over the network.
+
+As an example here is a trie containing 3 pairs:
+- 0x1201 -> val1
+- 0x1203 -> val2
+- 0x1224 -> val3
+- 0x12 -> val4
+
+ExtensionNode(0x0102), Next
+ _______________________|
+ |
+BranchNode [0, 1, 2, ...], Last -> Leaf(val4)
+            |     |
+            |     ExtensionNode [0x04], Next -> Leaf(val3)
+            |
+            BranchNode [0, 1, 2, 3, ...], Last -> HashNode(nil)
+                           |     |
+                           |     Leaf(val2)
+                           |
+                           Leaf(val1)
+
+There are 3 invariants that this implementation has:
+- Branch node cannot have <= 1 children
+- Extension node cannot have zero-length key
+- Extension node cannot have another Extension node in it's next field
+
+Thank to these restrictions, there is a single root hash for every set of key-value pairs
+irregardless of the order they were added/removed with.
+The actual trie structure can vary because of node -> HashNode compressing.
+
+There is also one optimization which cost us almost nothing in terms of complexity but is very beneficial:
+When we perform get/put/delete on a speficic path, every Hash node which was retreived from storage is
+replaced by its uncompressed form, so that subsequent hits of this not don't use storage.
+*/
+package mpt

pkg/core/mpt/extension.go

@@ -0,0 +1,94 @@
+package mpt
+
+import (
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+
+	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+)
+
+// MaxKeyLength is the max length of the extension node key.
+const MaxKeyLength = 1125
+
+// ExtensionNode represents MPT's extension node.
+type ExtensionNode struct {
+	hash  util.Uint256
+	valid bool
+
+	key  []byte
+	next Node
+}
+
+var _ Node = (*ExtensionNode)(nil)
+
+// NewExtensionNode returns hash node with the specified key and next node.
+// Note: because it is a part of Trie, key must be mangled, i.e. must contain only bytes with high half = 0.
+func NewExtensionNode(key []byte, next Node) *ExtensionNode {
+	return &ExtensionNode{
+		key:  key,
+		next: next,
+	}
+}
+
+// Type implements Node interface.
+func (e ExtensionNode) Type() NodeType { return ExtensionT }
+
+// Hash implements Node interface.
+func (e *ExtensionNode) Hash() util.Uint256 {
+	if !e.valid {
+		e.hash = hash.DoubleSha256(toBytes(e))
+		e.valid = true
+	}
+	return e.hash
+}
+
+// invalidateHash invalidates node hash.
+func (e *ExtensionNode) invalidateHash() {
+	e.valid = false
+}
+
+// DecodeBinary implements io.Serializable.
+func (e *ExtensionNode) DecodeBinary(r *io.BinReader) {
+	sz := r.ReadVarUint()
+	if sz > MaxKeyLength {
+		r.Err = fmt.Errorf("extension node key is too big: %d", sz)
+		return
+	}
+	e.valid = false
+	e.key = make([]byte, sz)
+	r.ReadBytes(e.key)
+	e.next = new(HashNode)
+	e.next.DecodeBinary(r)
+}
+
+// EncodeBinary implements io.Serializable.
+func (e ExtensionNode) EncodeBinary(w *io.BinWriter) {
+	w.WriteVarBytes(e.key)
+	n := NewHashNode(e.next.Hash())
+	n.EncodeBinary(w)
+}
+
+// MarshalJSON implements json.Marshaler.
+func (e *ExtensionNode) MarshalJSON() ([]byte, error) {
+	m := map[string]interface{}{
+		"key":  hex.EncodeToString(e.key),
+		"next": e.next,
+	}
+	return json.Marshal(m)
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (e *ExtensionNode) UnmarshalJSON(data []byte) error {
+	var obj NodeObject
+	if err := obj.UnmarshalJSON(data); err != nil {
+		return err
+	} else if u, ok := obj.Node.(*ExtensionNode); ok {
+		*e = *u
+		return nil
+	}
+	return errors.New("expected extension node")
+}

pkg/core/mpt/hash.go

@@ -0,0 +1,82 @@
+package mpt
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+)
+
+// HashNode represents MPT's hash node.
+type HashNode struct {
+	hash  util.Uint256
+	valid bool
+}
+
+var _ Node = (*HashNode)(nil)
+
+// NewHashNode returns hash node with the specified hash.
+func NewHashNode(h util.Uint256) *HashNode {
+	return &HashNode{
+		hash:  h,
+		valid: true,
+	}
+}
+
+// Type implements Node interface.
+func (h *HashNode) Type() NodeType { return HashT }
+
+// Hash implements Node interface.
+func (h *HashNode) Hash() util.Uint256 {
+	if !h.valid {
+		panic("can't get hash of an empty HashNode")
+	}
+	return h.hash
+}
+
+// IsEmpty returns true iff h is an empty node i.e. contains no hash.
+func (h *HashNode) IsEmpty() bool { return !h.valid }
+
+// DecodeBinary implements io.Serializable.
+func (h *HashNode) DecodeBinary(r *io.BinReader) {
+	sz := r.ReadVarUint()
+	switch sz {
+	case 0:
+		h.valid = false
+	case util.Uint256Size:
+		h.valid = true
+		r.ReadBytes(h.hash[:])
+	default:
+		r.Err = fmt.Errorf("invalid hash node size: %d", sz)
+	}
+}
+
+// EncodeBinary implements io.Serializable.
+func (h HashNode) EncodeBinary(w *io.BinWriter) {
+	if !h.valid {
+		w.WriteVarUint(0)
+		return
+	}
+	w.WriteVarBytes(h.hash[:])
+}
+
+// MarshalJSON implements json.Marshaler.
+func (h *HashNode) MarshalJSON() ([]byte, error) {
+	if !h.valid {
+		return []byte(`{}`), nil
+	}
+	return []byte(`{"hash":"` + h.hash.StringLE() + `"}`), nil
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (h *HashNode) UnmarshalJSON(data []byte) error {
+	var obj NodeObject
+	if err := obj.UnmarshalJSON(data); err != nil {
+		return err
+	} else if u, ok := obj.Node.(*HashNode); ok {
+		*h = *u
+		return nil
+	}
+	return errors.New("expected hash node")
+}