feat(pkg/trie): Read-only lazy loading TrieDB (#3855)

pkg/trie/inmemory/in_memory.go

@@ -100,7 +100,7 @@ func (t *InMemoryTrie) Snapshot() (newTrie *InMemoryTrie) {
 	}
 }
 
-// handleTrackedDeltas sets the pending deleted node hashes in
+// HandleTrackedDeltas sets the pending deleted node hashes in
 // the trie deltas tracker if and only if success is true.
 func (t *InMemoryTrie) HandleTrackedDeltas(success bool, pendingDeltas tracking.Getter) {
 	if !success || t.generation == 0 {

pkg/trie/inmemory/mocks_generate_test.go

@@ -0,0 +1,6 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package inmemory
+
+//go:generate mockgen -destination=db_getter_mocks_test.go -package=$GOPACKAGE github.com/ChainSafe/gossamer/pkg/trie/db DBGetter

pkg/trie/trie.go

@@ -13,17 +13,23 @@ import (
 // EmptyHash is the empty trie hash.
 var EmptyHash = common.MustBlake2bHash([]byte{0})
 
-type ChildTrieSupport interface {
+type ChildTriesRead interface {
 	GetChild(keyToChild []byte) (Trie, error)
 	GetFromChild(keyToChild, key []byte) ([]byte, error)
 	GetChildTries() map[common.Hash]Trie
+}
+
+type ChildTriesWrite interface {
 	PutIntoChild(keyToChild, key, value []byte) error
 	DeleteChild(keyToChild []byte) (err error)
 	ClearFromChild(keyToChild, key []byte) error
 }
 
-type KVStore interface {
+type KVStoreRead interface {
 	Get(key []byte) []byte
+}
+
+type KVStoreWrite interface {
 	Put(key, value []byte) error
 	Delete(key []byte) error
 }
@@ -33,8 +39,11 @@ type TrieIterator interface {
 	NextKey(key []byte) []byte
 }
 
-type PrefixTrie interface {
+type PrefixTrieRead interface {
 	GetKeysWithPrefix(prefix []byte) (keysLE [][]byte)
+}
+
+type PrefixTrieWrite interface {
 	ClearPrefix(prefix []byte) (err error)
 	ClearPrefixLimit(prefix []byte, limit uint32) (
 		deleted uint32, allDeleted bool, err error)
@@ -54,13 +63,21 @@ type Hashable interface {
 	Hash() (common.Hash, error)
 }
 
-type Trie interface {
-	PrefixTrie
-	KVStore
+type TrieRead interface {
+	fmt.Stringer
+
+	KVStoreRead
 	Hashable
-	ChildTrieSupport
+	ChildTriesRead
+	PrefixTrieRead
 	TrieIterator
-	TrieDeltas
+}
+
+type Trie interface {
+	TrieRead
+	ChildTriesWrite
+	PrefixTrieWrite
+	KVStoreWrite
 	Versioned
-	fmt.Stringer
+	TrieDeltas
 }

pkg/trie/triedb/child_tries.go

@@ -0,0 +1,21 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package triedb
+
+import (
+	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/pkg/trie"
+)
+
+func (t *TrieDB) GetChild(keyToChild []byte) (trie.Trie, error) {
+	panic("not implemented yet")
+}
+
+func (t *TrieDB) GetFromChild(keyToChild, key []byte) ([]byte, error) {
+	panic("not implemented yet")
+}
+
+func (t *TrieDB) GetChildTries() map[common.Hash]trie.Trie {
+	panic("not implemented yet")
+}

pkg/trie/triedb/codec/decode.go

@@ -0,0 +1,166 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package codec
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+
+	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/pkg/scale"
+)
+
+var (
+	ErrDecodeHashedStorageValue  = errors.New("cannot decode hashed storage value")
+	ErrDecodeHashedValueTooShort = errors.New("hashed storage value too short")
+	ErrReadChildrenBitmap        = errors.New("cannot read children bitmap")
+	// ErrDecodeChildHash is defined since no sentinel error is defined
+	// in the scale package.
+	ErrDecodeChildHash = errors.New("cannot decode child hash")
+	// ErrDecodeStorageValue is defined since no sentinel error is defined
+	// in the scale package.
+	ErrDecodeStorageValue = errors.New("cannot decode storage value")
+)
+
+const hashLength = common.HashLength
+
+// Decode decodes a node from a reader.
+// The encoding format is documented in the README.md
+// of this package, and specified in the Polkadot spec at
+// https://spec.polkadot.network/chap-state#defn-node-header
+func Decode(reader io.Reader) (n Node, err error) {
+	variant, partialKeyLength, err := decodeHeader(reader)
+	if err != nil {
+		return nil, fmt.Errorf("decoding header: %w", err)
+	}
+
+	if variant == emptyVariant {
+		return Empty{}, nil
+	}
+
+	partialKey, err := decodeKey(reader, partialKeyLength)
+	if err != nil {
+		return nil, fmt.Errorf("cannot decode key: %w", err)
+	}
+
+	switch variant {
+	case leafVariant, leafWithHashedValueVariant:
+		n, err = decodeLeaf(reader, variant, partialKey)
+		if err != nil {
+			return nil, fmt.Errorf("cannot decode leaf: %w", err)
+		}
+		return n, nil
+	case branchVariant, branchWithValueVariant, branchWithHashedValueVariant:
+		n, err = decodeBranch(reader, variant, partialKey)
+		if err != nil {
+			return nil, fmt.Errorf("cannot decode branch: %w", err)
+		}
+		return n, nil
+	default:
+		// this is a programming error, an unknown node variant should be caught by decodeHeader.
+		panic(fmt.Sprintf("not implemented for node variant %08b", variant))
+	}
+}
+
+// decodeBranch reads from a reader and decodes to a node branch.
+// Note that we are not decoding the children nodes.
+func decodeBranch(reader io.Reader, variant variant, partialKey []byte) (
+	node Branch, err error) {
+	node = Branch{
+		PartialKey: partialKey,
+	}
+
+	var childrenBitmap uint16
+	err = binary.Read(reader, binary.LittleEndian, &childrenBitmap)
+	if err != nil {
+		return Branch{}, fmt.Errorf("%w: %s", ErrReadChildrenBitmap, err)
+	}
+
+	sd := scale.NewDecoder(reader)
+
+	switch variant {
+	case branchWithValueVariant:
+		valueBytes := make([]byte, 0)
+		err := sd.Decode(&valueBytes)
+		if err != nil {
+			return Branch{}, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
+		}
+
+		node.Value = NewInlineValue(valueBytes)
+	case branchWithHashedValueVariant:
+		hashedValue, err := decodeHashedValue(reader)
+		if err != nil {
+			return Branch{}, err
+		}
+		node.Value = NewHashedValue(hashedValue)
+	default:
+		// Do nothing, branch without value
+	}
+
+	for i := 0; i < ChildrenCapacity; i++ {
+		// Skip this index if we don't have a child here
+		if (childrenBitmap>>i)&1 != 1 {
+			continue
+		}
+
+		var hash []byte
+		err := sd.Decode(&hash)
+		if err != nil {
+			return Branch{}, fmt.Errorf("%w: at index %d: %s",
+				ErrDecodeChildHash, i, err)
+		}
+
+		if len(hash) < hashLength {
+			node.Children[i] = NewInlineNode(hash)
+		} else {
+			node.Children[i] = NewHashedNode(hash)
+		}
+	}
+
+	return node, nil
+}
+
+// decodeLeaf reads from a reader and decodes to a leaf node.
+func decodeLeaf(reader io.Reader, variant variant, partialKey []byte) (node Leaf, err error) {
+	node = Leaf{
+		PartialKey: partialKey,
+	}
+
+	sd := scale.NewDecoder(reader)
+
+	if variant == leafWithHashedValueVariant {
+		hashedValue, err := decodeHashedValue(reader)
+		if err != nil {
+			return Leaf{}, err
+		}
+
+		node.Value = NewHashedValue(hashedValue)
+		return node, nil
+	}
+
+	valueBytes := make([]byte, 0)
+	err = sd.Decode(&valueBytes)
+	if err != nil {
+		return Leaf{}, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
+	}
+
+	node.Value = NewInlineValue(valueBytes)
+
+	return node, nil
+}
+
+func decodeHashedValue(reader io.Reader) ([]byte, error) {
+	buffer := make([]byte, hashLength)
+	n, err := reader.Read(buffer)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
+	}
+	if n < hashLength {
+		return nil, fmt.Errorf("%w: expected %d, got: %d", ErrDecodeHashedValueTooShort, hashLength, n)
+	}
+
+	return buffer, nil
+}