some cleanup

Author: c2akula

ba.go

@@ -16,16 +16,24 @@ const (
 // Bits externally represented as `bool` are stored internally as `uint64`s.
 // The total number of bits stored is set at creation and is immutable.
 type BitArray struct {
+	// buf is a backing array that bits writes into by default when the no. of bits requested to allocate is
+	// < 512. Only if more is asked, we'll skip buf and allocate directly into bits
+	buf  [8]Bit
 	bits []Bit
 	n    int // no. of bits
 }
 
-// New creates a new BitArray of `n` bits.
-func New(n int) BitArray {
-	return BitArray{
-		bits: make([]Bit, int(math.Ceil(float64(n)/64))),
-		n:    n,
+// New creates a new BitArray of `n` bits. If n <= 512, no allocation is done.
+func New(n int) (ba BitArray) {
+	nblk := nbitsToNblks(n)
+	ba.n = n
+	ba.bits = ba.buf[:]
+	if nblk <= 8 {
+		ba.bits = ba.buf[:nblk]
+		return
 	}
+	ba.bits = append(ba.bits, make([]Bit, nblk-len(ba.buf))...)
+	return
 }
 
 // Copy copies src into dst.
@@ -65,16 +73,8 @@ func FromUint64(u uint64) BitArray {
 // Size returns the no. of bits stored.
 func (ba *BitArray) Size() int { return ba.n }
 
-func biandsi(k int) (uint64, uint64) {
-	i := uint64(k)
-	return i / 64, i % 64
-}
-
-func (ba *BitArray) set1(bi, si uint64) { ba.bits[bi] |= 1 << si }
-func (ba *BitArray) set0(bi, si uint64) { ba.bits[bi] &= ^(1 << si) }
-
 // Set sets the bit at position k.
-func (ba *BitArray) Set(k int) { ba.set1(biandsi(k)) }
+func (ba *BitArray) Set(k int) { bi, si := biandsi(k); set(&ba.bits[bi], si) }
 
 // SetAll sets all the bits.
 func (ba *BitArray) SetAll() {
@@ -84,7 +84,7 @@ func (ba *BitArray) SetAll() {
 }
 
 // Clr clears the bit at position k.
-func (ba *BitArray) Clr(k int) { ba.set0(biandsi(k)) }
+func (ba *BitArray) Clr(k int) { bi, si := biandsi(k); clr(&ba.bits[bi], si) }
 
 // ClrAll clears all the bits.
 func (ba *BitArray) ClrAll() {
@@ -93,25 +93,24 @@ func (ba *BitArray) ClrAll() {
 	}
 }
 
-func (ba *BitArray) chk(bi, si uint64) bool { return (ba.bits[bi]>>si)&1 > 0 }
-
 // ChkSet returns the value of the bit at position k before setting it.
 func (ba *BitArray) ChkSet(k int) (b bool) {
 	bi, si := biandsi(k)
-	b = ba.chk(bi, si)
+	u := &ba.bits[bi]
+	b = chk(*u, si) != 0
 	if !b {
-		ba.set1(bi, si)
+		set(u, si)
 	}
 	return
-
 }
 
 // ChkClr returns the value of the bit at position k before clearing it.
 func (ba *BitArray) ChkClr(k int) (b bool) {
 	bi, si := biandsi(k)
-	b = ba.chk(bi, si)
+	u := &ba.bits[bi]
+	b = chk(*u, si) != 0
 	if b {
-		ba.set0(bi, si)
+		clr(u, si)
 	}
 	return
 }
@@ -125,29 +124,33 @@ func (ba *BitArray) Tgl(k int) {
 // Cnt returns the number of set bits.
 func (ba *BitArray) Cnt() (n int) {
 	for _, b := range ba.bits {
-		n += bits.OnesCount64(uint64(b))
+		n += bits.OnesCount64(b)
 	}
 	return
 }
 
 // Chk returns the value of the bit at position k.
 func (ba *BitArray) Chk(k int) bool {
 	bi, si := biandsi(k)
-	return (ba.bits[bi]>>si)&1 > 0
+	return chk(ba.bits[bi], si) != 0
 }
 
 // Put sets the value of the bit at position k to v.
 func (ba *BitArray) Put(k int, v Bit) {
 	bi, si := biandsi(k)
-	ba.bits[bi] = (ba.bits[bi] & ^(1 << si)) | (v << si)
+	put(&ba.bits[bi], si, v)
 }
 
 // Swap swaps the value of bit at position k with v. On return, v contains the old value.
-func (ba *BitArray) Swap(k int, v *Bit) {
+func (ba *BitArray) Swap(k int, b *Bit) {
 	bi, si := biandsi(k)
-	ob := (ba.bits[bi] >> si) & 1
-	ba.bits[bi] = (ba.bits[bi] & ^(1 << si)) | (*v << si)
-	*v = ob
+	t := &ba.bits[bi]
+	ob := chk(*t, si)
+	if ob == *b {
+		return
+	}
+	put(t, si, ob)
+	*b = ob
 }
 
 func (ba *BitArray) String() string {
@@ -160,3 +163,15 @@ func (ba *BitArray) String() string {
 	}
 	return string(sb)
 }
+
+func nbitsToNblks(n int) int { return int(math.Ceil(float64(n) / 64)) }
+
+func set(u *uint64, si uint64)        { *u |= 1 << si }
+func clr(u *uint64, si uint64)        { *u &= ^(1 << si) }
+func chk(u uint64, si uint64) Bit     { return (u >> si) & 1 }
+func put(u *uint64, si uint64, b Bit) { *u = (*u & ^(1 << si)) | (b << si) }
+
+func biandsi(k int) (uint64, uint64) {
+	i := uint64(k)
+	return i / 64, i % 64
+}

ba_test.go

@@ -25,6 +25,8 @@ func TestNew(t *testing.T) {
 		{65, 2},
 		{256, 4},
 		{257, 5},
+		{512, 8},
+		{768, 12},
 	}
 
 	for i, tt := range tests {
@@ -269,10 +271,19 @@ func TestBitArray(t *testing.T) {
 }
 
 func BenchmarkNew(b *testing.B) {
-	b.ReportAllocs()
-	for i := 0; i < b.N; i++ {
-		New(256)
-	}
+	b.Run("<=512", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			New(512)
+		}
+	})
+
+	b.Run(">512", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			New(768)
+		}
+	})
 }
 
 func BenchmarkBitArray(b *testing.B) {
@@ -397,7 +408,6 @@ func BenchmarkBitArray_Cnt(b *testing.B) {
 			ba.Cnt()
 		}
 	})
-
 }
 
 func BenchmarkBiandSi(b *testing.B) {