Merge pull request #39 from jsign/setopt

indexForBit optimization

Author: Stebalien

hamt_bench_test.go

@@ -3,6 +3,7 @@ package hamt
 import (
 	"context"
 	"encoding/hex"
+	"fmt"
 	"math/rand"
 	"runtime"
 	"testing"
@@ -47,6 +48,40 @@ func BenchmarkSerializeNode(b *testing.B) {
 	}
 }
 
+type benchSetCase struct {
+	count    int
+	bitwidth int
+}
+
+func BenchmarkSet(b *testing.B) {
+	kCounts := []int{1, 10, 100}
+	bitwidths := []int{5, 8}
+
+	var table []benchSetCase
+	for _, c := range kCounts {
+
+		for _, bw := range bitwidths {
+			table = append(table, benchSetCase{count: c * 1000, bitwidth: bw})
+		}
+
+	}
+	r := rander{rand.New(rand.NewSource(int64(42)))}
+	for _, t := range table {
+		b.Run(fmt.Sprintf("%d/%d", t.count, t.bitwidth), func(b *testing.B) {
+			ctx := context.Background()
+			n := NewNode(NewCborStore(), UseTreeBitWidth(t.bitwidth))
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				for j := 0; j < t.count; j++ {
+					if err := n.Set(ctx, r.randString(), r.randValue()); err != nil {
+						b.Fatal(err)
+					}
+				}
+			}
+		})
+	}
+}
+
 func BenchmarkFind(b *testing.B) {
 	b.Run("find-10k", doBenchmarkEntriesCount(10000, 8))
 	b.Run("find-100k", doBenchmarkEntriesCount(100000, 8))

uhamt.go

@@ -9,21 +9,19 @@ import (
 // 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) indexForBitPos(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)
+	return indexForBitPos(bp, n.Bitfield)
 }
 
-func popCount(i *big.Int) int {
-	var n int
-	for _, v := range i.Bits() {
-		n += bits.OnesCount64(uint64(v))
+func indexForBitPos(bp int, bitfield *big.Int) int {
+	var x uint
+	var count, i int
+	w := bitfield.Bits()
+	for x = uint(bp); x > bits.UintSize && i < len(w); x -= bits.UintSize {
+		count += bits.OnesCount(uint(w[i]))
+		i++
+	}
+	if i == len(w) {
+		return count
 	}
-	return n
+	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
+}