Num prefixes

ipv4/prefix.go

@@ -2,6 +2,7 @@ package ipv4
 
 import (
 	"fmt"
+	"math"
 	"net"
 )
 
@@ -173,7 +174,22 @@ func (me Prefix) Contains(other SetI) bool {
 // NumAddresses returns the number of addresses in the prefix, including network and
 // broadcast addresses. It ignores any bits set in the host part of the address.
 func (me Prefix) NumAddresses() int64 {
-	return 1 << (addressSize - me.Length())
+	count, _ := me.NumPrefixes(32)
+	return int64(count)
+}
+
+// NumPrefixes returns the number of prefixes of the given length contained in
+// this prefix.
+func (me Prefix) NumPrefixes(length uint32) (count uint64, err error) {
+	switch {
+	case 32 < length:
+		err = fmt.Errorf("length is greater than 32")
+	case length < me.length:
+		count = 0
+	default:
+		count = uint64(math.Pow(2, float64(length-me.length)))
+	}
+	return
 }
 
 // String returns the string representation of this prefix in dotted-quad cidr

ipv4/prefix_test.go

@@ -534,3 +534,46 @@ func TestPrefixAsMapKey(t *testing.T) {
 
 	assert.True(t, m[_p("203.0.113.1/32")])
 }
+
+func TestPrefixNumPrefixes(t *testing.T) {
+	tests := []struct {
+		description string
+		prefix      Prefix
+		length      uint32
+		count       uint64
+		error       bool
+	}{
+		{
+			description: "bad length",
+			length:      33,
+			error:       true,
+		}, {
+			description: "same size",
+			prefix:      _p("203.0.113.0/24"),
+			length:      24,
+			count:       1,
+		}, {
+			description: "too big",
+			prefix:      _p("203.0.113.0/24"),
+			length:      23,
+		}, {
+			description: "28",
+			prefix:      _p("203.0.113.0/24"),
+			length:      30,
+			count:       0x40,
+		}, {
+			description: "26",
+			prefix:      _p("203.0.113.0/24"),
+			length:      26,
+			count:       4,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.description, func(t *testing.T) {
+			count, err := tt.prefix.NumPrefixes(tt.length)
+			assert.Equal(t, tt.error, err != nil)
+			assert.Equal(t, tt.count, count)
+		})
+	}
+}

ipv4/range.go

@@ -28,9 +28,16 @@ func RangeFromAddresses(first, last Address) (r Range, empty bool) {
 	}, false
 }
 
+// NumPrefixes returns the number of prefixes of the given prefix length in
+// this range.
+func (me Range) NumPrefixes(length uint32) (count uint64, err error) {
+	return me.Set().NumPrefixes(length)
+}
+
 // NumAddresses returns the number of addresses in the range
 func (me Range) NumAddresses() int64 {
-	return 1 + int64(me.last.ui-me.first.ui)
+	c, _ := me.NumPrefixes(32)
+	return int64(c)
 }
 
 // First returns the first address in the range

ipv4/range_test.go

@@ -398,3 +398,33 @@ func TestRangeAsMapKey(t *testing.T) {
 
 	assert.True(t, m[_r(_a("203.0.113.0"), _a("203.0.113.127"))])
 }
+
+func TestRangeNumPrefixes(t *testing.T) {
+	tests := []struct {
+		description string
+		r           Range
+		length      uint32
+		count       uint64
+		error       bool
+	}{
+		{
+			description: "simple",
+			r:           _r(_a("203.0.113.0"), _a("203.0.114.0")),
+			length:      31,
+			count:       128,
+		}, {
+			description: "bonus",
+			r:           _r(_a("203.0.113.0"), _a("203.0.114.0")),
+			length:      32,
+			count:       257,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.description, func(t *testing.T) {
+			count, err := tt.r.NumPrefixes(tt.length)
+			assert.Equal(t, tt.error, err != nil)
+			assert.Equal(t, tt.count, count)
+		})
+	}
+}

ipv4/set.go

@@ -244,6 +244,19 @@ func (me Set) WalkAddresses(callback func(Address) bool) bool {
 	})
 }
 
+// NumPrefixes returns the number of prefixes of the given prefix length in
+// this set.
+func (me Set) NumPrefixes(length uint32) (count uint64, err error) {
+	// NOTE This could be done more efficiently with its own recursive
+	// implementation that stops descending when the prefixes are too small.
+	me.WalkPrefixes(func(p Prefix) bool {
+		c, _ := p.NumPrefixes(length)
+		count += c
+		return true
+	})
+	return
+}
+
 // WalkRanges calls `callback` for each IP range in lexographical order. It
 // stops iteration immediately if callback returns false. It always uses the
 // largest ranges possible so if two ranges are adjacent and can be combined,

ipv4/set_test.go

@@ -2,6 +2,7 @@ package ipv4
 
 import (
 	"math/rand"
+	"strconv"
 	"sync"
 	"testing"
 
@@ -860,3 +861,95 @@ func TestOldEqualAllIPv4(t *testing.T) {
 	assert.True(t, c.Equal(a))
 	assert.True(t, c.Equal(b))
 }
+
+func TestSetNumPrefixes(t *testing.T) {
+	tests := []struct {
+		description string
+		prefixes    []SetI
+		length      uint32
+		count       uint64
+		error       bool
+	}{
+		{
+			description: "empty set",
+			prefixes:    []SetI{},
+			length:      32,
+			count:       0,
+		}, {
+			description: "single prefix",
+			prefixes:    []SetI{_p("203.0.113.0/8")},
+			length:      16,
+			count:       256,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.description, func(t *testing.T) {
+			s := Set{}.Build(func(s Set_) bool {
+				for _, prefix := range tt.prefixes {
+					s.Insert(prefix)
+				}
+				return true
+			})
+
+			count, err := s.NumPrefixes(tt.length)
+			assert.Equal(t, tt.error, err != nil)
+			if !tt.error {
+				assert.Equal(t, tt.count, count)
+			}
+		})
+	}
+}
+
+func TestSetNumPrefixesStairs(t *testing.T) {
+	tests := []struct {
+		description string
+		length      uint32
+		count       uint64
+	}{
+		{length: 8, count: 0x00001},
+		{length: 9, count: 0x00003},
+		{length: 10, count: 0x00007},
+		{length: 11, count: 0x0000f},
+		{length: 12, count: 0x0001f},
+		{length: 13, count: 0x0003f},
+		{length: 14, count: 0x0007f},
+		{length: 15, count: 0x000ff},
+		{length: 16, count: 0x001ff},
+		{length: 17, count: 0x003ff},
+		{length: 18, count: 0x007ff},
+		{length: 19, count: 0x00fff},
+		{length: 20, count: 0x01fff},
+		{length: 21, count: 0x03fff},
+		{length: 22, count: 0x07fff},
+		{length: 23, count: 0x0fffe},
+		{length: 24, count: 0x1fffc},
+	}
+
+	s := Set{}.Build(func(s Set_) bool {
+		s.Insert(_p("10.0.0.0/8"))
+		s.Insert(_p("11.0.0.0/9"))
+		s.Insert(_p("11.128.0.0/10"))
+		s.Insert(_p("11.192.0.0/11"))
+		s.Insert(_p("11.224.0.0/12"))
+		s.Insert(_p("11.240.0.0/13"))
+		s.Insert(_p("11.248.0.0/14"))
+		s.Insert(_p("11.252.0.0/15"))
+		s.Insert(_p("11.254.0.0/16"))
+		s.Insert(_p("11.255.0.0/17"))
+		s.Insert(_p("11.255.128.0/18"))
+		s.Insert(_p("11.255.192.0/19"))
+		s.Insert(_p("11.255.224.0/20"))
+		s.Insert(_p("11.255.240.0/21"))
+		s.Insert(_p("11.255.248.0/22"))
+		return true
+	})
+
+	for _, tt := range tests {
+		t.Run(strconv.FormatUint(uint64(tt.length), 10), func(t *testing.T) {
+			count, err := s.NumPrefixes(tt.length)
+			assert.Nil(t, err)
+			assert.Equal(t, tt.count, count)
+		})
+	}
+}

ipv6/prefix.go

@@ -2,6 +2,7 @@ package ipv6
 
 import (
 	"fmt"
+	"math"
 	"net"
 )
 
@@ -116,6 +117,22 @@ func (me Prefix) Length() int {
 	return int(me.length)
 }
 
+// NumPrefixes returns the number of prefixes of the given length contained in
+// this prefix.
+func (me Prefix) NumPrefixes(length uint32) (count uint64, err error) {
+	switch {
+	case 128 < length:
+		err = fmt.Errorf("length is greater than 128")
+	case length < me.length:
+		count = 0
+	case length-me.length > 63:
+		err = fmt.Errorf("overflow")
+	default:
+		count = uint64(math.Pow(2, float64(length-me.length)))
+	}
+	return
+}
+
 // Mask returns a new Address with 1s in the first `length` bits and then 0s
 // representing the network mask for this prefix.
 func (me Prefix) Mask() Mask {

ipv6/prefix_test.go

@@ -488,3 +488,50 @@ func TestPrefixAsMapKey(t *testing.T) {
 
 	assert.True(t, m[_p("2001::/56")])
 }
+
+func TestPrefixNumPrefixes(t *testing.T) {
+	tests := []struct {
+		description string
+		prefix      Prefix
+		length      uint32
+		count       uint64
+		error       bool
+	}{
+		{
+			description: "overflow",
+			length:      64,
+			error:       true,
+		}, {
+			description: "bad length",
+			length:      129,
+			error:       true,
+		}, {
+			description: "same size",
+			prefix:      _p("2001:db8::/64"),
+			length:      64,
+			count:       1,
+		}, {
+			description: "too big",
+			prefix:      _p("2001:db8::/64"),
+			length:      32,
+		}, {
+			description: "32",
+			prefix:      _p("2001:db8::/32"),
+			length:      64,
+			count:       0x100000000,
+		}, {
+			description: "56",
+			prefix:      _p("2001:db8::/56"),
+			length:      64,
+			count:       256,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.description, func(t *testing.T) {
+			count, err := tt.prefix.NumPrefixes(tt.length)
+			assert.Equal(t, tt.error, err != nil)
+			assert.Equal(t, tt.count, count)
+		})
+	}
+}

ipv6/range.go

@@ -47,6 +47,12 @@ func (me Range) Contains(other SetI) bool {
 	return me.Set().Contains(other)
 }
 
+// NumPrefixes returns the number of prefixes of the given prefix length in
+// this range.
+func (me Range) NumPrefixes(length uint32) (count uint64, err error) {
+	return me.Set().NumPrefixes(length)
+}
+
 // Minus returns a slice of ranges resulting from subtracting the given range
 // The slice will contain from 0 to 2 new ranges depending on how they overlap
 func (me Range) Minus(other Range) []Range {

ipv6/range_test.go

@@ -394,3 +394,33 @@ func TestRangeAsMapKey(t *testing.T) {
 
 	assert.True(t, m[_r(_a("2001::"), _a("2001::1000:0"))])
 }
+
+func TestRangeNumPrefixes(t *testing.T) {
+	tests := []struct {
+		description string
+		r           Range
+		length      uint32
+		count       uint64
+		error       bool
+	}{
+		{
+			description: "simple",
+			r:           _r(_a("2001::"), _a("2002::")),
+			length:      24,
+			count:       256,
+		}, {
+			description: "bonus",
+			r:           _r(_a("2001::"), _a("2001::1:0")),
+			length:      128,
+			count:       0x10001,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.description, func(t *testing.T) {
+			count, err := tt.r.NumPrefixes(tt.length)
+			assert.Equal(t, tt.error, err != nil)
+			assert.Equal(t, tt.count, count)
+		})
+	}
+}