Skip to content

Commit

Permalink
chore(lib/trie): insert inserts value byte slice
Browse files Browse the repository at this point in the history
  • Loading branch information
qdm12 committed Feb 16, 2022
1 parent 52cb3f3 commit b3f1729
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
Original file line number Diff line number Diff line change
Expand Up @@ -300,22 +300,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 @@ -334,13 +331,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 @@ -358,7 +354,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 @@ -371,8 +367,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 @@ -384,15 +378,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 @@ -406,7 +406,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 @@ -417,6 +417,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 @@ -432,10 +433,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
Original file line number Diff line number Diff line change
Expand Up @@ -939,18 +939,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 @@ -965,10 +967,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 @@ -993,10 +993,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 @@ -1012,10 +1010,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 @@ -1030,10 +1026,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 @@ -1044,6 +1038,7 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{},
Value: []byte("leaf"),
Generation: 1,
Dirty: true,
},
},
},
Expand All @@ -1056,10 +1051,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 @@ -1076,6 +1069,7 @@ func Test_Trie_insert(t *testing.T) {
Key: []byte{3},
Value: []byte("leaf"),
Generation: 1,
Dirty: true,
},
},
},
Expand All @@ -1087,10 +1081,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 @@ -1105,10 +1097,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 @@ -1148,7 +1138,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 @@ -1159,10 +1149,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 @@ -1180,10 +1168,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 @@ -1201,10 +1187,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 @@ -1234,10 +1218,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 @@ -1268,10 +1250,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 @@ -1288,6 +1268,7 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{
Key: []byte{5, 6},
Value: []byte{6},
Dirty: true,
},
},
},
Expand All @@ -1300,10 +1281,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 @@ -1320,6 +1299,7 @@ func Test_Trie_insertInBranch(t *testing.T) {
&node.Leaf{
Key: []byte{},
Value: []byte{6},
Dirty: true,
},
},
},
Expand All @@ -1332,10 +1312,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 b3f1729

Please sign in to comment.