Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

trie: fix issue insert wrong path in stack trie and remove the offset… #604

Merged
merged 1 commit into from
Oct 17, 2024
Merged

trie: fix issue insert wrong path in stack trie and remove the offset… #604

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 37 additions & 44 deletions trie/stacktrie.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,12 @@ func returnToPool(st *StackTrie) {
// in order. Once it determines that a subtree will no longer be inserted
// into, it will hash it and free up the memory it uses.
type StackTrie struct {
owner common.Hash // the owner of the trie
nodeType uint8 // node type (as in branch, ext, leaf)
val []byte // value contained by this node if it's a leaf
key []byte // key chunk covered by this (full|ext) node
keyOffset int // offset of the key chunk inside a full key
children [16]*StackTrie // list of children (for fullnodes and exts)
writeFn NodeWriteFunc // function for commiting nodes, can be nil
owner common.Hash // the owner of the trie
nodeType uint8 // node type (as in branch, ext, leaf)
val []byte // value contained by this node if it's a leaf
key []byte // key chunk covered by this (full|ext) node
children [16]*StackTrie // list of children (for fullnodes and exts)
writeFn NodeWriteFunc // function for commiting nodes, can be nil
}

// NewStackTrie allocates and initializes an empty trie.
Expand Down Expand Up @@ -105,17 +104,15 @@ func (st *StackTrie) MarshalBinary() (data []byte, err error) {
w = bufio.NewWriter(&b)
)
if err := gob.NewEncoder(w).Encode(struct {
Owner common.Hash
NodeType uint8
Val []byte
Key []byte
KeyOffset uint8
Owner common.Hash
NodeType uint8
Val []byte
Key []byte
}{
st.owner,
st.nodeType,
st.val,
st.key,
uint8(st.keyOffset),
}); err != nil {
return nil, err
}
Expand Down Expand Up @@ -143,18 +140,16 @@ func (st *StackTrie) UnmarshalBinary(data []byte) error {

func (st *StackTrie) unmarshalBinary(r io.Reader) error {
var dec struct {
Owner common.Hash
NodeType uint8
Val []byte
Key []byte
KeyOffset uint8
Owner common.Hash
NodeType uint8
Val []byte
Key []byte
}
gob.NewDecoder(r).Decode(&dec)
st.owner = dec.Owner
st.nodeType = dec.NodeType
st.val = dec.Val
st.key = dec.Key
st.keyOffset = int(dec.KeyOffset)

var hasChild = make([]byte, 1)
for i := range st.children {
Expand All @@ -179,20 +174,18 @@ func (st *StackTrie) setWriteFunc(writeFn NodeWriteFunc) {
}
}

func newLeaf(owner common.Hash, ko int, key, val []byte, writeFn NodeWriteFunc) *StackTrie {
func newLeaf(owner common.Hash, key, val []byte, writeFn NodeWriteFunc) *StackTrie {
st := stackTrieFromPool(writeFn, owner)
st.nodeType = leafNode
st.keyOffset = ko
st.key = append(st.key, key[ko:]...)
st.key = append(st.key, key...)
st.val = val
return st
}

func newExt(owner common.Hash, ko int, key []byte, child *StackTrie, writeFn NodeWriteFunc) *StackTrie {
func newExt(owner common.Hash, key []byte, child *StackTrie, writeFn NodeWriteFunc) *StackTrie {
st := stackTrieFromPool(writeFn, owner)
st.nodeType = extNode
st.keyOffset = ko
st.key = append(st.key, key[ko:]...)
st.key = append(st.key, key...)
st.children[0] = child
return st
}
Expand Down Expand Up @@ -231,25 +224,26 @@ func (st *StackTrie) Reset() {
st.children[i] = nil
}
st.nodeType = emptyNode
st.keyOffset = 0
}

// Helper function that, given a full key, determines the index
// at which the chunk pointed by st.keyOffset is different from
// the same chunk in the full key.
func (st *StackTrie) getDiffIndex(key []byte) int {
diffindex := 0
for ; diffindex < len(st.key) && st.key[diffindex] == key[st.keyOffset+diffindex]; diffindex++ {
for idx, nibble := range st.key {
if nibble != key[idx] {
return idx
}
}
return diffindex
return len(st.key)
}

// Helper function to that inserts a (key, value) pair into
// the trie. Adding the prefix when inserting too.
func (st *StackTrie) insert(key, value []byte, prefix []byte) {
func (st *StackTrie) insert(key, value, prefix []byte) {
switch st.nodeType {
case branchNode: /* Branch */
idx := int(key[st.keyOffset])
idx := int(key[0])
// Unresolve elder siblings
for i := idx - 1; i >= 0; i-- {
if st.children[i] != nil {
Expand All @@ -261,10 +255,11 @@ func (st *StackTrie) insert(key, value []byte, prefix []byte) {
}
// Add new child
if st.children[idx] == nil {
st.children[idx] = stackTrieFromPool(st.writeFn, st.owner)
st.children[idx].keyOffset = st.keyOffset + 1
st.children[idx] = newLeaf(st.owner, key[1:], value, st.writeFn)
} else {
st.children[idx].insert(key[1:], value, append(prefix, key[0]))
}
st.children[idx].insert(key, value, append(prefix, key[st.keyOffset]))

case extNode: /* Ext */
// Compare both key chunks and see where they differ
diffidx := st.getDiffIndex(key)
Expand All @@ -277,7 +272,7 @@ func (st *StackTrie) insert(key, value []byte, prefix []byte) {
if diffidx == len(st.key) {
// Ext key and key segment are identical, recurse into
// the child node.
st.children[0].insert(key, value, append(prefix, key[:diffidx]...))
st.children[0].insert(key[diffidx:], value, append(prefix, key[:diffidx]...))
return
}
// Save the original part. Depending if the break is
Expand All @@ -289,7 +284,7 @@ func (st *StackTrie) insert(key, value []byte, prefix []byte) {
// Break on the non-last byte, insert an intermediate
// extension. The path prefix of the newly-inserted
// extension should also contain the different byte.
n = newExt(st.owner, diffidx+1, st.key, st.children[0], st.writeFn)
n = newExt(st.owner, st.key[diffidx+1:], st.children[0], st.writeFn)
n.hash(append(prefix, st.key[:diffidx+1]...))
} else {
// an extension node: reuse the current node.
Expand All @@ -313,15 +308,14 @@ func (st *StackTrie) insert(key, value []byte, prefix []byte) {
// node.
st.children[0] = stackTrieFromPool(st.writeFn, st.owner)
st.children[0].nodeType = branchNode
st.children[0].keyOffset = st.keyOffset + diffidx
p = st.children[0]
}
// Create a leaf for the inserted part
o := newLeaf(st.owner, st.keyOffset+diffidx+1, key, value, st.writeFn)
o := newLeaf(st.owner, key[diffidx+1:], value, st.writeFn)

// Insert both child leaves where they belong:
origIdx := st.key[diffidx]
newIdx := key[diffidx+st.keyOffset]
newIdx := key[diffidx]
p.children[origIdx] = n
p.children[newIdx] = o
st.key = st.key[:diffidx]
Expand Down Expand Up @@ -355,7 +349,6 @@ func (st *StackTrie) insert(key, value []byte, prefix []byte) {
st.nodeType = extNode
st.children[0] = NewStackTrieWithOwner(st.writeFn, st.owner)
st.children[0].nodeType = branchNode
st.children[0].keyOffset = st.keyOffset + diffidx
p = st.children[0]
}

Expand All @@ -364,19 +357,19 @@ func (st *StackTrie) insert(key, value []byte, prefix []byte) {
// The child leave will be hashed directly in order to
// free up some memory.
origIdx := st.key[diffidx]
p.children[origIdx] = newLeaf(st.owner, diffidx+1, st.key, st.val, st.writeFn)
p.children[origIdx] = newLeaf(st.owner, st.key[diffidx+1:], st.val, st.writeFn)
p.children[origIdx].hash(append(prefix, st.key[:diffidx+1]...))

newIdx := key[diffidx+st.keyOffset]
p.children[newIdx] = newLeaf(st.owner, p.keyOffset+1, key, value, st.writeFn)
newIdx := key[diffidx]
p.children[newIdx] = newLeaf(st.owner, key[diffidx+1:], value, st.writeFn)

// Finally, cut off the key part that has been passed
// over to the children.
st.key = st.key[:diffidx]
st.val = nil
case emptyNode: /* Empty */
st.nodeType = leafNode
st.key = key[st.keyOffset:]
st.key = key
st.val = value
case hashedNode:
panic("trying to insert into hash")
Expand Down
Loading