Better trie implementations
This commit is contained in:
248
dataset/nettrie/trie_test.go
Normal file
248
dataset/nettrie/trie_test.go
Normal file
@@ -0,0 +1,248 @@
|
||||
package nettrie
|
||||
|
||||
import (
|
||||
"net/netip"
|
||||
"sort"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestTrie_InsertAndContains_IPv4 tests insertion and verifies presence
|
||||
// using ContainsPrefix and Contains for IPv4.
|
||||
func TestTrie_InsertAndContains_IPv4(t *testing.T) {
|
||||
trie := New()
|
||||
|
||||
prefixes := []string{
|
||||
"0.0.0.0/0",
|
||||
"10.0.0.0/8",
|
||||
"192.168.0.0/16",
|
||||
"192.168.1.0/24",
|
||||
"192.168.1.128/25",
|
||||
"192.168.1.200/32",
|
||||
}
|
||||
for _, s := range prefixes {
|
||||
trie.Insert(netip.MustParsePrefix(s))
|
||||
}
|
||||
|
||||
t.Run("Check existing prefixes", func(t *testing.T) {
|
||||
for _, s := range prefixes {
|
||||
p := netip.MustParsePrefix(s)
|
||||
if !trie.ContainsPrefix(p) {
|
||||
t.Errorf("expected trie to contain prefix %s, but it did not", p)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Check non-existing prefixes", func(t *testing.T) {
|
||||
nonExistentPrefixes := []string{
|
||||
"10.0.0.0/9",
|
||||
"192.168.1.0/23",
|
||||
"1.2.3.4/32",
|
||||
}
|
||||
for _, s := range nonExistentPrefixes {
|
||||
p := netip.MustParsePrefix(s)
|
||||
if trie.ContainsPrefix(p) {
|
||||
t.Errorf("expected trie to not contain prefix %s, but it did", p)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Check host addresses with Contains", func(t *testing.T) {
|
||||
if !trie.Contains(netip.MustParseAddr("192.168.1.200")) {
|
||||
t.Error("expected trie to contain host address 192.168.1.200, but it did not")
|
||||
}
|
||||
if trie.Contains(netip.MustParseAddr("192.168.1.201")) {
|
||||
t.Error("expected trie to not contain host address 192.168.1.201, but it did")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestTrie_InsertAndContains_IPv6 tests insertion and verifies presence
|
||||
// using ContainsPrefix and Contains for IPv6.
|
||||
func TestTrie_InsertAndContains_IPv6(t *testing.T) {
|
||||
trie := New()
|
||||
|
||||
prefixes := []string{
|
||||
"::/0",
|
||||
"2001:db8::/32",
|
||||
"2001:db8:acad::/48",
|
||||
"2001:db8:acad:1::/64",
|
||||
"2001:db8:acad:1::1/128",
|
||||
}
|
||||
for _, s := range prefixes {
|
||||
trie.Insert(netip.MustParsePrefix(s))
|
||||
}
|
||||
|
||||
t.Run("Check existing prefixes", func(t *testing.T) {
|
||||
for _, s := range prefixes {
|
||||
p := netip.MustParsePrefix(s)
|
||||
if !trie.ContainsPrefix(p) {
|
||||
t.Errorf("expected trie to contain prefix %s, but it did not", p)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Check non-existing prefixes", func(t *testing.T) {
|
||||
nonExistentPrefixes := []string{
|
||||
"2001:db8::/33",
|
||||
"2001:db8:acad::/47",
|
||||
"2002::/16",
|
||||
}
|
||||
for _, s := range nonExistentPrefixes {
|
||||
p := netip.MustParsePrefix(s)
|
||||
if trie.ContainsPrefix(p) {
|
||||
t.Errorf("expected trie to not contain prefix %s, but it did", p)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Check host addresses with Contains", func(t *testing.T) {
|
||||
if !trie.Contains(netip.MustParseAddr("2001:db8:acad:1::1")) {
|
||||
t.Error("expected trie to contain host address 2001:db8:acad:1::1, but it did not")
|
||||
}
|
||||
if trie.Contains(netip.MustParseAddr("2001:db8:acad:1::2")) {
|
||||
t.Error("expected trie to not contain host address 2001:db8:acad:1::2, but it did")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestTrie_Delete(t *testing.T) {
|
||||
trie := New()
|
||||
|
||||
p8 := netip.MustParsePrefix("10.0.0.0/8")
|
||||
p24 := netip.MustParsePrefix("10.0.0.0/24")
|
||||
pOther := netip.MustParsePrefix("10.0.1.0/24")
|
||||
|
||||
trie.Insert(p8)
|
||||
trie.Insert(p24)
|
||||
trie.Insert(pOther)
|
||||
|
||||
t.Run("Delete Non-Existent", func(t *testing.T) {
|
||||
if trie.Delete(netip.MustParsePrefix("1.2.3.4/32")) {
|
||||
t.Error("Delete returned true for a non-existent prefix")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Delete Leaf Node", func(t *testing.T) {
|
||||
if !trie.Delete(p24) {
|
||||
t.Fatal("Delete returned false for an existing prefix")
|
||||
}
|
||||
if trie.ContainsPrefix(p24) {
|
||||
t.Error("ContainsPrefix should be false after delete")
|
||||
}
|
||||
if !trie.ContainsPrefix(p8) {
|
||||
t.Error("parent prefix should not be affected by child delete")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Delete Intermediate Node", func(t *testing.T) {
|
||||
// Re-insert p24 so we can delete its parent p8
|
||||
trie.Insert(p24)
|
||||
|
||||
if !trie.Delete(p8) {
|
||||
t.Fatal("Delete returned false for an existing intermediate prefix")
|
||||
}
|
||||
if trie.ContainsPrefix(p8) {
|
||||
t.Error("ContainsPrefix for deleted intermediate node should be false")
|
||||
}
|
||||
if !trie.ContainsPrefix(p24) {
|
||||
t.Error("child prefix should still exist after parent was deleted")
|
||||
}
|
||||
if !trie.ContainsPrefix(pOther) {
|
||||
t.Error("other prefix should still exist after parent was deleted")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestTrie_Walk(t *testing.T) {
|
||||
trie := New()
|
||||
prefixes := []string{
|
||||
"10.0.0.0/8",
|
||||
"192.168.1.0/24",
|
||||
"2001:db8::/32",
|
||||
"172.16.0.0/12",
|
||||
"2001:db8:acad::/48",
|
||||
}
|
||||
for _, s := range prefixes {
|
||||
trie.Insert(netip.MustParsePrefix(s))
|
||||
}
|
||||
|
||||
t.Run("Walk all prefixes", func(t *testing.T) {
|
||||
var walkedPrefixes []string
|
||||
trie.Walk(func(p netip.Prefix) bool {
|
||||
walkedPrefixes = append(walkedPrefixes, p.String())
|
||||
return true
|
||||
})
|
||||
|
||||
if len(walkedPrefixes) != len(prefixes) {
|
||||
t.Fatalf("expected to walk %d prefixes, but got %d", len(prefixes), len(walkedPrefixes))
|
||||
}
|
||||
|
||||
sort.Strings(prefixes)
|
||||
sort.Strings(walkedPrefixes)
|
||||
|
||||
for i := range prefixes {
|
||||
if prefixes[i] != walkedPrefixes[i] {
|
||||
t.Errorf("walked prefixes mismatch: expected %v, got %v", prefixes, walkedPrefixes)
|
||||
break
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Stop walk early", func(t *testing.T) {
|
||||
count := 0
|
||||
trie.Walk(func(p netip.Prefix) bool {
|
||||
count++
|
||||
return count < 3 // Stop after visiting 3 prefixes
|
||||
})
|
||||
|
||||
if count != 3 {
|
||||
t.Errorf("expected walk to stop after 3 prefixes, but it visited %d", count)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestTrie_Merge(t *testing.T) {
|
||||
trieA := New()
|
||||
trieA.Insert(netip.MustParsePrefix("10.0.0.0/8"))
|
||||
trieA.Insert(netip.MustParsePrefix("192.168.1.0/24"))
|
||||
trieA.Insert(netip.MustParsePrefix("2001:db8::/32")) // v6
|
||||
|
||||
trieB := New()
|
||||
trieB.Insert(netip.MustParsePrefix("10.1.0.0/16"))
|
||||
trieB.Insert(netip.MustParsePrefix("192.168.1.0/24")) // Overlap
|
||||
trieB.Insert(netip.MustParsePrefix("172.16.0.0/12"))
|
||||
trieB.Insert(netip.MustParsePrefix("2001:db8:acad::/48")) // v6
|
||||
|
||||
trieA.Merge(trieB)
|
||||
|
||||
expectedPrefixes := []string{
|
||||
// From A
|
||||
"10.0.0.0/8",
|
||||
"192.168.1.0/24",
|
||||
"2001:db8::/32",
|
||||
// From B
|
||||
"10.1.0.0/16",
|
||||
"172.16.0.0/12",
|
||||
"2001:db8:acad::/48",
|
||||
}
|
||||
|
||||
for _, s := range expectedPrefixes {
|
||||
p := netip.MustParsePrefix(s)
|
||||
if !trieA.ContainsPrefix(p) {
|
||||
t.Errorf("after merge, trieA is missing prefix %s", p)
|
||||
}
|
||||
}
|
||||
|
||||
// Check a prefix that should not be there
|
||||
if trieA.ContainsPrefix(netip.MustParsePrefix("9.9.9.9/32")) {
|
||||
t.Error("trieA contains a prefix it should not have")
|
||||
}
|
||||
|
||||
// Verify trieB is unchanged
|
||||
if !trieB.ContainsPrefix(netip.MustParsePrefix("172.16.0.0/12")) {
|
||||
t.Error("trieB was modified during merge")
|
||||
}
|
||||
if trieB.ContainsPrefix(netip.MustParsePrefix("10.0.0.0/8")) {
|
||||
t.Error("trieB contains a prefix from trieA")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user