Skip to content

Commit

Permalink
chore(lib/trie): insert inserts value byte slice instead of node (#…
Browse files Browse the repository at this point in the history 
…2296)
  • Loading branch information
@qdm12
qdm12 committed Feb 18, 2022
1 parent 9ac6642 commit 0eacba3
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 87 deletions.
61 changes: 33 additions & 28 deletions lib/trie/trie.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -306,22 +306,19 @@ func (t *Trie) Put(keyLE, value []byte) {
}

func (t *Trie) put(key, value []byte) {
nodeToInsert := &node.Leaf{
Value: value,
Generation: t.generation,
Dirty: true,
}
t.root = t.insert(t.root, key, nodeToInsert)
t.root = t.insert(t.root, key, value)
}

// insert attempts to insert a key with value into the trie
func (t *Trie) insert(parent Node, key []byte, value Node) (newParent Node) {
// TODO change value node to be value []byte?
value.SetGeneration(t.generation) // just in case it's not set by the caller.

// insert inserts a value in the trie at the key specified.
// It may create one or more new nodes or update an existing node.
func (t *Trie) insert(parent Node, key, value []byte) (newParent Node) {
if parent == nil {
value.SetKey(key)
return value
return &node.Leaf{
Key: key,
Value: value,
Generation: t.generation,
Dirty: true,
}
}

// TODO ensure all values have dirty set to true
Expand All @@ -340,13 +337,12 @@ func (t *Trie) insert(parent Node, key []byte, value Node) (newParent Node) {
}
}

func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key []byte,
value Node) (newParent Node) {
newValue := value.(*node.Leaf).Value

func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key,
value []byte) (newParent Node) {
if bytes.Equal(parentLeaf.Key, key) {
if !bytes.Equal(newValue, parentLeaf.Value) {
parentLeaf.Value = newValue
if !bytes.Equal(value, parentLeaf.Value) {
parentLeaf.Value = value
parentLeaf.Generation = t.generation
parentLeaf.SetDirty(true)
}
return parentLeaf
Expand All @@ -364,7 +360,7 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key []byte,

if len(key) == commonPrefixLength {
// key is included in parent leaf key
newBranchParent.Value = newValue
newBranchParent.Value = value

if len(key) < len(parentLeafKey) {
// Move the current leaf parent as a child to the new branch.
Expand All @@ -377,8 +373,6 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key []byte,
return newBranchParent
}

value.SetKey(key[commonPrefixLength+1:])

if len(parentLeaf.Key) == commonPrefixLength {
// the key of the parent leaf is at this new branch
newBranchParent.Value = parentLeaf.Value
Expand All @@ -390,15 +384,21 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key []byte,
newBranchParent.Children[childIndex] = parentLeaf
}
childIndex := key[commonPrefixLength]
newBranchParent.Children[childIndex] = value
newBranchParent.Children[childIndex] = &node.Leaf{
Key: key[commonPrefixLength+1:],
Value: value,
Generation: t.generation,
Dirty: true,
}

return newBranchParent
}

func (t *Trie) insertInBranch(parentBranch *node.Branch, key []byte, value Node) (newParent Node) {
func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (newParent Node) {
if bytes.Equal(key, parentBranch.Key) {
parentBranch.SetDirty(true)
parentBranch.Value = value.GetValue()
parentBranch.Generation = t.generation
parentBranch.Value = value
return parentBranch
}

Expand All @@ -412,7 +412,7 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key []byte, value Node)
if child == nil {
child = &node.Leaf{
Key: remainingKey,
Value: value.GetValue(),
Value: value,
Generation: t.generation,
Dirty: true,
}
Expand All @@ -423,6 +423,7 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key []byte, value Node)

parentBranch.Children[childIndex] = child
parentBranch.SetDirty(true)
parentBranch.Generation = t.generation
return parentBranch
}

Expand All @@ -438,10 +439,14 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key []byte, value Node)

oldParentIndex := parentBranch.Key[commonPrefixLength]
remainingOldParentKey := parentBranch.Key[commonPrefixLength+1:]
newParentBranch.Children[oldParentIndex] = t.insert(nil, remainingOldParentKey, parentBranch)

parentBranch.Dirty = true
parentBranch.Key = remainingOldParentKey
parentBranch.Generation = t.generation
newParentBranch.Children[oldParentIndex] = parentBranch

if len(key) <= commonPrefixLength {
newParentBranch.Value = value.(*node.Leaf).Value
newParentBranch.Value = value
} else {
childIndex := key[commonPrefixLength]
remainingKey := key[commonPrefixLength+1:]
Expand Down
96 changes: 37 additions & 59 deletions lib/trie/trie_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1020,18 +1020,20 @@ func Test_Trie_insert(t *testing.T) {
trie Trie
parent Node
key []byte
value Node
value []byte
newNode Node
}{
"nil parent": {
trie: Trie{
generation: 1,
},
key: []byte{1},
value: &node.Leaf{},
value: []byte("leaf"),
newNode: &node.Leaf{
Key: []byte{1},
Value: []byte("leaf"),
Generation: 1,
Dirty: true,
},
},
"branch parent": {
Expand All @@ -1046,10 +1048,8 @@ func Test_Trie_insert(t *testing.T) {
&node.Leaf{Key: []byte{2}},
},
},
key: []byte{1, 0},
value: &node.Leaf{
Value: []byte("leaf"),
},
key: []byte{1, 0},
value: []byte("leaf"),
newNode: &node.Branch{
Key: []byte{1},
Value: []byte("branch"),
Expand All @@ -1074,10 +1074,8 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{1},
Value: []byte("original leaf"),
},
key: []byte{1},
value: &node.Leaf{
Value: []byte("new leaf"),
},
key: []byte{1},
value: []byte("new leaf"),
newNode: &node.Leaf{
Key: []byte{1},
Value: []byte("new leaf"),
Expand All @@ -1093,10 +1091,8 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{1},
Value: []byte("same"),
},
key: []byte{1},
value: &node.Leaf{
Value: []byte("same"),
},
key: []byte{1},
value: []byte("same"),
newNode: &node.Leaf{
Key: []byte{1},
Value: []byte("same"),
Expand All @@ -1111,10 +1107,8 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{1},
Value: []byte("original leaf"),
},
key: []byte{1, 0},
value: &node.Leaf{
Value: []byte("leaf"),
},
key: []byte{1, 0},
value: []byte("leaf"),
newNode: &node.Branch{
Key: []byte{1},
Value: []byte("original leaf"),
Expand All @@ -1125,6 +1119,7 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{},
Value: []byte("leaf"),
Generation: 1,
Dirty: true,
},
},
},
Expand All @@ -1137,10 +1132,8 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{1, 2},
Value: []byte("original leaf"),
},
key: []byte{2, 3},
value: &node.Leaf{
Value: []byte("leaf"),
},
key: []byte{2, 3},
value: []byte("leaf"),
newNode: &node.Branch{
Key: []byte{},
Dirty: true,
Expand All @@ -1157,6 +1150,7 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{3},
Value: []byte("leaf"),
Generation: 1,
Dirty: true,
},
},
},
Expand All @@ -1168,10 +1162,8 @@ func Test_Trie_insert(t *testing.T) {
parent: &node.Leaf{
Key: []byte{1},
},
key: []byte{1},
value: &node.Leaf{
Value: []byte("leaf"),
},
key: []byte{1},
value: []byte("leaf"),
newNode: &node.Leaf{
Key: []byte{1},
Value: []byte("leaf"),
Expand All @@ -1186,10 +1178,8 @@ func Test_Trie_insert(t *testing.T) {
parent: &node.Leaf{
Key: []byte{1, 2},
},
key: []byte{1},
value: &node.Leaf{
Value: []byte("leaf"),
},
key: []byte{1},
value: []byte("leaf"),
newNode: &node.Branch{
Key: []byte{1},
Value: []byte("leaf"),
Expand Down Expand Up @@ -1229,7 +1219,7 @@ func Test_Trie_insertInBranch(t *testing.T) {
testCases := map[string]struct {
parent *node.Branch
key []byte
value Node
value []byte
newNode Node
}{
"update with branch": {
Expand All @@ -1240,10 +1230,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{Key: []byte{1}},
},
},
key: []byte{2},
value: &node.Branch{
Value: []byte("new"),
},
key: []byte{2},
value: []byte("new"),
newNode: &node.Branch{
Key: []byte{2},
Value: []byte("new"),
Expand All @@ -1261,10 +1249,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{Key: []byte{1}},
},
},
key: []byte{2},
value: &node.Leaf{
Value: []byte("new"),
},
key: []byte{2},
value: []byte("new"),
newNode: &node.Branch{
Key: []byte{2},
Value: []byte("new"),
Expand All @@ -1282,10 +1268,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{Key: []byte{1}},
},
},
key: []byte{2, 3, 4, 5},
value: &node.Leaf{
Value: []byte{6},
},
key: []byte{2, 3, 4, 5},
value: []byte{6},
newNode: &node.Branch{
Key: []byte{2},
Value: []byte{5},
Expand Down Expand Up @@ -1315,10 +1299,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
},
},
},
key: []byte{2, 3, 4, 5, 6},
value: &node.Leaf{
Value: []byte{6},
},
key: []byte{2, 3, 4, 5, 6},
value: []byte{6},
newNode: &node.Branch{
Key: []byte{2},
Value: []byte{5},
Expand Down Expand Up @@ -1349,10 +1331,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{Key: []byte{1}},
},
},
key: []byte{2, 4, 5, 6},
value: &node.Leaf{
Value: []byte{6},
},
key: []byte{2, 4, 5, 6},
value: []byte{6},
newNode: &node.Branch{
Key: []byte{2},
Dirty: true,
Expand All @@ -1369,6 +1349,7 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{
Key: []byte{5, 6},
Value: []byte{6},
Dirty: true,
},
},
},
Expand All @@ -1381,10 +1362,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{Key: []byte{1}},
},
},
key: []byte{3},
value: &node.Leaf{
Value: []byte{6},
},
key: []byte{3},
value: []byte{6},
newNode: &node.Branch{
Key: []byte{},
Dirty: true,
Expand All @@ -1401,6 +1380,7 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{
Key: []byte{},
Value: []byte{6},
Dirty: true,
},
},
},
Expand All @@ -1413,10 +1393,8 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{Key: []byte{1}},
},
},
key: []byte{},
value: &node.Leaf{
Value: []byte{6},
},
key: []byte{},
value: []byte{6},
newNode: &node.Branch{
Key: []byte{},
Value: []byte{6},
Expand Down

0 comments on commit 0eacba3

Please sign in to comment.