From 1f6645b19ff2deb470311b7d29f42b190633144c Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Sun, 16 Jun 2024 20:33:43 -0700
Subject: [PATCH] net/ipset: skip the loop over Prefixes when there's only one

For pprof cosmetic/confusion reasons more than performance, but it
might have tiny speed benefit.

Updates #12486

Change-Id: I40e03714f3afa3a7e7f5e1fa99b81c7e889b91b6
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 net/ipset/ipset.go      | 23 +++++++++++++----------
 net/ipset/ipset_test.go |  2 +-
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/net/ipset/ipset.go b/net/ipset/ipset.go
index 58a120432..70068d0a2 100644
--- a/net/ipset/ipset.go
+++ b/net/ipset/ipset.go
@@ -70,19 +70,22 @@ func NewContainsIPFunc(addrs views.Slice[netip.Prefix]) func(ip netip.Addr) bool
 	// If any addr is a prefix with more than a single IP, then do either a
 	// linear scan or a bart table, depending on the number of addrs.
 	if addrs.ContainsFunc(func(p netip.Prefix) bool { return !p.IsSingleIP() }) {
-		if addrs.Len() > 6 {
-			pathForTest("bart")
-			// Built a bart table.
-			t := &bart.Table[struct{}]{}
-			for i := range addrs.Len() {
-				t.Insert(addrs.At(i), struct{}{})
-			}
-			return bartLookup(t)
-		} else {
-			pathForTest("linear-contains")
+		if addrs.Len() == 1 {
+			pathForTest("one-prefix")
+			return addrs.At(0).Contains
+		}
+		if addrs.Len() <= 6 {
 			// Small enough to do a linear search.
+			pathForTest("linear-contains")
 			return prefixContainsLoop(addrs.AsSlice())
 		}
+		pathForTest("bart")
+		// Built a bart table.
+		t := &bart.Table[struct{}]{}
+		for i := range addrs.Len() {
+			t.Insert(addrs.At(i), struct{}{})
+		}
+		return bartLookup(t)
 	}
 	// Fast paths for 1 and 2 IPs:
 	if addrs.Len() == 1 {
diff --git a/net/ipset/ipset_test.go b/net/ipset/ipset_test.go
index 7060df5b3..2df4939cb 100644
--- a/net/ipset/ipset_test.go
+++ b/net/ipset/ipset_test.go
@@ -41,7 +41,7 @@ func aa(ss ...string) (ret []netip.Addr) {
 	{
 		name:    "cidr-list-1",
 		pfx:     pp("10.0.0.0/8"),
-		want:    "linear-contains",
+		want:    "one-prefix",
 		wantIn:  aa("10.0.0.1", "10.2.3.4"),
 		wantOut: aa("8.8.8.8"),
 	},