Better trie implementations

dataset/network.go

@@ -2,10 +2,37 @@ package dataset
 
 import (
 	"net"
+	"net/netip"
 
+	"git.maze.io/maze/styx/dataset/nettrie"
 	"github.com/yl2chen/cidranger"
 )
 
+type NetworkTrie struct {
+	*nettrie.Trie
+}
+
+func NewNetworkTrie(prefixes ...netip.Prefix) *NetworkTrie {
+	trie := &NetworkTrie{
+		Trie: nettrie.New(),
+	}
+	for _, prefix := range prefixes {
+		trie.Insert(prefix)
+	}
+	return trie
+}
+
+func (trie *NetworkTrie) ContainsIP(ip net.IP) bool {
+	if ip == nil {
+		return false
+	}
+	addr, ok := netip.AddrFromSlice(ip)
+	if !ok {
+		return false
+	}
+	return trie.Contains(addr)
+}
+
 type NetworkTree struct {
 	ranger cidranger.Ranger
 }