styx/dataset/nettrie/trie_test.go

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")
	}
}