support bitmap

datastruct/bitmap/bitmap.go

@@ -0,0 +1,97 @@
+package bitmap
+
+type BitMap []byte
+
+func New() *BitMap {
+	b := BitMap(make([]byte, 0))
+	return &b
+}
+
+func toByteSize(bitSize int64) int64 {
+	if bitSize%8 == 0 {
+		return bitSize / 8
+	}
+	return bitSize/8 + 1
+}
+
+func (b *BitMap) grow(bitSize int64) {
+	byteSize := toByteSize(bitSize)
+	gap := byteSize - int64(len(*b))
+	if gap <= 0 {
+		return
+	}
+	*b = append(*b, make([]byte, gap)...)
+}
+
+func (b *BitMap) BitSize() int {
+	return len(*b) * 8
+}
+
+func FromBytes(bytes []byte) *BitMap {
+	bm := BitMap(bytes)
+	return &bm
+}
+
+func (b *BitMap) ToBytes() []byte {
+	return *b
+}
+
+func (b *BitMap) SetBit(offset int64, val byte) {
+	byteIndex := offset / 8
+	bitOffset := offset % 8
+	mask := byte(1 << bitOffset)
+	b.grow(offset + 1)
+	if val > 0 {
+		// set bit
+		(*b)[byteIndex] |= mask
+	} else {
+		// clear bit
+		(*b)[byteIndex] &^= mask
+	}
+}
+
+func (b *BitMap) GetBit(offset int64) byte {
+	byteIndex := offset / 8
+	bitOffset := offset % 8
+	if byteIndex >= int64(len(*b)) {
+		return 0
+	}
+	return ((*b)[byteIndex] >> bitOffset) & 0x01
+}
+
+type Callback func(offset int64, val byte) bool
+
+func (b *BitMap) ForEachBit(begin int64, end int64, cb Callback) {
+	offset := begin
+	byteIndex := offset / 8
+	bitOffset := offset % 8
+	for byteIndex < int64(len(*b)) {
+		b := (*b)[byteIndex]
+		for bitOffset < 8 {
+			bit := byte(b >> bitOffset & 0x01)
+			if !cb(offset, bit) {
+				return
+			}
+			bitOffset++
+			offset++
+		}
+		byteIndex++
+		bitOffset = 0
+		if end > 0 && offset == end {
+			break
+		}
+	}
+}
+
+func (b *BitMap) ForEachByte(begin int, end int, cb Callback) {
+	if end == 0 {
+		end = len(*b)
+	} else if end > len(*b) {
+		end = len(*b)
+	}
+	for i := begin; i < end; i++ {
+		if !cb(int64(i), (*b)[i]) {
+			return
+		}
+	}
+}

datastruct/bitmap/bitmap_test.go

@@ -0,0 +1,168 @@
+package bitmap
+
+import (
+	"bytes"
+	"math/rand"
+	"testing"
+)
+
+func TestSetBit(t *testing.T) {
+	// generate distinct rand offsets
+	size := 1000
+	offsets := make([]int64, size)
+	for i := 0; i < size; i++ {
+		offsets[i] = int64(i)
+	}
+	rand.Shuffle(size, func(i, j int) {
+		offsets[i], offsets[j] = offsets[j], offsets[i]
+	})
+	offsets = offsets[0 : size/5]
+	// set bit
+	offsetMap := make(map[int64]struct{})
+	bm := New()
+	for _, offset := range offsets {
+		offsetMap[offset] = struct{}{}
+		bm.SetBit(offset, 1)
+	}
+	// get bit
+	for i := 0; i < bm.BitSize(); i++ {
+		offset := int64(i)
+		_, expect := offsetMap[offset]
+		actual := bm.GetBit(offset) > 0
+		if expect != actual {
+			t.Errorf("wrong value at %d", offset)
+		}
+	}
+
+	bm2 := New()
+	bm2.SetBit(15, 1)
+	if bm2.GetBit(15) != 1 {
+		t.Error("wrong value")
+	}
+	if bm2.GetBit(16) != 0 {
+		t.Error("wrong value")
+	}
+}
+
+func TestFromBytes(t *testing.T) {
+	bs := []byte{0xff, 0xff}
+	bm := FromBytes(bs)
+	bm.SetBit(8, 0)
+	expect := []byte{0xff, 0xfe}
+	if !bytes.Equal(bs, expect) {
+		t.Error("wrong value")
+	}
+	ret := bm.ToBytes()
+	if !bytes.Equal(ret, expect) {
+		t.Error("wrong value")
+	}
+}
+
+func TestForEachBit(t *testing.T) {
+	bm := New()
+	for i := 0; i < 1000; i++ {
+		if i%2 == 0 {
+			bm.SetBit(int64(i), 1)
+		}
+	}
+	expectOffset := int64(100)
+	count := 0
+	bm.ForEachBit(100, 200, func(offset int64, val byte) bool {
+		if offset != expectOffset {
+			t.Error("wrong offset")
+		}
+		expectOffset++
+		if offset%2 == 0 && val == 0 {
+			t.Error("expect 1")
+		}
+		count++
+		return true
+	})
+	if count != 100 {
+		t.Error("wrong count")
+	}
+	bm = New()
+	size := 1000
+	offsets := make([]int64, size)
+	for i := 0; i < size; i++ {
+		offsets[i] = int64(i)
+	}
+	rand.Shuffle(size, func(i, j int) {
+		offsets[i], offsets[j] = offsets[j], offsets[i]
+	})
+	offsets = offsets[0 : size/5]
+	offsetMap := make(map[int64]struct{})
+	for _, offset := range offsets {
+		bm.SetBit(offset, 1)
+		offsetMap[offset] = struct{}{}
+	}
+	bm.ForEachBit(int64(size/20+1), 0, func(offset int64, val byte) bool {
+		_, expect := offsetMap[offset]
+		actual := bm.GetBit(offset) > 0
+		if expect != actual {
+			t.Errorf("wrong value at %d", offset)
+		}
+		return true
+	})
+	count = 0
+	bm.ForEachBit(0, 0, func(offset int64, val byte) bool {
+		count++
+		return false
+	})
+	if count != 1 {
+		t.Error("break failed")
+	}
+}
+
+func TestBitMap_ForEachByte(t *testing.T) {
+	bm := New()
+	for i := 0; i < 1000; i++ {
+		if i%16 == 0 {
+			bm.SetBit(int64(i), 1)
+		}
+	}
+	bm.ForEachByte(0, 0, func(offset int64, val byte) bool {
+		if offset%2 == 0 {
+			if val != 1 {
+				t.Error("wrong value")
+			}
+		} else {
+			if val != 0 {
+				t.Error("wrong value")
+			}
+		}
+		return true
+	})
+	bm.ForEachByte(0, 2000, func(offset int64, val byte) bool {
+		if offset%2 == 0 {
+			if val != 1 {
+				t.Error("wrong value")
+			}
+		} else {
+			if val != 0 {
+				t.Error("wrong value")
+			}
+		}
+		return true
+	})
+	bm.ForEachByte(0, 500, func(offset int64, val byte) bool {
+		if offset%2 == 0 {
+			if val != 1 {
+				t.Error("wrong value")
+			}
+		} else {
+			if val != 0 {
+				t.Error("wrong value")
+			}
+		}
+		return true
+	})
+	count := 0
+	bm.ForEachByte(0, 0, func(offset int64, val byte) bool {
+		count++
+		return false
+	})
+	if count != 1 {
+		t.Error("break failed")
+	}
+}