Fixes for block handling

ip/network/blocks.go

@@ -2,14 +2,15 @@ package network
 
 import (
 	"math"
+	"math/bits"
 
 	"github.com/ipfreely-uk/go/ip"
 	"github.com/ipfreely-uk/go/ip/compare"
 	"github.com/ipfreely-uk/go/ip/subnet"
 )
 
 // TODO: can replace with constant
-var LOG_2 = math.Log2(2.0)
+var LOG_2 = math.Log(2.0)
 
 // Subdivides range into valid CIDR blocks
 func Blocks[A ip.Address[A]](r Range[A]) Iterator[Block[A]] {
@@ -33,15 +34,13 @@ func blockIterator[A ip.Address[A]](start, end A) Iterator[Block[A]] {
 		if done {
 			return false, nil
 		}
-		max := maxMask(current)
-		size := ip.Next(end.Subtract(start))
+		maxSize := maxMask(current)
+		size := ip.Next(end.Subtract(current))
 		x := log(size) / LOG_2
 		width := start.Family().Width()
 		maxDiff := int(width - int(math.Floor(x)))
-		if max > maxDiff {
-			max = maxDiff
-		}
-		block := NewBlock(current, max)
+		mask := max(maxSize, maxDiff)
+		block := NewBlock(current, mask)
 		last := block.Last()
 		if compare.Eq(last, end) {
 			done = true
@@ -52,43 +51,63 @@ func blockIterator[A ip.Address[A]](start, end A) Iterator[Block[A]] {
 	}
 }
 
+func max(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}
+
 func log[A ip.Address[A]](address A) float64 {
 	MAX_DIGITS_2 := 977
 
-	bi := ip.ToBigInt(address)
-	blex := bi.BitLen() - MAX_DIGITS_2
+	bitlen := bitLen(address)
+	blex := bitlen - MAX_DIGITS_2
+	a := address
 	if blex > 0 {
-		a := address.Shift(blex)
-		bi = ip.ToBigInt(a)
+		a = address.Shift(blex)
 	}
-	double, _ := bi.Float64()
-	res := math.Log2(double)
-	if res > 0 {
+	double := toFloat64(a)
+	res := math.Log(double)
+	if blex > 0 {
 		res = res + float64(blex)*LOG_2
 	}
 	return res
 }
 
+func toFloat64[A ip.Address[A]](address A) float64 {
+	bi := ip.ToBigInt(address)
+	f, _ := bi.Float64()
+	return f
+}
+
+func bitLen[A ip.Address[A]](address A) int {
+	// TODO: replace all this with leading/trailing zero func
+	bytes := address.Bytes()
+	lead0Bits := 0
+	for _, b := range bytes {
+		if b == 0 {
+			lead0Bits = lead0Bits + 8
+		} else {
+			lead0Bits = lead0Bits + bits.LeadingZeros8(b)
+			break
+		}
+	}
+	return address.Family().Width() - lead0Bits
+}
+
 func maxMask[A ip.Address[A]](address A) int {
+	// TODO: replace all this with leading/trailing zero func
 	bytes := address.Bytes()
-	bits := address.Family().Width()
+	mask := address.Family().Width()
 	for i := len(bytes) - 1; i >= 0; i-- {
 		b := bytes[i]
 		if b == 0 {
-			bits = bits - 8
+			mask = mask - 8
 		} else {
-			n := 0
-			var m byte = 1
-			for j := 0; j <= 8; j++ {
-				if b&m != 0 {
-					n = j
-					break
-				}
-				m = m << 1
-			}
-			bits = bits - n
+			mask = mask - bits.TrailingZeros8(b)
 			break
 		}
 	}
-	return bits
+	return mask
 }

ip/network/blocks_test.go

@@ -0,0 +1,138 @@
+package network_test
+
+import (
+	"testing"
+
+	"github.com/ipfreely-uk/go/ip"
+	"github.com/ipfreely-uk/go/ip/network"
+	"github.com/ipfreely-uk/go/ip/subnet"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBlocks(t *testing.T) {
+	{
+		family := ip.V6()
+		expectedFirst0, _ := ip.Parse(family, "fe80::")
+		mask := subnet.Mask(family, 64)
+		inverse := mask.Not()
+		expectedLast0 := inverse.Or(expectedFirst0)
+		expectedFirst1 := ip.Next(expectedLast0)
+		expectedLast1 := expectedFirst1
+		input := network.NewRange(expectedFirst0, expectedLast1)
+
+		nextBlock := network.Blocks(input)
+
+		ok, actual := nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst0, actual.First())
+		assert.Equal(t, expectedLast0, actual.Last())
+
+		ok, actual = nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst1, actual.First())
+		assert.Equal(t, expectedLast1, actual.Last())
+
+		ok, _ = nextBlock()
+		assert.False(t, ok)
+	}
+	{
+		family := ip.V6()
+		expectedFirst0, _ := ip.Parse(family, "f000::")
+		mask := subnet.Mask(family, 8)
+		inverse := mask.Not()
+		expectedLast0 := inverse.Or(expectedFirst0)
+		expectedFirst1 := ip.Next(expectedLast0)
+		expectedLast1 := expectedFirst1
+		input := network.NewRange(expectedFirst0, expectedLast1)
+
+		nextBlock := network.Blocks(input)
+
+		ok, actual := nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst0, actual.First())
+		assert.Equal(t, expectedLast0, actual.Last())
+
+		ok, actual = nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst1, actual.First())
+		assert.Equal(t, expectedLast1, actual.Last())
+
+		ok, _ = nextBlock()
+		assert.False(t, ok)
+	}
+	{
+		family := ip.V6()
+		expectedFirst0, _ := ip.Parse(family, "::1")
+		mask := subnet.Mask(family, 128)
+		inverse := mask.Not()
+		expectedLast0 := inverse.Or(expectedFirst0)
+		expectedFirst1 := ip.Next(expectedLast0)
+		expectedLast1 := expectedFirst1
+		input := network.NewRange(expectedFirst0, expectedLast1)
+
+		nextBlock := network.Blocks(input)
+
+		ok, actual := nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst0, actual.First())
+		assert.Equal(t, expectedLast0, actual.Last())
+
+		ok, actual = nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst1, actual.First())
+		assert.Equal(t, expectedLast1, actual.Last())
+
+		ok, _ = nextBlock()
+		assert.False(t, ok)
+	}
+	{
+		family := ip.V4()
+		expectedFirst0, _ := ip.Parse(family, "10.0.0.0")
+		mask := subnet.Mask(family, 8)
+		inverse := mask.Not()
+		expectedLast0 := inverse.Or(expectedFirst0)
+		expectedFirst1 := ip.Next(expectedLast0)
+		expectedLast1 := expectedFirst1
+		input := network.NewRange(expectedFirst0, expectedLast1)
+
+		nextBlock := network.Blocks(input)
+
+		ok, actual := nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst0, actual.First())
+		assert.Equal(t, expectedLast0, actual.Last())
+
+		ok, actual = nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst1, actual.First())
+		assert.Equal(t, expectedLast1, actual.Last())
+
+		ok, _ = nextBlock()
+		assert.False(t, ok)
+	}
+	{
+		family := ip.V4()
+		expectedFirst0, _ := ip.Parse(family, "127.0.0.1")
+		mask := subnet.Mask(family, 32)
+		inverse := mask.Not()
+		expectedLast0 := inverse.Or(expectedFirst0)
+		expectedFirst1 := ip.Next(expectedLast0)
+		expectedLast1 := expectedFirst1
+		input := network.NewRange(expectedFirst0, expectedLast1)
+
+		nextBlock := network.Blocks(input)
+
+		ok, actual := nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst0, actual.First())
+		assert.Equal(t, expectedLast0, actual.Last())
+
+		ok, actual = nextBlock()
+		assert.True(t, ok)
+		assert.Equal(t, expectedFirst1, actual.First())
+		assert.Equal(t, expectedLast1, actual.Last())
+
+		ok, _ = nextBlock()
+		assert.False(t, ok)
+	}
+}