Skip to content

Commit

Permalink
move NodeOwned
Browse files Browse the repository at this point in the history
  • Loading branch information
timwu20 committed Oct 4, 2024
1 parent cb769bc commit cbb91aa
Show file tree
Hide file tree
Showing 5 changed files with 193 additions and 183 deletions.
164 changes: 0 additions & 164 deletions pkg/trie/triedb/codec/node.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
package codec

import (
"bytes"
"fmt"
"io"

Expand Down Expand Up @@ -172,166 +171,3 @@ func EncodeHeader(partialKey []byte, partialKeyLength uint, kind NodeKind, write

return nil
}

type ValueOwnedTypes[H any] interface {
ValueOwnedInline[H] | ValueOwnedNode[H]
ValueOwned
}
type ValueOwned interface {
isValueOwned()
}

type (
// Value bytes as stored in a trie node and its hash.
ValueOwnedInline[H any] struct {
Value []byte
Hash H
}
// Hash byte slice as stored in a trie node.
ValueOwnedNode[H any] struct {
Hash H
}
)

func (ValueOwnedInline[H]) isValueOwned() {}
func (ValueOwnedNode[H]) isValueOwned() {}

var (
_ ValueOwned = ValueOwnedInline[string]{}
_ ValueOwned = ValueOwnedNode[string]{}
)

func ValueOwnedFromEncodedValue[H hash.Hash, Hasher hash.Hasher[H]](encVal EncodedValue) ValueOwned {
switch encVal := encVal.(type) {
case InlineValue:
return ValueOwnedInline[H]{
Value: encVal,
Hash: (*(new(Hasher))).Hash(encVal),
}
case HashedValue[H]:
return ValueOwnedNode[H](encVal)
default:
panic("unreachable")
}
}

type NodeHandleOwnedTypes[H any] interface {
NodeHandleOwnedHash[H] | NodeHandleOwnedInline[H]
}

type NodeHandleOwned interface {
isNodeHandleOwned()
}

type (
NodeHandleOwnedHash[H any] struct {
Hash H
}
NodeHandleOwnedInline[H any] struct {
NodeOwned
}
)

func (NodeHandleOwnedHash[H]) isNodeHandleOwned() {}
func (NodeHandleOwnedInline[H]) isNodeHandleOwned() {}

var (
_ NodeHandleOwned = NodeHandleOwnedHash[string]{}
_ NodeHandleOwned = NodeHandleOwnedInline[string]{}
)

func NodeHandleOwnedFromMerkleValue[H hash.Hash, Hasher hash.Hasher[H]](mv MerkleValue) (NodeHandleOwned, error) {
switch mv := mv.(type) {
case HashedNode[H]:
return NodeHandleOwnedHash[H](mv), nil
case InlineNode:
buf := bytes.NewBuffer(mv)
node, err := Decode[H](buf)
if err != nil {
return nil, err
}
nodeOwned, err := NodeOwnedFromNode[H, Hasher](node)
if err != nil {
return nil, err
}
return NodeHandleOwnedInline[H]{nodeOwned}, nil
default:
panic("unreachable")
}
}

type NodeOwnedTypes[H any] interface {
NodeOwnedEmpty | NodeOwnedLeaf[H] | NodeOwnedBranch[H] | NodeOwnedValue[H]
NodeOwned
}
type NodeOwned interface {
isNodeOwned()
}

type (
// Null trie node; could be an empty root or an empty branch entry.
NodeOwnedEmpty struct{}
// Leaf node; has key slice and value. Value may not be empty.
NodeOwnedLeaf[H any] struct {
PartialKey nibbles.Nibbles
Value ValueOwned
}
// Branch node; has slice of child nodes (each possibly null)
// and an optional immediate node data.
NodeOwnedBranch[H any] struct {
PartialKey nibbles.Nibbles
Children [ChildrenCapacity]NodeHandleOwned // can be nil to represent no child
Value ValueOwned
}
// Node that represents a value.
//
// This variant is only constructed when working with a [`crate::TrieCache`]. It is only
// used to cache a raw value.
NodeOwnedValue[H any] struct {
Value []byte
Hash H
}
)

func (NodeOwnedEmpty) isNodeOwned() {}
func (NodeOwnedLeaf[H]) isNodeOwned() {}
func (NodeOwnedBranch[H]) isNodeOwned() {}
func (NodeOwnedValue[H]) isNodeOwned() {}

var (
_ NodeOwned = NodeOwnedEmpty{}
_ NodeOwned = NodeOwnedLeaf[string]{}
_ NodeOwned = NodeOwnedBranch[string]{}
_ NodeOwned = NodeOwnedValue[string]{}
)

func NodeOwnedFromNode[H hash.Hash, Hasher hash.Hasher[H]](n EncodedNode) (NodeOwned, error) {
switch n := n.(type) {
case Empty:
return NodeOwnedEmpty{}, nil
case Leaf:
return NodeOwnedLeaf[H]{
PartialKey: n.PartialKey,
Value: ValueOwnedFromEncodedValue[H, Hasher](n.Value),
}, nil
case Branch:
var childrenOwned [ChildrenCapacity]NodeHandleOwned
for i, child := range n.Children {
if child == nil {
continue
}
var err error
childrenOwned[i], err = NodeHandleOwnedFromMerkleValue[H, Hasher](child)
if err != nil {
return nil, err
}
}
return NodeOwnedBranch[H]{
PartialKey: n.PartialKey,
Children: childrenOwned,
Value: ValueOwnedFromEncodedValue[H, Hasher](n.Value),
}, nil
default:
panic("unreachable")
}
}
35 changes: 19 additions & 16 deletions pkg/trie/triedb/lookup.go
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ func (l *TrieLookup[H, Hasher, QueryItem]) lookupValueWithCache(fullKey []byte,
data, err := loadValueOwned[H](
// If we only have the hash cached, this can only be a value node.
// For inline nodes we cache them directly as `CachedValue::Existing`.
codec.ValueOwned(codec.ValueOwnedNode[H]{Hash: cachedVal.Hash}),
ValueOwned(ValueOwnedNode[H]{Hash: cachedVal.Hash}),
keyNibbles, // nibble_key.original_data_as_prefix(),
fullKey,
cache,
Expand Down Expand Up @@ -266,7 +266,7 @@ func (l *TrieLookup[H, Hasher, QueryItem]) lookupValueWithCache(fullKey []byte,
}

type loadValueFunc[H hash.Hash, R any] func(
v codec.ValueOwned,
v ValueOwned,
prefix nibbles.Nibbles,
fullKey []byte,
cache TrieCache[H],
Expand All @@ -285,9 +285,12 @@ func lookupValueWithCacheInternal[H hash.Hash, Hasher hash.Hasher[H], R, QueryIt
hash := l.hash
var keyNibbles uint

// TODO: remove this
_ = partial

var depth uint
for {
node, err := cache.GetOrInsertNode(hash, func() (codec.NodeOwned, error) {
node, err := cache.GetOrInsertNode(hash, func() (NodeOwned, error) {
prefixedKey := append(nibbleKey.Mid(keyNibbles).Left().JoinedBytes(), hash.Bytes()...)
nodeData, err := l.db.Get(prefixedKey)
if err != nil {
Expand All @@ -303,7 +306,7 @@ func lookupValueWithCacheInternal[H hash.Hash, Hasher hash.Hasher[H], R, QueryIt
return nil, err
}

return codec.NodeOwnedFromNode[H, Hasher](decoded)
return NodeOwnedFromNode[H, Hasher](decoded)
})
if err != nil {
return nil, err
Expand All @@ -315,9 +318,9 @@ func lookupValueWithCacheInternal[H hash.Hash, Hasher hash.Hasher[H], R, QueryIt
// this loop iterates through all inline children (usually max 1)
// without incrementing the depth.
for {
var nextNode codec.NodeHandleOwned
var nextNode NodeHandleOwned
switch node := node.(type) {
case codec.NodeOwnedLeaf[H]:
case NodeOwnedLeaf[H]:
if partial.Equal(node.PartialKey) {
value := node.Value
r, err := loadValue(value, nibbleKey, fullKey, cache, l.db, l.recorder)
Expand All @@ -329,7 +332,7 @@ func lookupValueWithCacheInternal[H hash.Hash, Hasher hash.Hasher[H], R, QueryIt
l.recordAccess(NonExistingNodeAccess{fullKey})
return nil, nil
}
case codec.NodeOwnedBranch[H]:
case NodeOwnedBranch[H]:
if partial.Len() == 0 {
value := node.Value
r, err := loadValue(value, nibbleKey, fullKey, cache, l.db, l.recorder)
Expand All @@ -348,18 +351,18 @@ func lookupValueWithCacheInternal[H hash.Hash, Hasher hash.Hasher[H], R, QueryIt
return nil, nil
}
}
case codec.NodeOwnedEmpty:
case NodeOwnedEmpty:
l.recordAccess(NonExistingNodeAccess{FullKey: fullKey})
default:
panic("unreachable")
}

// check if new node data is inline or hash.
switch nextNode := nextNode.(type) {
case codec.NodeHandleOwnedHash[H]:
case NodeHandleOwnedHash[H]:
hash = nextNode.Hash
break inlineLoop
case codec.NodeHandleOwnedInline[H]:
case NodeHandleOwnedInline[H]:
node = nextNode.NodeOwned
default:
panic("unreachable")
Expand Down Expand Up @@ -392,38 +395,38 @@ func (vh *valueHash[H]) CachedValue() CachedValue {
//
// Returns the bytes representing the value and its hash.
func loadValueOwned[H hash.Hash](
v codec.ValueOwned,
v ValueOwned,
prefix nibbles.Nibbles,
fullKey []byte,
cache TrieCache[H],
db db.DBGetter,
recorder TrieRecorder) (valueHash[H], error) {

switch v := v.(type) {
case codec.ValueOwnedInline[H]:
case ValueOwnedInline[H]:
if recorder != nil {
recorder.Record(InlineValueAccess{fullKey})
}
return valueHash[H]{
Value: v.Value,
Hash: v.Hash,
}, nil
case codec.ValueOwnedNode[H]:
node, err := cache.GetOrInsertNode(v.Hash, func() (codec.NodeOwned, error) {
case ValueOwnedNode[H]:
node, err := cache.GetOrInsertNode(v.Hash, func() (NodeOwned, error) {
prefixedKey := append(prefix.Left().JoinedBytes(), v.Hash.Bytes()...)
val, err := db.Get(prefixedKey)
if err != nil {
return nil, err
}
return codec.NodeOwnedValue[H]{Value: val, Hash: v.Hash}, nil
return NodeOwnedValue[H]{Value: val, Hash: v.Hash}, nil
})
if err != nil {
return valueHash[H]{}, err
}

var value []byte
switch node := node.(type) {
case codec.NodeOwnedValue[H]:
case NodeOwnedValue[H]:
value = node.Value
default:
panic("we are caching a `NodeOwnedValue` for a value node hash and this cached node has always data attached")
Expand Down
Loading

0 comments on commit cbb91aa

Please sign in to comment.