增加bitmap

ellyn_common/collections/bitmap.go

@@ -0,0 +1,49 @@
+package collections
+
+const slotShift = 6
+
+type bitMap struct {
+	slots []uint64
+	cap   uint
+	size  int
+}
+
+func NewBitMap(cap uint) *bitMap {
+	return &bitMap{
+		slots: make([]uint64, ((uint64(cap)-1)>>slotShift)+1),
+		cap:   cap,
+	}
+}
+
+func (m *bitMap) Set(pos uint) {
+	m.checkPos(pos)
+	old := m.slots[pos>>slotShift]
+	m.slots[pos>>slotShift] |= uint64(1) << pos
+	if m.slots[pos>>slotShift] != old {
+		m.size++
+	}
+}
+
+func (m *bitMap) Get(pos uint) bool {
+	m.checkPos(pos)
+	return m.slots[pos>>slotShift]&(uint64(1)<<pos) != 0
+}
+
+func (m *bitMap) Clear(pos uint) {
+	m.checkPos(pos)
+	old := m.slots[pos>>slotShift]
+	m.slots[pos>>slotShift] &= ^(uint64(1) << pos)
+	if m.slots[pos>>slotShift] != old {
+		m.size--
+	}
+}
+
+func (m *bitMap) Size() int {
+	return m.size
+}
+
+func (m *bitMap) checkPos(pos uint) {
+	if pos >= m.cap {
+		panic("Position exceeds bitmap length")
+	}
+}

ellyn_common/collections/bitmap_test.go

@@ -0,0 +1,39 @@
+package collections
+
+import (
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestNewBitMap(t *testing.T) {
+	m := NewBitMap(1)
+	require.Equal(t, 1, len(m.slots))
+	m = NewBitMap(64)
+	require.Equal(t, 1, len(m.slots))
+	m = NewBitMap(65)
+	require.Equal(t, 2, len(m.slots))
+	m = NewBitMap(128)
+	require.Equal(t, 2, len(m.slots))
+	m = NewBitMap(129)
+	require.Equal(t, 3, len(m.slots))
+}
+
+func TestBitMapBasic(t *testing.T) {
+	m := NewBitMap(10)
+	m.Set(1)
+	m.Set(2)
+	require.True(t, m.Get(1))
+	require.True(t, m.Get(2))
+	require.Equal(t, 2, m.Size())
+	require.False(t, m.Get(3))
+	m.Clear(1)
+	require.False(t, m.Get(1))
+	require.Equal(t, 1, m.Size())
+}
+
+func BenchmarkBitMap(b *testing.B) {
+	m := NewBitMap(1024)
+	for i := 0; i < b.N; i++ {
+		m.Set(uint(i) & 1023)
+	}
+}