index for bit optimization

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

Author: jsign

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

@@ -8,7 +8,7 @@ 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) indexForBitPos(bp int) int {
+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.
@@ -27,3 +27,17 @@ func popCount(i *big.Int) int {
 	}
 	return n
 }
+
+func (n *Node) indexForBitPos(bp int) int {
+	var x uint64
+	var count, i int
+	w := n.Bitfield.Bits()
+	for x = uint64(bp); x > 64 && i < len(w); x -= 64 {
+		count += bits.OnesCount64(uint64(w[i]))
+		i++
+	}
+	if i == len(w) {
+		return count
+	}
+	return count + bits.OnesCount64(uint64(w[i])&((1<<x)-1))
+}