diff --git a/lib/trie/trie.go b/lib/trie/trie.go index 35f742d385c..0d85db0815c 100644 --- a/lib/trie/trie.go +++ b/lib/trie/trie.go @@ -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 @@ -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 @@ -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. @@ -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 @@ -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 } @@ -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, } @@ -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 } @@ -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:] diff --git a/lib/trie/trie_test.go b/lib/trie/trie_test.go index 4212f34cc18..366f6c70d05 100644 --- a/lib/trie/trie_test.go +++ b/lib/trie/trie_test.go @@ -939,7 +939,7 @@ func Test_Trie_insert(t *testing.T) { trie Trie parent Node key []byte - value Node + value []byte newNode Node }{ "nil parent": { @@ -947,10 +947,12 @@ func Test_Trie_insert(t *testing.T) { 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": { @@ -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"), @@ -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"), @@ -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"), @@ -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"), @@ -1044,6 +1038,7 @@ func Test_Trie_insert(t *testing.T) { Key: []byte{}, Value: []byte("leaf"), Generation: 1, + Dirty: true, }, }, }, @@ -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, @@ -1076,6 +1069,7 @@ func Test_Trie_insert(t *testing.T) { Key: []byte{3}, Value: []byte("leaf"), Generation: 1, + Dirty: true, }, }, }, @@ -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"), @@ -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"), @@ -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": { @@ -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"), @@ -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"), @@ -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}, @@ -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}, @@ -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, @@ -1288,6 +1268,7 @@ func Test_Trie_insertInBranch(t *testing.T) { &node.Leaf{ Key: []byte{5, 6}, Value: []byte{6}, + Dirty: true, }, }, }, @@ -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, @@ -1320,6 +1299,7 @@ func Test_Trie_insertInBranch(t *testing.T) { &node.Leaf{ Key: []byte{}, Value: []byte{6}, + Dirty: true, }, }, }, @@ -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},