use OneCount to decide uint size, add regression tests

Signed-off-by: jsign <[email protected]>

Author: jsign

uhamt.go

@@ -8,36 +8,20 @@ import (
 // indexForBitPos returns the index within the collapsed array corresponding to
 // the given bit in the bitset.  The collapsed array contains only one entry
 // per bit set in the bitfield, and this function is used to map the indices.
-func (n *Node) indexForBitPosOld(bp int) int {
-	// TODO: an optimization could reuse the same 'mask' here and change the size
-	//       as needed. This isnt yet done as the bitset package doesnt make it easy
-	//       to do.
-
-	// make a bitmask (all bits set) 'bp' bits long
-	mask := new(big.Int).Sub(new(big.Int).Exp(big.NewInt(2), big.NewInt(int64(bp)), nil), big.NewInt(1))
-	mask.And(mask, n.Bitfield)
-
-	return popCount(mask)
-}
-
-func popCount(i *big.Int) int {
-	var n int
-	for _, v := range i.Bits() {
-		n += bits.OnesCount64(uint64(v))
-	}
-	return n
+func (n *Node) indexForBitPos(bp int) int {
+	return indexForBitPos(bp, n.Bitfield)
 }
 
-func (n *Node) indexForBitPos(bp int) int {
+func indexForBitPos(bp int, bitfield *big.Int) int {
 	var x uint
 	var count, i int
-	w := n.Bitfield.Bits()
+	w := bitfield.Bits()
 	for x = uint(bp); x > bits.UintSize && i < len(w); x -= bits.UintSize {
-		count += bits.OnesCount64(uint64(w[i]))
+		count += bits.OnesCount(uint(w[i]))
 		i++
 	}
 	if i == len(w) {
 		return count
 	}
-	return count + bits.OnesCount64(uint64(w[i])&((1<<x)-1))
+	return count + bits.OnesCount(uint(w[i])&((1<<x)-1))
 }

uhamt_test.go

@@ -0,0 +1,57 @@
+package hamt
+
+import (
+	"math/big"
+	"math/bits"
+	"math/rand"
+	"testing"
+)
+
+func TestIndexForBitRandom(t *testing.T) {
+	t.Parallel()
+	r := rand.New(rand.NewSource(int64(42)))
+
+	count := 100000
+	slot := make([]byte, 32)
+	for i := 0; i < count; i++ {
+		_, err := r.Read(slot)
+		if err != nil {
+			t.Fatal("couldn't create random bitfield")
+		}
+		bi := big.NewInt(0).SetBytes(slot)
+		for k := 0; k < 256; k++ {
+			if indexForBitPos(k, bi) != indexForBitPosOriginal(k, bi) {
+				t.Fatalf("indexForBit doesn't match with original")
+			}
+		}
+	}
+}
+
+func TestIndexForBitLinear(t *testing.T) {
+	t.Parallel()
+	var i int64
+	for i = 0; i < 1<<16-1; i++ {
+		bi := big.NewInt(i)
+		for k := 0; k < 16; k++ {
+			if indexForBitPos(k, bi) != indexForBitPosOriginal(k, bi) {
+				t.Fatalf("indexForBit doesn't match with original")
+			}
+		}
+	}
+}
+
+// Original implementation of indexForBit, before #39.
+func indexForBitPosOriginal(bp int, bitfield *big.Int) int {
+	mask := new(big.Int).Sub(new(big.Int).Exp(big.NewInt(2), big.NewInt(int64(bp)), nil), big.NewInt(1))
+	mask.And(mask, bitfield)
+
+	return popCount(mask)
+}
+
+func popCount(i *big.Int) int {
+	var n int
+	for _, v := range i.Bits() {
+		n += bits.OnesCount64(uint64(v))
+	}
+	return n
+}