Better trie implementations

dataset/dnstrie/valuetrie_test.go

@@ -0,0 +1,182 @@
+package dnstrie
+
+import (
+	"testing"
+)
+
+func TestValueTrie(t *testing.T) {
+	t.Run("InsertAndSearchStrings", func(t *testing.T) {
+		trie := NewValue[string]()
+		domain := "www.example.com"
+		value := "192.0.2.1"
+		err := trie.Insert(domain, value)
+		if err != nil {
+			t.Fatalf("Expected no error on insert, got %v", err)
+		}
+
+		foundValue, found := trie.Search(domain)
+		if !found {
+			t.Fatalf("Expected to find domain '%s', but did not", domain)
+		}
+		if foundValue != value {
+			t.Errorf("Expected value '%s', got '%s'", value, foundValue)
+		}
+	})
+
+	t.Run("SearchNotFound", func(t *testing.T) {
+		trie := NewValue[string]()
+		err := trie.Insert("example.com", "value")
+		if err != nil {
+			t.Fatalf("Insert failed: %v", err)
+		}
+
+		_, found := trie.Search("nonexistent.com")
+		if found {
+			t.Error("Expected not to find domain 'nonexistent.com', but did")
+		}
+
+		// Search for a path that exists but is not a terminal node
+		_, found = trie.Search("com")
+		if found {
+			t.Error("Expected not to find non-terminal path 'com', but did")
+		}
+	})
+
+	t.Run("InsertInvalidDomain", func(t *testing.T) {
+		trie := NewValue[string]()
+		err := trie.Insert("not-a-valid-domain-", "value")
+		if err == nil {
+			t.Error("Expected an error when inserting an invalid domain, but got nil")
+		}
+	})
+
+	t.Run("OverwriteValue", func(t *testing.T) {
+		trie := NewValue[string]()
+		domain := "overwrite.com"
+		initialValue := "first"
+		overwriteValue := "second"
+
+		err := trie.Insert(domain, initialValue)
+		if err != nil {
+			t.Fatalf("Initial insert failed: %v", err)
+		}
+		err = trie.Insert(domain, overwriteValue)
+		if err != nil {
+			t.Fatalf("Overwrite insert failed: %v", err)
+		}
+
+		foundValue, found := trie.Search(domain)
+		if !found {
+			t.Fatalf("Expected to find domain '%s' after overwrite", domain)
+		}
+		if foundValue != overwriteValue {
+			t.Errorf("Expected overwritten value '%s', got '%s'", overwriteValue, foundValue)
+		}
+	})
+
+	t.Run("Canonicalization", func(t *testing.T) {
+		trie := NewValue[string]()
+		value := "canonical"
+		// Insert lowercase with trailing dot
+		err := trie.Insert("case.example.org.", value)
+		if err != nil {
+			t.Fatalf("Insert failed: %v", err)
+		}
+
+		// Search uppercase without trailing dot
+		foundValue, found := trie.Search("CASE.EXAMPLE.ORG")
+		if !found {
+			t.Fatal("Failed to find domain with different case and no trailing dot")
+		}
+		if foundValue != value {
+			t.Errorf("Expected value '%s' for canonical search, got '%s'", value, foundValue)
+		}
+	})
+
+	t.Run("InsertAndSearchIntegers", func(t *testing.T) {
+		trie := NewValue[int]()
+		domain := "int.example.com"
+		value := 12345
+
+		err := trie.Insert(domain, value)
+		if err != nil {
+			t.Fatalf("Expected no error on insert for int trie, got %v", err)
+		}
+
+		foundValue, found := trie.Search(domain)
+		if !found {
+			t.Fatalf("Expected to find domain '%s' in int trie, but did not", domain)
+		}
+		if foundValue != value {
+			t.Errorf("Expected int value %d, got %d", value, foundValue)
+		}
+
+		// Search for a non-existent domain in the int trie
+		_, found = trie.Search("nonexistent.int.example.com")
+		if found {
+			t.Error("Found a nonexistent domain in the int trie")
+		}
+	})
+
+	t.Run("Contains", func(t *testing.T) {
+		trie := NewValue[string]()
+		domain := "contains.example.com"
+		err := trie.Insert(domain, "some-value")
+		if err != nil {
+			t.Fatalf("Insert failed: %v", err)
+		}
+
+		if !trie.Contains(domain) {
+			t.Errorf("Expected Contains('%s') to be true, but it was false", domain)
+		}
+
+		if trie.Contains("nonexistent." + domain) {
+			t.Errorf("Expected Contains for nonexistent domain to be false, but it was true")
+		}
+
+		if trie.Contains("example.com") {
+			t.Error("Expected Contains for a non-terminal path to be false, but it was true")
+		}
+	})
+
+	t.Run("MergeTries", func(t *testing.T) {
+		trie1 := NewValue[int]()
+		trie1.Insert("example.com", 100)
+		trie1.Insert("sub.example.com", 200)
+
+		trie2 := NewValue[int]()
+		trie2.Insert("google.com", 300)
+		trie2.Insert("sub.example.com", 999) // Overlapping domain, new value
+		trie2.Insert("another.net", 400)
+
+		trie1.Merge(trie2)
+
+		// Test domains from both tries are present
+		if !trie1.Contains("example.com") {
+			t.Error("Merge failed: trie1 should contain 'example.com'")
+		}
+		if !trie1.Contains("google.com") {
+			t.Error("Merge failed: trie1 should contain 'google.com'")
+		}
+		if !trie1.Contains("another.net") {
+			t.Error("Merge failed: trie1 should contain 'another.net'")
+		}
+
+		// Test that overlapping value was updated from trie2
+		val, found := trie1.Search("sub.example.com")
+		if !found || val != 999 {
+			t.Errorf("Expected value for overlapping domain to be 999, but got %d", val)
+		}
+
+		// Test that non-overlapping value from trie1 is intact
+		val, found = trie1.Search("example.com")
+		if !found || val != 100 {
+			t.Errorf("Expected value for 'example.com' to be 100, but got %d", val)
+		}
+
+		// Ensure trie2 is not modified
+		if _, found := trie2.Search("example.com"); found {
+			t.Error("Source trie (trie2) was modified after merge")
+		}
+	})
+}