impl. Size()

README.md

@@ -77,15 +77,19 @@ The API has changed since v0.4.2, 0.5.3. and v0.6.3
   func (t *Table[V]) Overlaps(o *Table[V]) bool
   func (t *Table[V]) OverlapsPrefix(pfx netip.Prefix) bool
   
-  func (t *Table[V]) String() string
-  func (t *Table[V]) Fprint(w io.Writer) error
-  func (t *Table[V]) MarshalText() ([]byte, error)
-  func (t *Table[V]) MarshalJSON() ([]byte, error)
+  func (t *Table[V]) Size() int
+  func (t *Table[V]) Size4() int
+  func (t *Table[V]) Size6() int
 
   func (t *Table[V]) All(yield func(pfx netip.Prefix, val V) bool) bool
   func (t *Table[V]) All4(yield func(pfx netip.Prefix, val V) bool) bool
   func (t *Table[V]) All6(yield func(pfx netip.Prefix, val V) bool) bool
 
+  func (t *Table[V]) String() string
+  func (t *Table[V]) Fprint(w io.Writer) error
+  func (t *Table[V]) MarshalText() ([]byte, error)
+  func (t *Table[V]) MarshalJSON() ([]byte, error)
+
   func (t *Table[V]) DumpList4() []DumpListNode[V]
   func (t *Table[V]) DumpList6() []DumpListNode[V]
 ```

node.go

@@ -65,21 +65,23 @@ func (n *node[V]) prefixRank(baseIdx uint) int {
 
 // insertPrefix adds the route octet/prefixLen, with value val.
 // Just an adapter for insertIdx.
-func (n *node[V]) insertPrefix(octet byte, prefixLen int, val V) {
-	n.insertIdx(prefixToBaseIndex(octet, prefixLen), val)
+func (n *node[V]) insertPrefix(octet byte, prefixLen int, val V) (wasPresent bool) {
+	return n.insertIdx(prefixToBaseIndex(octet, prefixLen), val)
 }
 
 // insertIdx adds the route for baseIdx, with value val.
-func (n *node[V]) insertIdx(baseIdx uint, val V) {
+// incSize reports if the sie counter must incremented.
+func (n *node[V]) insertIdx(baseIdx uint, val V) (wasPresent bool) {
 	// prefix exists, overwrite val
 	if n.prefixesBitset.Test(baseIdx) {
 		n.prefixes[n.prefixRank(baseIdx)] = val
-		return
+		return true
 	}
 
 	// new, insert into bitset and slice
 	n.prefixesBitset.Set(baseIdx)
 	n.prefixes = slices.Insert(n.prefixes, n.prefixRank(baseIdx), val)
+	return false
 }
 
 // deletePrefix removes the route octet/prefixLen. Reports whether the
@@ -105,26 +107,25 @@ func (n *node[V]) deletePrefix(octet byte, prefixLen int) (wasPresent bool) {
 }
 
 // updatePrefix, update or set the value at prefix via callback.
-func (n *node[V]) updatePrefix(octet byte, prefixLen int, cb func(V, bool) V) (val V) {
+func (n *node[V]) updatePrefix(octet byte, prefixLen int, cb func(V, bool) V) (val V, wasPresent bool) {
 	// calculate idx once
 	baseIdx := prefixToBaseIndex(octet, prefixLen)
 
-	var ok bool
 	var rnk int
 
 	// if prefix is set, get current value
-	if ok = n.prefixesBitset.Test(baseIdx); ok {
+	if wasPresent = n.prefixesBitset.Test(baseIdx); wasPresent {
 		rnk = n.prefixRank(baseIdx)
 		val = n.prefixes[rnk]
 	}
 
 	// callback function to get updated or new value
-	val = cb(val, ok)
+	val = cb(val, wasPresent)
 
 	// prefix is already set, update and return value
-	if ok {
+	if wasPresent {
 		n.prefixes[rnk] = val
-		return val
+		return
 	}
 
 	// new prefix, insert into bitset ...
@@ -136,7 +137,7 @@ func (n *node[V]) updatePrefix(octet byte, prefixLen int, cb func(V, bool) V) (v
 	// ... and insert value into slice
 	n.prefixes = slices.Insert(n.prefixes, rnk, val)
 
-	return val
+	return
 }
 
 // lpmByIndex does a route lookup for idx in the 8-bit (stride) routing table

table.go

@@ -29,6 +29,10 @@ type Table[V any] struct {
 	rootV4 *node[V]
 	rootV6 *node[V]
 
+	// number of prefixes in trie
+	sizeV4 int
+	sizeV6 int
+
 	// BitSets have to be initialized.
 	initOnce sync.Once
 }
@@ -102,7 +106,10 @@ func (t *Table[V]) Insert(pfx netip.Prefix, val V) {
 	}
 
 	// insert prefix into node
-	n.insertPrefix(lastOctet, lastOctetBits, val)
+	wasPresent := n.insertPrefix(lastOctet, lastOctetBits, val)
+	if !wasPresent {
+		t.incDecSize(+1, is4)
+	}
 }
 
 // Update or set the value at pfx with a callback function.
@@ -147,7 +154,11 @@ func (t *Table[V]) Update(pfx netip.Prefix, cb func(val V, ok bool) V) V {
 	}
 
 	// update/insert prefix into node
-	return n.updatePrefix(lastOctet, lastOctetBits, cb)
+	val, wasPresent := n.updatePrefix(lastOctet, lastOctetBits, cb)
+	if !wasPresent {
+		t.incDecSize(+1, is4)
+	}
+	return val
 }
 
 // Get returns the associated payload for prefix and true, or false if
@@ -235,10 +246,14 @@ func (t *Table[V]) Delete(pfx netip.Prefix) {
 	}
 
 	// try to delete prefix in trie node
-	if !n.deletePrefix(lastOctet, lastOctetBits) {
-		return // nothing deleted
+	if wasPresent := n.deletePrefix(lastOctet, lastOctetBits); !wasPresent {
+		// nothing deleted
+		return
 	}
 
+	// prefix deleted, decrement the size
+	t.incDecSize(-1, is4)
+
 	// purge dangling nodes after successful deletion
 	for i > 0 {
 		if n.isEmpty() {
@@ -601,6 +616,32 @@ func (t *Table[V]) All6(yield func(pfx netip.Prefix, val V) bool) {
 	t.rootV6.allRec(nil, false, yield)
 }
 
+// Size returns the sum of the IPv4 and IPv6 refixes.
+func (t *Table[V]) Size() int {
+	t.init()
+	return t.sizeV4 + t.sizeV6
+}
+
+// Size4 returns the number of IPv4 refixes.
+func (t *Table[V]) Size4() int {
+	t.init()
+	return t.sizeV4
+}
+
+// Size6 returns the number of IPv6 refixes.
+func (t *Table[V]) Size6() int {
+	t.init()
+	return t.sizeV6
+}
+
+func (t *Table[V]) incDecSize(val int, is4 bool) {
+	if is4 {
+		t.sizeV4 = t.sizeV4 + val
+	} else {
+		t.sizeV6 = t.sizeV6 + val
+	}
+}
+
 // ipToOctets, be careful, do not allocate!
 //
 // intended use: SA4009: argument octets is overwritten before first use

table_test.go

@@ -1874,6 +1874,42 @@ func TestAll(t *testing.T) {
 	})
 }
 
+func TestSize(t *testing.T) {
+	rtbl := new(Table[any])
+	pfxs := randomPrefixes(10_000)
+
+	for _, pfx := range pfxs {
+		switch rand.Intn(3) {
+		case 0:
+			rtbl.Insert(pfx.pfx, nil)
+		case 1:
+			rtbl.Delete(pfx.pfx)
+		case 2:
+			rtbl.Update(pfx.pfx, func(any, bool) any { return nil })
+		}
+	}
+
+	var size4 int
+	var size6 int
+
+	rtbl.All4(func(netip.Prefix, any) bool {
+		size4++
+		return true
+	})
+
+	rtbl.All6(func(netip.Prefix, any) bool {
+		size6++
+		return true
+	})
+
+	if size4 != rtbl.Size4() {
+		t.Errorf("Size4: want: %d, got: %d", size4, rtbl.Size4())
+	}
+	if size6 != rtbl.Size6() {
+		t.Errorf("Size6: want: %d, got: %d", size6, rtbl.Size6())
+	}
+}
+
 // ############################################################################
 
 type tableOverlapsTest struct {