From 8a304357c5f95d3355f8f615695072bd39e2c0b5 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Fri, 21 Aug 2020 11:58:59 -0700
Subject: [PATCH] reduce allocations

---
 cidranger.go |  2 +-
 net/ip.go    |  5 ++++-
 trie.go      | 26 +++++++++++++++++++-------
 trie_test.go |  6 +++---
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/cidranger.go b/cidranger.go
index 2e8f118..f638eed 100644
--- a/cidranger.go
+++ b/cidranger.go
@@ -95,5 +95,5 @@ type Ranger interface {
 // NewPCTrieRanger returns a versionedRanger that supports both IPv4 and IPv6
 // using the path compressed trie implemention.
 func NewPCTrieRanger() Ranger {
-	return newVersionedRanger(newPrefixTree)
+	return newVersionedRanger(newRanger)
 }
diff --git a/net/ip.go b/net/ip.go
index 1ce30e8..de700cf 100644
--- a/net/ip.go
+++ b/net/ip.go
@@ -236,7 +236,10 @@ func (n Network) LeastCommonBitPosition(n1 Network) (uint, error) {
 
 // Equal is the equality test for 2 networks.
 func (n Network) Equal(n1 Network) bool {
-	return n.String() == n1.String()
+	nones, nbits := n.IPNet.Mask.Size()
+	n1ones, n1bits := n1.IPNet.Mask.Size()
+
+	return nones == n1ones && nbits == n1bits && n.IPNet.IP.Equal(n1.IPNet.IP)
 }
 
 func (n Network) String() string {
diff --git a/trie.go b/trie.go
index 2efe60e..6c1c543 100644
--- a/trie.go
+++ b/trie.go
@@ -35,7 +35,7 @@ import (
 // TODO: Implement level-compressed component of the LPC trie.
 type prefixTrie struct {
 	parent   *prefixTrie
-	children []*prefixTrie
+	children [2]*prefixTrie
 
 	numBitsSkipped uint
 	numBitsHandled uint
@@ -46,17 +46,29 @@ type prefixTrie struct {
 	size int // This is only maintained in the root trie.
 }
 
+var ip4ZeroCIDR, ip6ZeroCIDR net.IPNet
+
+func init() {
+	_, v4, _ := net.ParseCIDR("0.0.0.0/0")
+	_, v6, _ := net.ParseCIDR("0::0/0")
+	ip4ZeroCIDR = *v4
+	ip6ZeroCIDR = *v6
+}
+
+func newRanger(version rnet.IPVersion) Ranger {
+	return newPrefixTree(version)
+}
+
 // newPrefixTree creates a new prefixTrie.
-func newPrefixTree(version rnet.IPVersion) Ranger {
-	_, rootNet, _ := net.ParseCIDR("0.0.0.0/0")
+func newPrefixTree(version rnet.IPVersion) *prefixTrie {
+	rootNet := ip4ZeroCIDR
 	if version == rnet.IPv6 {
-		_, rootNet, _ = net.ParseCIDR("0::0/0")
+		rootNet = ip6ZeroCIDR
 	}
 	return &prefixTrie{
-		children:       make([]*prefixTrie, 2, 2),
 		numBitsSkipped: 0,
 		numBitsHandled: 1,
-		network:        rnet.NewNetwork(*rootNet),
+		network:        rnet.NewNetwork(rootNet),
 	}
 }
 
@@ -65,7 +77,7 @@ func newPathprefixTrie(network rnet.Network, numBitsSkipped uint) *prefixTrie {
 	if len(network.Number) == rnet.IPv6Uint32Count {
 		version = rnet.IPv6
 	}
-	path := newPrefixTree(version).(*prefixTrie)
+	path := newPrefixTree(version)
 	path.numBitsSkipped = numBitsSkipped
 	path.network = network.Masked(int(numBitsSkipped))
 	return path
diff --git a/trie_test.go b/trie_test.go
index 1461702..66b1fd0 100644
--- a/trie_test.go
+++ b/trie_test.go
@@ -73,7 +73,7 @@ func TestPrefixTrieInsert(t *testing.T) {
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			trie := newPrefixTree(tc.version).(*prefixTrie)
+			trie := newPrefixTree(tc.version)
 			for _, insert := range tc.inserts {
 				_, network, _ := net.ParseCIDR(insert)
 				err := trie.Insert(NewBasicRangerEntry(*network))
@@ -104,7 +104,7 @@ func TestPrefixTrieInsert(t *testing.T) {
 
 func TestPrefixTrieString(t *testing.T) {
 	inserts := []string{"192.168.0.1/24", "192.168.1.1/24", "192.168.1.1/30"}
-	trie := newPrefixTree(rnet.IPv4).(*prefixTrie)
+	trie := newPrefixTree(rnet.IPv4)
 	for _, insert := range inserts {
 		_, network, _ := net.ParseCIDR(insert)
 		trie.Insert(NewBasicRangerEntry(*network))
@@ -195,7 +195,7 @@ func TestPrefixTrieRemove(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			trie := newPrefixTree(tc.version).(*prefixTrie)
+			trie := newPrefixTree(tc.version)
 			for _, insert := range tc.inserts {
 				_, network, _ := net.ParseCIDR(insert)
 				err := trie.Insert(NewBasicRangerEntry(*network))
-- 
GitLab