Optimize CIDR aggregation to improve performance and reduce memory usage

pkg/provider/loadbalancer/fnutil/map.go

@@ -16,6 +16,14 @@ limitations under the License.
 
 package fnutil
 
+func Keys[K comparable, V any](m map[K]V) []K {
+	rv := make([]K, 0, len(m))
+	for k := range m {
+		rv = append(rv, k)
+	}
+	return rv
+}
+
 func Values[K comparable, V any](m map[K]V) []V {
 	rv := make([]V, 0, len(m))
 	for _, v := range m {

pkg/provider/loadbalancer/iputil/bits.go

@@ -0,0 +1,37 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package iputil
+
+// setBitAt sets the bit at the i-th position in the byte slice to the given value.
+// Panics if the index is out of bounds.
+// For example,
+// - setBitAt([0x00, 0x00], 8, 1) returns [0x00, 0b1000_0000].
+// - setBitAt([0xff, 0xff], 0, 0) returns [0b0111_1111, 0xff].
+func setBitAt(bytes []byte, i int, bit uint8) {
+	if bit == 1 {
+		bytes[i/8] |= 1 << (7 - i%8)
+	} else {
+		bytes[i/8] &^= 1 << (7 - i%8)
+	}
+}
+
+// bitAt returns the bit at the i-th position in the byte slice.
+// The return value is either 0 or 1 as uint8.
+// Panics if the index is out of bounds.
+func bitAt(bytes []byte, i int) uint8 {
+	return bytes[i/8] >> (7 - i%8) & 1
+}

pkg/provider/loadbalancer/iputil/bits_test.go

@@ -0,0 +1,103 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package iputil
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_bitAt(t *testing.T) {
+	bytes := []byte{0b1010_1010, 0b0101_0101}
+	assert.Equal(t, uint8(1), bitAt(bytes, 0))
+	assert.Equal(t, uint8(0), bitAt(bytes, 1))
+	assert.Equal(t, uint8(1), bitAt(bytes, 2))
+	assert.Equal(t, uint8(0), bitAt(bytes, 3))
+
+	assert.Equal(t, uint8(1), bitAt(bytes, 4))
+	assert.Equal(t, uint8(0), bitAt(bytes, 5))
+	assert.Equal(t, uint8(1), bitAt(bytes, 6))
+	assert.Equal(t, uint8(0), bitAt(bytes, 7))
+
+	assert.Equal(t, uint8(0), bitAt(bytes, 8))
+	assert.Equal(t, uint8(1), bitAt(bytes, 9))
+	assert.Equal(t, uint8(0), bitAt(bytes, 10))
+	assert.Equal(t, uint8(1), bitAt(bytes, 11))
+
+	assert.Equal(t, uint8(0), bitAt(bytes, 12))
+	assert.Equal(t, uint8(1), bitAt(bytes, 13))
+	assert.Equal(t, uint8(0), bitAt(bytes, 14))
+	assert.Equal(t, uint8(1), bitAt(bytes, 15))
+
+	assert.Panics(t, func() { bitAt(bytes, 16) })
+}
+
+func Test_setBitAt(t *testing.T) {
+	tests := []struct {
+		name     string
+		initial  []byte
+		index    int
+		bit      uint8
+		expected []byte
+	}{
+		{
+			name:     "Set first bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    0,
+			bit:      1,
+			expected: []byte{0b1000_0000},
+		},
+		{
+			name:     "Set last bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    7,
+			bit:      1,
+			expected: []byte{0b0000_0001},
+		},
+		{
+			name:     "Set middle bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    4,
+			bit:      1,
+			expected: []byte{0b0000_1000},
+		},
+		{
+			name:     "Set bit to 0",
+			initial:  []byte{0b1111_1111},
+			index:    3,
+			bit:      0,
+			expected: []byte{0b1110_1111},
+		},
+		{
+			name:     "Set bit in second byte",
+			initial:  []byte{0b0000_0000, 0b0000_0000},
+			index:    9,
+			bit:      1,
+			expected: []byte{0b0000_0000, 0b0100_0000},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			setBitAt(tt.initial, tt.index, tt.bit)
+			assert.Equal(t, tt.expected, tt.initial)
+		})
+	}
+
+	assert.Panics(t, func() { setBitAt([]byte{0x00}, 8, 1) })
+}

pkg/provider/loadbalancer/iputil/prefix.go

@@ -17,8 +17,10 @@ limitations under the License.
 package iputil
 
 import (
+	"bytes"
 	"fmt"
 	"net/netip"
+	"sort"
 )
 
 // IsPrefixesAllowAll returns true if one of the prefixes allows all addresses.
@@ -61,9 +63,108 @@ func GroupPrefixesByFamily(vs []netip.Prefix) ([]netip.Prefix, []netip.Prefix) {
 	return v4, v6
 }
 
-// AggregatePrefixes aggregates prefixes.
-// Overlapping prefixes are merged.
+// ContainsPrefix checks if prefix p fully contains prefix o.
+// It returns true if o is a subset of p, meaning all addresses in o are also in p.
+// This is true when p overlaps with o and p has fewer or equal number of bits than o.
+func ContainsPrefix(p netip.Prefix, o netip.Prefix) bool {
+	return p.Bits() <= o.Bits() && p.Overlaps(o)
+}
+
+// mergeAdjacentPrefixes attempts to merge two adjacent prefixes into a single prefix.
+// It returns the merged prefix and a boolean indicating success.
+// Note: This function only merges adjacent prefixes, not overlapping ones.
+func mergeAdjacentPrefixes(p1, p2 netip.Prefix) (netip.Prefix, bool) {
+	// Merge neighboring prefixes if possible
+	if p1.Bits() != p2.Bits() || p1.Bits() == 0 {
+		return netip.Prefix{}, false
+	}
+
+	var (
+		bits    = p1.Bits()
+		p1Bytes = p1.Addr().AsSlice()
+		p2Bytes = p2.Addr().AsSlice()
+	)
+	if bitAt(p1Bytes, bits-1) == 0 {
+		setBitAt(p1Bytes, bits-1, 1)
+	} else {
+		setBitAt(p2Bytes, bits-1, 1)
+	}
+	if !bytes.Equal(p1Bytes, p2Bytes) {
+		return netip.Prefix{}, false
+	}
+
+	rv, _ := p1.Addr().Prefix(bits - 1)
+	return rv, true
+}
+
+// aggregatePrefixesForSingleIPFamily merges overlapping or adjacent prefixes into a single prefix.
+// The input prefixes must be the same IP family (IPv4 or IPv6).
+// For example,
+// - [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31] (adjacent)
+// - [192.168.0.0/24, 192.168.0.1/32] -> [192.168.1.0/24] (overlapping)
+func aggregatePrefixesForSingleIPFamily(prefixes []netip.Prefix) []netip.Prefix {
+	if len(prefixes) <= 1 {
+		return prefixes
+	}
+
+	sort.Slice(prefixes, func(i, j int) bool {
+		addrCmp := prefixes[i].Addr().Compare(prefixes[j].Addr())
+		if addrCmp == 0 {
+			return prefixes[i].Bits() < prefixes[j].Bits()
+		}
+		return addrCmp < 0
+	})
+
+	var rv = []netip.Prefix{prefixes[0]}
+
+	for i := 1; i < len(prefixes); i++ {
+		last, p := rv[len(rv)-1], prefixes[i]
+		if ContainsPrefix(last, p) {
+			// Skip overlapping prefixes
+			continue
+		}
+		rv = append(rv, p)
+
+		// Merge adjacent prefixes if possible
+		for len(rv) >= 2 {
+			// Merge the last two prefixes if they are adjacent
+			p, ok := mergeAdjacentPrefixes(rv[len(rv)-2], rv[len(rv)-1])
+			if !ok {
+				break
+			}
+
+			// Replace the last two prefixes with the merged prefix
+			rv = rv[:len(rv)-2]
+			rv = append(rv, p)
+		}
+	}
+	return rv
+}
+
+// AggregatePrefixes merges overlapping or adjacent prefixes into a single prefix.
+// It combines prefixes that can be represented by a larger, more inclusive prefix.
+//
+// Examples:
+//   - Adjacent:    [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31]
+//   - Overlapping: [192.168.0.0/24, 192.168.0.1/32] -> [192.168.0.0/24]
 func AggregatePrefixes(prefixes []netip.Prefix) []netip.Prefix {
+	var (
+		v4, v6 = GroupPrefixesByFamily(prefixes)
+	)
+
+	return append(aggregatePrefixesForSingleIPFamily(v4), aggregatePrefixesForSingleIPFamily(v6)...)
+}
+
+// AggregatePrefixesWithPrefixTree merges overlapping or adjacent prefixes into a single prefix.
+//
+// This function uses a prefix tree to aggregate the input prefixes. While it achieves
+// the same result as AggregatePrefixes, it is less efficient. For better performance,
+// use AggregatePrefixes instead.
+//
+// Examples:
+//   - Adjacent:   [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31]
+//   - Overlapping: [192.168.0.0/24, 192.168.0.1/32] -> [192.168.0.0/24]
+func AggregatePrefixesWithPrefixTree(prefixes []netip.Prefix) []netip.Prefix {
 	var (
 		v4, v6 = GroupPrefixesByFamily(prefixes)
 		v4Tree = newPrefixTreeForIPv4()