From ddffb6cb1a150922e73d140eb96ab44d0e889e03 Mon Sep 17 00:00:00 2001 From: Daniel Jaglowski Date: Wed, 25 Oct 2023 07:56:44 -0600 Subject: [PATCH] [chore][pkg/stanza] Allow trie to store a value (#28592) This PR enhances the internal `trie` struct (which is not yet in use anywhere in the codebase) such that a node may contain any value. The motivation for this is that I believe we may soon be able to migrate `knownFiles` and `previousPollFiles` to tries, which will prove out the functionality of the trie and also should improve efficiency of the package overall. --- pkg/stanza/fileconsumer/internal/trie/trie.go | 97 ++--- .../fileconsumer/internal/trie/trie_test.go | 341 ++++-------------- 2 files changed, 119 insertions(+), 319 deletions(-) diff --git a/pkg/stanza/fileconsumer/internal/trie/trie.go b/pkg/stanza/fileconsumer/internal/trie/trie.go index a1e53cbfdc8e..3f93bc679ea9 100644 --- a/pkg/stanza/fileconsumer/internal/trie/trie.go +++ b/pkg/stanza/fileconsumer/internal/trie/trie.go @@ -12,93 +12,74 @@ package trie // import "github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/fileconsumer/internal/trie" -type Trie struct { - isEnd bool - children map[byte]*Trie +type Trie[T any] struct { + value *T + children map[byte]*Trie[T] } // NewTrie allocates and returns a new *Trie. -func NewTrie() *Trie { - return &Trie{} +func NewTrie[T any]() *Trie[T] { + return &Trie[T]{children: make(map[byte]*Trie[T])} } -func (trie *Trie) HasKey(key []byte) bool { +func (trie *Trie[T]) Get(key []byte) T { node := trie - isEnd := false - for _, r := range key { - node = node.children[r] + var value T + for _, b := range key { + node = node.children[b] if node == nil { - return isEnd + return value // end of trie } - // We have reached end of the current path and all the previous characters have matched - // Return if current node is leaf and it is not root - if node.isLeaf() && node != trie { - return true + if node.value != nil { + value = *node.value } - // check for any ending node in our current path - isEnd = isEnd || node.isEnd } - return isEnd + return value // end of fingerprint } -// Put inserts the key into the trie -func (trie *Trie) Put(key []byte) { +// TODO []byte, T value parameters +// Put inserts a value into the trie +func (trie *Trie[T]) Put(key []byte, v T) { node := trie - for _, r := range key { - child, ok := node.children[r] + for _, b := range key { + _, ok := node.children[b] if !ok { - if node.children == nil { - node.children = map[byte]*Trie{} - } - child = NewTrie() - node.children[r] = child + node.children[b] = NewTrie[T]() } - node = child + node = node.children[b] } - node.isEnd = true + node.value = &v } -// Delete removes keys from the Trie. Returns true if node was found for the given key. -// If the node or any of its ancestors -// becomes childless as a result, it is removed from the trie. -func (trie *Trie) Delete(key []byte) bool { - var path []*Trie // record ancestors to check later +// Delete removes a value from the Trie. Returns true if the value was found. +// Any empty nodes which become childless as a result are removed from the trie. +func (trie *Trie[T]) Delete(key []byte) bool { + var path []*Trie[T] // record ancestors to check later node := trie for _, b := range key { - path = append(path, node) node = node.children[b] if node == nil { - // node does not exist return false } + path = append(path, node) } - // someonce called Delete() on the node which is not end of current path - if !node.isEnd { + if node.value == nil { return false } - node.isEnd = false - // if leaf, remove it from its parent's children map. Repeat for ancestor path. - if node.isLeaf() { - // iterate backwards over path - for i := len(path) - 1; i >= 0; i-- { - parent := path[i] - b := key[i] - delete(parent.children, b) - if !parent.isLeaf() { - // parent has other children, stop - break - } - parent.children = nil - if parent.isEnd { - // Parent has a value, stop - break - } + + // Remove the value + node.value = nil + + // Iterate back up the path and remove empty nodes + for i := len(path) - 1; i >= 0; i-- { + node := path[i] + b := key[i] + delete(node.children, b) + if len(node.children) > 0 || node.value != nil { + break // node has other children or a value, leave it } } - return true // node (internal or not) existed and its value was nil'd -} -func (trie *Trie) isLeaf() bool { - return len(trie.children) == 0 + return true } diff --git a/pkg/stanza/fileconsumer/internal/trie/trie_test.go b/pkg/stanza/fileconsumer/internal/trie/trie_test.go index e1976b0d1c25..e58a16844803 100644 --- a/pkg/stanza/fileconsumer/internal/trie/trie_test.go +++ b/pkg/stanza/fileconsumer/internal/trie/trie_test.go @@ -9,217 +9,31 @@ import ( "github.com/stretchr/testify/assert" ) -type testCase struct { - value []byte - matchExpected bool - delete bool // If we should delete the given value from the trie - deleteExpected bool -} - -type trieTest struct { - initialItems []string - testCases []testCase - name string -} - func TestTrie(t *testing.T) { - // Run all the test cases in sequential order on the same trie to test the expected behavior - testCases := []trieTest{ - { - name: "TrieCase_Normal", - initialItems: []string{"ABCD", "XYZ"}, - testCases: []testCase{ - { - value: []byte("ABCDEFG"), - matchExpected: true, - }, - { - value: []byte("ABCD"), - matchExpected: false, - delete: true, - deleteExpected: true, - }, - { - value: []byte("ABCDEFG"), - matchExpected: false, - }, - { - value: []byte("XYZ"), - matchExpected: true, - }, - { - value: []byte("XYZBlaBla"), - matchExpected: true, - }, - { - value: []byte("X"), - matchExpected: false, - }, - }, - }, - { - name: "TrieCase_SimilarKeys_1", - initialItems: []string{"ABCDEFG", "ABCD"}, - testCases: []testCase{ - { - value: []byte("ABCDEFG"), - matchExpected: false, - delete: true, - deleteExpected: true, - }, - { - value: []byte("ABCD"), - matchExpected: true, - }, - { - value: []byte("ABCDEFG"), - matchExpected: true, - }, - { - value: []byte("ABCDEFGHI"), - matchExpected: true, - }, - }, - }, - { - name: "TrieCase_SimilarKeys_2", - initialItems: []string{"ABCDEFG", "ABCD"}, - testCases: []testCase{ - { - value: []byte("ABCD"), - delete: true, - deleteExpected: true, - }, - { - value: []byte("ABCD"), - matchExpected: false, - }, - { - value: []byte("ABCDEF"), - matchExpected: true, - }, - { - value: []byte("ABCDEFGHI"), - matchExpected: true, - }, - }, - }, - { - name: "TrieCase_Different", - initialItems: []string{"ABCD", "XYZ"}, - testCases: []testCase{ - { - value: []byte("ABCEFG"), - }, - { - value: []byte("ABCDXYZ"), - matchExpected: true, - }, - { - value: []byte("ABXE"), - }, - }, - }, - { - name: "TrieCase_Exact", - initialItems: []string{"ABCDEFG", "ABCD"}, - testCases: []testCase{ - { - value: []byte("ABCDEFG"), - matchExpected: true, - }, - { - value: []byte("ABCD"), - matchExpected: true, - }, - { - value: []byte("ABCDE"), - matchExpected: true, - }, - }, - }, - { - name: "TrieCase_DeleteFalse", - initialItems: []string{"ABCDEFG"}, - testCases: []testCase{ - { - value: []byte("ABCDEFG"), - delete: true, - deleteExpected: true, - }, - { - value: []byte("ABCD"), - }, - { - value: []byte("XYZ"), - delete: true, - // it should be false, as we haven't inserted such values - deleteExpected: false, - }, - }, - }, - { - name: "TrieCase_Complex", - initialItems: []string{"ABCDE", "ABC"}, - testCases: []testCase{ - { - value: []byte("ABCDEXYZ"), - matchExpected: true, - }, - { - value: []byte("ABCXYZ"), - matchExpected: true, - }, - }, - }, - } - - for _, tc := range testCases { - trie := NewTrie() - for _, k := range tc.initialItems { - trie.Put([]byte(k)) - } - t.Run(tc.name, func(t *testing.T) { - for _, step := range tc.testCases { - if step.delete { - // Delete the value and check if it was deleted successfully - assert.Equal(t, trie.Delete(step.value), step.deleteExpected) - } else { - assert.Equal(t, trie.HasKey(step.value), step.matchExpected) - if !step.matchExpected { - trie.Put(step.value) - } - } - } - }) - } -} - -func TestTrieOpSequences(t *testing.T) { opTree{ continuations: map[string]opTree{ "Found:ABC": { // First poll finds only ABC. ops: []testOp{ - has("ABC", false, "empty trie"), - put("ABC"), + get("ABC", nil, "empty trie"), + put("ABC", 3), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC and remove from trie ops: []testOp{ del("ABC", "was just added"), - has("ABC", false, "was just deleted"), + get("ABC", nil, "was just deleted"), }, }, "Found:ABCDEF": { // Next poll finds ABCDEF ops: []testOp{ - has("ABCDEF", true, "recognize ABC w/ DEF appended"), + get("ABCDEF", 3, "recognize ABC w/ DEF appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Done reading the file, remove it as ABC ops: []testOp{ del("ABC", "should be deleted"), - has("ABC", false, "should have been deleted"), + get("ABC", nil, "should have been deleted"), }, }, }, @@ -230,24 +44,24 @@ func TestTrieOpSequences(t *testing.T) { ops: []testOp{ // In order to avoid overwriting ABC with ABCDEF, we need to add ABCDEF first. // TODO Should poll results be sorted by decreasing length before adding to trie? - has("ABCDEF", false, "empty trie"), - put("ABCDEF"), - has("ABC", false, "ABC should not be in trie yet"), - put("ABC"), - has("ABCDEF", true, "this would pass if either ABC or ABCDEF were added, but make sure nothing changed"), + get("ABCDEF", nil, "empty trie"), + put("ABCDEF", 6), + get("ABC", nil, "ABC should not be in trie yet"), + put("ABC", 3), + get("ABCDEF", 6, "make sure adding ABC did not change value of ABCDEF"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC exists"), - has("ABC", false, "ABC should have been deleted"), - has("ABCDEF", true, "ABCDEF should not have been deleted"), + get("ABC", nil, "ABC should have been deleted"), + get("ABCDEF", 6, "ABCDEF should not have been deleted"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF exists"), - has("ABCDEF", false, "ABCDEF should have been deleted"), + get("ABCDEF", nil, "ABCDEF should have been deleted"), }, }, }, @@ -255,46 +69,46 @@ func TestTrieOpSequences(t *testing.T) { "Done:ABCDEF": { // Finish reading ABCDEF and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF exists"), - has("ABC", true, "should not have been deleted"), - // has(ABCDEF) will still return true because ABC is still in trie + get("ABC", 3, "should not have been deleted"), + // get(ABCDEF) will still return true because ABC is still in trie }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC exists"), - has("ABC", false, "just deleted ABC"), + get("ABC", nil, "just deleted ABC"), }, }, }, }, "Found:ABCxyz,ABCDEF": { // Next poll finds ABCxyz and ABCDEF ops: []testOp{ - has("ABCxyz", true, "recognize ABC w/ xyz appended"), - has("ABCDEF", true, "recognize ABCDEF"), + get("ABCxyz", 3, "recognize ABC w/ xyz appended"), + get("ABCDEF", 6, "recognize ABCDEF"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "should still be known as ABC"), - has("ABC", false, "just deleted ABC"), - has("ABCDEF", true, "ABCDEF should not have been affected"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEF", 6, "ABCDEF should not have been affected"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF exists"), - has("ABCDEF", false, "ABCDEF should have been deleted"), + get("ABCDEF", nil, "ABCDEF should have been deleted"), }, }, "Found:ABCDEFxyz": { // Next poll finds ABCDEFxyz ops: []testOp{ - has("ABCDEFxyz", true, "recognize ABCDEF w/ xyz appended"), + get("ABCDEFxyz", 6, "recognize ABCDEF w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEFxyz exists"), - has("ABCDEF", false, "ABCDEFxyz should have been deleted"), + get("ABCDEF", nil, "ABCDEFxyz should have been deleted"), }, }, }, @@ -304,26 +118,26 @@ func TestTrieOpSequences(t *testing.T) { "Done:ABCDEF": { // Finish reading ABC and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF exists"), - has("ABCxyz", true, "should still exist as ABC"), + get("ABCxyz", 3, "should still exist as ABC"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABCDEF and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEF", false, "deleted this earlier but has(ABCDEF) was true until after del(ABC)"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEF", nil, "deleted this earlier but get(ABCDEF) was true until after del(ABC)"), }, }, "Found:ABCxyz": { // Next poll finds ABCDEFxyz ops: []testOp{ - has("ABCxyz", true, "recognize ABC w/ xyz appended"), + get("ABCxyz", 3, "recognize ABC w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "still known as ABC"), - has("ABC", false, "just deleted ABC"), - has("ABCDEF", false, "deleted this earlier but has(ABCDEF) was true until after del(ABC)"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEF", nil, "deleted this earlier but get(ABCDEF) was true until after del(ABC)"), }, }, }, @@ -332,21 +146,21 @@ func TestTrieOpSequences(t *testing.T) { }, "Found:ABCxyz,ABCDEFxyz": { // Next poll finds ABCxyz and ABCDEFxyz ops: []testOp{ - has("ABCxyz", true, "recognize ABC w/ xyz appended"), - has("ABCDEFxyz", true, "recognize ABCDEF w/ xyz appended"), + get("ABCxyz", 3, "recognize ABC w/ xyz appended"), + get("ABCDEFxyz", 6, "recognize ABCDEF w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "should still be present as ABC"), - has("ABC", false, "just deleted ABC"), - has("ABCDEFxyz", true, "should still exist as ABCDEF"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEFxyz", 6, "should still exist as ABCDEF"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABCDEF", false, "just deleted ABCDEF"), + get("ABCDEF", nil, "just deleted ABCDEF"), }, }, }, @@ -354,13 +168,13 @@ func TestTrieOpSequences(t *testing.T) { "Done:ABCDEF": { // Finish reading ABCDEFxyz and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABCxyz", true, "should still exist as ABC"), + get("ABCxyz", 3, "should still exist as ABC"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "just confirmed ABCxyz exists"), - has("ABC", false, "just deleted ABC"), + get("ABC", nil, "just deleted ABC"), }, }, }, @@ -371,21 +185,21 @@ func TestTrieOpSequences(t *testing.T) { }, "Found:ABC,ABCDEFxyz": { // Next poll finds ABC and ABCDEFxyz ops: []testOp{ - has("ABC", true, "recognize ABC"), - has("ABCDEFxyz", true, "recognize ABCDEF w/ xyz appended"), + get("ABC", 3, "recognize ABC"), + get("ABCDEFxyz", 6, "recognize ABCDEF w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEFxyz", true, "should still exist as ABCDEF"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEFxyz", 6, "should still exist as ABCDEF"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABCDEF", false, "just deleted ABCDEF"), + get("ABCDEF", nil, "just deleted ABCDEF"), }, }, }, @@ -393,25 +207,25 @@ func TestTrieOpSequences(t *testing.T) { "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABC", true, "ABC should not have been deleted"), + get("ABC", 3, "ABC should not have been deleted"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC exists"), - has("ABC", false, "just deleted ABC"), + get("ABC", nil, "just deleted ABC"), }, }, "Found:ABCxyz": { // Next poll finds ABCxyz ops: []testOp{ - has("ABCxyz", true, "recognize ABC w/ xyz appended"), + get("ABCxyz", 3, "recognize ABC w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC(xyz) exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEF", false, "deleted this earlier but has(ABCDEF) was true until after del(ABC)"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEF", nil, "deleted this earlier but get(ABCDEF) was true until after del(ABC)"), }, }, }, @@ -420,21 +234,21 @@ func TestTrieOpSequences(t *testing.T) { }, "Found:ABCxyz,ABCDEFxyz": { // Next poll finds ABCxyz and ABCDEFxyz ops: []testOp{ - has("ABCxyz", true, "recognize ABC w/ xyz appended"), - has("ABCDEFxyz", true, "recognize ABCDEF w/ xyz appended"), + get("ABCxyz", 3, "recognize ABC w/ xyz appended"), + get("ABCDEFxyz", 6, "recognize ABCDEF w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC(xyz) exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEFxyz", true, "ABCDEF(xyz) should not have been affected"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEFxyz", 6, "ABCDEF(xyz) should not have been affected"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEFxyz exists"), - has("ABCDEF", false, "just deleted ABCDEF"), + get("ABCDEF", nil, "just deleted ABCDEF"), }, }, }, @@ -442,14 +256,14 @@ func TestTrieOpSequences(t *testing.T) { "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABCxyz", true, "should still exist as ABC"), + get("ABCxyz", 3, "should still exist as ABC"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "just confirmed ABCxyz exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEF", false, "deleted this earlier but has(ABCDEF) was true until after del(ABC)"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEF", nil, "deleted this earlier but get(ABCDEF) was true until after del(ABC)"), }, }, }, @@ -461,21 +275,21 @@ func TestTrieOpSequences(t *testing.T) { "Found:ABCxyz,ABCDEFxyz": { // Next poll finds ABCxyz and ABCDEFxyz ops: []testOp{ // Process longer string first - has("ABCDEFxyz", true, "recognize ABCDEF w/ xyz appended"), - has("ABCxyz", true, "recognize ABC w/ xyz appended"), + get("ABCDEFxyz", 6, "recognize ABCDEF w/ xyz appended"), + get("ABCxyz", 3, "recognize ABC w/ xyz appended"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC(xyz) exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEFxyz", true, "ABCDEF(xyz) should not have been deleted"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEFxyz", 6, "ABCDEF(xyz) should not have been deleted"), }, continuations: map[string]opTree{ "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABCDEF", false, "just deleted ABCDEF"), + get("ABCDEF", nil, "just deleted ABCDEF"), }, }, }, @@ -483,14 +297,14 @@ func TestTrieOpSequences(t *testing.T) { "Done:ABCDEF": { // Finish reading ABCDEF(xyz) and remove from trie ops: []testOp{ del("ABCDEF", "just confirmed ABCDEF(xyz) exists"), - has("ABCxyz", true, "ABC(xyz) should not have been deleted"), + get("ABCxyz", 3, "ABC(xyz) should not have been deleted"), }, continuations: map[string]opTree{ "Done:ABC": { // Finish reading ABC(xyz) and remove from trie ops: []testOp{ del("ABC", "just confirmed ABC(xyz) exists"), - has("ABC", false, "just deleted ABC"), - has("ABCDEF", false, "deleted this earlier but has(ABCDEF) was true until after del(ABC)"), + get("ABC", nil, "just deleted ABC"), + get("ABCDEF", nil, "deleted this earlier but get(ABCDEF) was true until after del(ABC)"), }, }, }, @@ -503,28 +317,33 @@ func TestTrieOpSequences(t *testing.T) { }.run([]testOp{})(t) } -// testOp is one HasKey, Put, or Delete call to the trie, -// along with validation of expectations. -type testOp func(t *testing.T, trie *Trie) +// testOp is one Get, Put, or Delete call to the trie, along with validation of expectations. +type testOp func(t *testing.T, trie *Trie[any]) -func has(key string, expect bool, why string) testOp { - return func(t *testing.T, trie *Trie) { - assert.Equalf(t, trie.HasKey([]byte(key)), expect, why) +func get(key string, expect any, why string) testOp { + return func(t *testing.T, trie *Trie[any]) { + assert.Equalf(t, expect, trie.Get([]byte(key)), why) } } // put automatically asserts that the trie contains the key after adding. -func put(key string) testOp { - return func(t *testing.T, trie *Trie) { - trie.Put([]byte(key)) - assert.Truef(t, trie.HasKey([]byte(key)), "called Put(%s) but HasKey(%s) is still false", key, key) +func put(key string, val any) testOp { + return func(t *testing.T, trie *Trie[any]) { + trie.Put([]byte(key), val) + assert.Equalf(t, val, trie.Get([]byte(key)), "called Put(%s, %d) but HasKey(%s) does not return %d", key, key) } } // del automatically asserts that the trie no longer contains the key after deleting it. func del(key string, why string) testOp { - return func(t *testing.T, trie *Trie) { - assert.Equalf(t, trie.Delete([]byte(key)), true, why) + return func(t *testing.T, trie *Trie[any]) { + val := trie.Get([]byte(key)) + if val == nil { + assert.Falsef(t, trie.Delete([]byte(key)), why) + } else { + assert.Truef(t, trie.Delete([]byte(key)), why) + assert.Falsef(t, trie.Delete([]byte(key)), "called Del(%s) twice in a row and got true both times") + } } } @@ -539,7 +358,7 @@ type opTree struct { func (ot opTree) run(opSequence []testOp) func(*testing.T) { return func(t *testing.T) { - trie := NewTrie() + trie := NewTrie[any]() opSequence = append(opSequence, ot.ops...) for _, op := range opSequence { op(t, trie)