Optimize CIDR aggregation to improve performance and reduce memory usage

pkg/util/collectionutil/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/util/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/util/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/util/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()