BREAKING CHANGE: introduce Bitmap for handling bitfield operations

Closes: https://github.com/ipfs/go-hamt-ipld/issues/54

bitmap.go

@@ -0,0 +1,118 @@
+package hamt
+
+import (
+	"fmt"
+	"math"
+	"math/bits"
+)
+
+// Bitmap is a managed bitmap, primarily for the purpose of tracking the
+// presence or absence of elements in an associated array. It can set and unset
+// individual bits and perform limited popcount for a given index to calculate
+// the position in the associated compacted array.
+type Bitmap struct {
+	Bytes []byte
+}
+
+// NewBitmap creates a new bitmap for a given bitWidth. The bitmap will hold
+// 2^bitWidth bytes.
+func NewBitmap(bitWidth int) *Bitmap {
+	bc := (1 << uint(bitWidth)) / 8
+	if bc == 0 {
+		panic("bitWidth too small")
+	}
+
+	return NewBitmapFrom(make([]byte, bc))
+}
+
+// NewBitmapFrom creates a new Bitmap from an existing byte array. It is
+// assumed that bytes is the correct length for the bitWidth of this Bitmap.
+func NewBitmapFrom(bytes []byte) *Bitmap {
+	if len(bytes) == 0 {
+		panic("can't form Bitmap from zero bytes")
+	}
+	bm := Bitmap{Bytes: bytes}
+	return &bm
+}
+
+// BitWidth calculates the bitWidth of this Bitmap by performing a
+// log2(bits). The bitWidth is the minimum number of bits required to
+// form indexes that address all of this Bitmap. e.g. a bitWidth of 5 can form
+// indexes of 0 to 31, i.e. 4 bytes.
+func (bm *Bitmap) BitWidth() int {
+	return int(math.Log2(float64(len(bm.Bytes) * 8)))
+}
+
+func (bm *Bitmap) bindex(in int) int {
+	// Return `in` to flip the byte addressing order to LE. For BE we address
+	// from the last byte backward.
+	bi := len(bm.Bytes) - 1 - in
+	if bi > len(bm.Bytes) || bi < 0 {
+		panic(fmt.Sprintf("invalid index for this Bitmap (index: %v, bytes: %v)", in, len(bm.Bytes)))
+	}
+	return bi
+}
+
+// IsSet indicates whether the bit at the provided position is set or not.
+func (bm *Bitmap) IsSet(position int) bool {
+	byt := bm.bindex(position / 8)
+	offset := position % 8
+	return (bm.Bytes[byt]>>offset)&1 == 1
+}
+
+// Set sets or unsets the bit at the given position according. If set is true,
+// the bit will be set. If set is false, the bit will be unset. Returns a
+// reference to this Bitmap.
+func (bm *Bitmap) Set(position int, set bool) *Bitmap {
+	has := bm.IsSet(position)
+	byt := bm.bindex(position / 8)
+	offset := position % 8
+
+	if set && !has {
+		bm.Bytes[byt] |= 1 << offset
+	} else if !set && has {
+		bm.Bytes[byt] ^= 1 << offset
+	}
+
+	return bm
+}
+
+// Index performs a limited popcount up to the given position. This calculates
+// the number of set bits up to the index of the bitmap. Useful for calculating
+// the position of an element in an associated compacted array.
+func (bm *Bitmap) Index(position int) int {
+	t := 0
+	eb := position / 8
+	byt := 0
+	for ; byt < eb; byt++ {
+		// quick popcount for the full bytes
+		t += bits.OnesCount(uint(bm.Bytes[bm.bindex(byt)]))
+	}
+	eb = eb * 8
+	if position > eb {
+		for i := byt * 8; i < position; i++ {
+			// manual per-bit check for the remainder <8 bits
+			if bm.IsSet(i) {
+				t++
+			}
+		}
+	}
+	return t
+}
+
+// Copy creates a clone of the Bitmap, creating a new byte array with the same
+// contents as the original.
+func (bm *Bitmap) Copy() *Bitmap {
+	ba := make([]byte, len(bm.Bytes))
+	copy(ba, bm.Bytes)
+	return NewBitmapFrom(ba)
+}
+
+// BitsSetCount counts how many bits are set in the bitmap.
+func (bm *Bitmap) BitsSetCount() int {
+	count := 0
+	for _, b := range bm.Bytes {
+		count += bits.OnesCount(uint(b))
+	}
+	return count
+}

bitmap_test.go

@@ -0,0 +1,230 @@
+package hamt
+
+import (
+	"bytes"
+	"testing"
+)
+
+// many cases taken from https://github.com/rvagg/iamap/blob/fad95295b013c8b4f0faac6dd5d9be175f6e606c/test/bit-utils-test.js
+// but rev() is used to reverse the data in most instances
+
+// reverse for BE format
+func rev(in []byte) []byte {
+	out := make([]byte, len(in))
+	for i := 0; i < len(in); i++ {
+		out[len(in)-1-i] = in[i]
+	}
+	return out
+}
+
+func TestBitmapHas(t *testing.T) {
+	type tcase struct {
+		bytes []byte
+		pos   int
+		set   bool
+	}
+	cases := []tcase{
+		{b(0b0), 0, false},
+		{b(0b1), 0, true},
+		{b(0b101010), 2, false},
+		{b(0b101010), 3, true},
+		{b(0b101010), 4, false},
+		{b(0b101010), 5, true},
+		{b(0b100000), 5, true},
+		{b(0b0100000), 5, true},
+		{b(0b00100000), 5, true},
+		{[]byte{0x0, 0b00100000}, 8 + 5, true},
+		{[]byte{0x0, 0x0, 0b00100000}, 8*2 + 5, true},
+		{[]byte{0x0, 0x0, 0x0, 0b00100000}, 8*3 + 5, true},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0b00100000}, 8*4 + 5, true},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0x0, 0b00100000}, 8*5 + 5, true},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0x0, 0b00100000}, 8*4 + 5, false},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0x0, 0b00100000}, 8*3 + 5, false},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0x0, 0b00100000}, 8*2 + 5, false},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0x0, 0b00100000}, 8 + 5, false},
+		{[]byte{0x0, 0x0, 0x0, 0x0, 0x0, 0b00100000}, 5, false},
+	}
+
+	for _, c := range cases {
+		bm := NewBitmapFrom(rev(c.bytes))
+		if bm.IsSet(c.pos) != c.set {
+			t.Fatalf("bitmap %v IsSet(%v) should be %v", c.bytes, c.pos, c.set)
+		}
+	}
+}
+
+func TestBitmapBitWidth(t *testing.T) {
+	for i := 3; i <= 16; i++ {
+		if NewBitmap(i).BitWidth() != i {
+			t.Fatal("incorrect bitWidth calculation")
+		}
+		if NewBitmapFrom(make([]byte, (1<<i)/8)).BitWidth() != i {
+			t.Fatal("incorrect bitWidth calculation")
+		}
+	}
+}
+
+func TestBitmapIndex(t *testing.T) {
+	type tcase struct {
+		bytes    []byte
+		pos      int
+		expected int
+	}
+	cases := []tcase{
+		{b(0b111111), 0, 0},
+		{b(0b111111), 1, 1},
+		{b(0b111111), 2, 2},
+		{b(0b111111), 4, 4},
+		{b(0b111100), 2, 0},
+		{b(0b111101), 4, 3},
+		{b(0b111001), 4, 2},
+		{b(0b111000), 4, 1},
+		{b(0b110000), 4, 0},
+		{b(0b000000), 0, 0},
+		{b(0b000000), 1, 0},
+		{b(0b000000), 2, 0},
+		{b(0b000000), 3, 0},
+		{[]byte{0x0, 0x0, 0x0}, 20, 0},
+		{[]byte{0xff, 0xff, 0xff}, 5, 5},
+		{[]byte{0xff, 0xff, 0xff}, 7, 7},
+		{[]byte{0xff, 0xff, 0xff}, 8, 8},
+		{[]byte{0xff, 0xff, 0xff}, 10, 10},
+		{[]byte{0xff, 0xff, 0xff}, 20, 20},
+	}
+
+	for _, c := range cases {
+		bm := NewBitmapFrom(rev(c.bytes))
+		if bm.Index(c.pos) != c.expected {
+			t.Fatalf("bitmap %v Index(%v) should be %v", c.bytes, c.pos, c.expected)
+		}
+	}
+}
+
+func TestBitmap_32bitFixed(t *testing.T) {
+	// a 32-byte bitmap and a list of all the bits that are set
+	byts := []byte{
+		0b00100101, 0b10000000, 0b00000000, 0b01000000,
+		0b00000000, 0b01000000, 0b00000000, 0b01000000,
+		0b00000000, 0b00100000, 0b00000000, 0b01000000,
+		0b00000000, 0b00010000, 0b00000000, 0b01000000,
+		0b00000000, 0b00001000, 0b00000000, 0b01000000,
+		0b00000000, 0b00000100, 0b00000000, 0b01000000,
+		0b00000000, 0b00000010, 0b00000000, 0b01000000,
+		0b00000000, 0b00000001, 0b00000000, 0b01000000,
+	}
+	bm := NewBitmapFrom(rev(byts))
+	set := []int{
+		0, 2, 5, 8 + 7, 8*3 + 6,
+		8*5 + 6, 8*7 + 6,
+		8*9 + 5, 8*11 + 6,
+		8*13 + 4, 8*15 + 6,
+		8*17 + 3, 8*19 + 6,
+		8*21 + 2, 8*23 + 6,
+		8*25 + 1, 8*27 + 6,
+		8 * 29, 8*31 + 6}
+
+	c := 0
+	for i := 0; i < 256; i++ {
+		if c < len(set) && i == set[c] {
+			if !bm.IsSet(i) {
+				t.Fatalf("IsSet(%v) should be true", i)
+			}
+			// the index c of `set` also gives us the translation of Index(i)
+			if bm.Index(i) != c {
+				t.Fatalf("Index(%v) should be %v", i, c)
+			}
+			c++
+		} else {
+			if bm.IsSet(i) {
+				t.Fatalf("IsSet(%v) should be false", i)
+			}
+		}
+	}
+}
+
+func TestBitmapSetBytes(t *testing.T) {
+	if !bytes.Equal(NewBitmap(3).Set(0, true).Bytes, rev([]byte{0b00000001})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmap(3).Set(1, true).Bytes, rev([]byte{0b00000010})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmap(3).Set(7, true).Bytes, rev([]byte{0b10000000})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11111111})).Set(0, true).Bytes, rev([]byte{0b11111111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11111111})).Set(7, true).Bytes, rev([]byte{0b11111111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b01010101})).Set(1, true).Bytes, rev([]byte{0b01010111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b01010101})).Set(7, true).Bytes, rev([]byte{0b11010101})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11111111})).Set(0, false).Bytes, rev([]byte{0b11111110})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11111111})).Set(1, false).Bytes, rev([]byte{0b11111101})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11111111})).Set(7, false).Bytes, rev([]byte{0b01111111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b11111111})).Set(8+0, true).Bytes, rev([]byte{0, 0b11111111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b11111111})).Set(8+7, true).Bytes, rev([]byte{0, 0b11111111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b01010101})).Set(8+1, true).Bytes, rev([]byte{0, 0b01010111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b01010101})).Set(8+7, true).Bytes, rev([]byte{0, 0b11010101})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b11111111})).Set(8+0, false).Bytes, rev([]byte{0, 0b11111110})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b11111111})).Set(8+1, false).Bytes, rev([]byte{0, 0b11111101})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0b11111111})).Set(8+7, false).Bytes, rev([]byte{0, 0b01111111})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0})).Set(0, false).Bytes, rev([]byte{0b00000000})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0})).Set(7, false).Bytes, rev([]byte{0b00000000})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b01010101})).Set(0, false).Bytes, rev([]byte{0b01010100})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b01010101})).Set(6, false).Bytes, rev([]byte{0b00010101})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})).Set(0, false).Bytes, rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})).Set(0, true).Bytes, rev([]byte{0b11000011, 0b11010010, 0b01001010, 0b0000001})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})).Set(12, false).Bytes, rev([]byte{0b11000010, 0b11000010, 0b01001010, 0b0000001})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})).Set(12, true).Bytes, rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})).Set(24, false).Bytes, rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000000})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})).Set(24, true).Bytes, rev([]byte{0b11000010, 0b11010010, 0b01001010, 0b0000001})) {
+		t.Fatal("Failed bytes comparison")
+	}
+	if !bytes.Equal(NewBitmapFrom(rev([]byte{0, 0, 0, 0})).Set(31, true).Bytes, rev([]byte{0, 0, 0, 0b10000000})) {
+		t.Fatal("Failed bytes comparison")
+	}
+}