Improve bitvector performance

License: MIT
Signed-off-by: Jakub Sztandera <kubuxu@protonmail.ch>

lib/rlepluslazy/internal/bitvector.go

@@ -121,6 +121,22 @@ func (v *BitVector) Take(index uint, count uint, order BitNumbering) (out byte)
 		log.Panicf("invalid count")
 	}
 
+	if order == LSB0 && v.BytePacking == LSB0 {
+		x := index >> 3
+		r := index & 7
+		var o uint16
+		l := uint(len(v.Buf))
+
+		if x+1 < l {
+			o = uint16(v.Buf[x]) | uint16(v.Buf[x+1])<<8
+		} else if x < l {
+			o = uint16(v.Buf[x])
+		}
+
+		o = o >> r
+		return byte(o & (0xFF >> (8 - count)))
+	}
+
 	for i := uint(0); i < count; i++ {
 		val, _ := v.Get(index + i)
 
@@ -134,6 +150,18 @@ func (v *BitVector) Take(index uint, count uint, order BitNumbering) (out byte)
 	return
 }
 
+var masks = [9]byte{
+	0x0,
+	0x1,
+	0x3,
+	0x7,
+	0xF,
+	0x1F,
+	0x3F,
+	0x7F,
+	0xFF,
+}
+
 // Iterator returns a function, which when invoked, returns the number
 // of bits requested, and increments an internal cursor.
 //
@@ -141,14 +169,60 @@ func (v *BitVector) Take(index uint, count uint, order BitNumbering) (out byte)
 //
 // Panics if count is out of range
 func (v *BitVector) Iterator(order BitNumbering) func(uint) byte {
-	cursor := uint(0)
-	return func(count uint) (out byte) {
-		if count > 8 {
-			log.Panicf("invalid count")
+	if order == LSB0 && v.BytePacking == LSB0 {
+		// Here be dragons
+		// This is about 10x faster
+		index := int(0)
+		bitIdx := uint(0)
+
+		var rest uint64
+		for n := 7; n >= 0; n-- {
+			var o uint64
+			if len(v.Buf) > n {
+				o = uint64(v.Buf[n])
+			}
+
+			rest = rest<<8 | o
+			index++
 		}
 
-		out = v.Take(cursor, count, order)
-		cursor += count
-		return
+		return func(bits uint) (out byte) {
+			if bits > 8 {
+				log.Panicf("invalid count")
+			}
+			res := byte(rest) & masks[bits]
+			rest = rest >> bits
+			bitIdx = bitIdx + bits
+
+			if bitIdx > (64 - 8) {
+				bitIdx = bitIdx & 7
+				var add uint64
+
+				for n := 6; n >= 0; n-- {
+					var o uint64
+					if len(v.Buf) > index+n {
+						o = uint64(v.Buf[index+n])
+					}
+
+					add = add<<8 | o
+				}
+				index = index + 7
+				rest = rest | add<<(8-bitIdx)
+			}
+
+			return res
+		}
+
+	} else {
+		cursor := uint(0)
+		return func(count uint) (out byte) {
+			if count > 8 {
+				log.Panicf("invalid count")
+			}
+
+			out = v.Take(cursor, count, order)
+			cursor += count
+			return
+		}
 	}
 }

lib/rlepluslazy/internal/bitvector_test.go

@@ -1,6 +1,7 @@
 package bitvector_test
 
 import (
+	"runtime"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -100,7 +101,7 @@ func TestBitVector(t *testing.T) {
 
 		// make a bitvector of 256 sample bits
 		for i := 0; i < 32; i++ {
-			buf = append(buf, 128+32)
+			buf = append(buf, byte(128+i))
 		}
 
 		v := bitvector.NewBitVector(buf, bitvector.LSB0)
@@ -110,7 +111,7 @@ func TestBitVector(t *testing.T) {
 		// compare to Get()
 		for i := uint(0); i < v.Len; i++ {
 			expected, _ := v.Get(i)
-			assert.Equal(t, expected, next(1))
+			assert.Equal(t, expected, next(1), "at index %d", i)
 		}
 
 		// out of range should return zero
@@ -134,3 +135,19 @@ func assertBitVector(t *testing.T, expectedBits []byte, actual bitvector.BitVect
 		assert.Equal(t, bit, actualBit)
 	}
 }
+
+func BenchmarkTake(b *testing.B) {
+	buf := make([]byte, 0, 128)
+	for i := 0; i < 128; i++ {
+		buf = append(buf, byte(128+i))
+	}
+
+	v := bitvector.NewBitVector(buf, bitvector.LSB0)
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	next := v.Iterator(bitvector.LSB0)
+	for i := 0; i < b.N; i++ {
+		runtime.KeepAlive(next(8))
+	}
+}

lib/rlepluslazy/rleplus_test.go

@@ -73,46 +73,36 @@ func TestGolden(t *testing.T) {
 }
 
 func BenchmarkIterator(b *testing.B) {
-	decoded := make([]uint64, 0, 1000)
-
 	for i := 0; i < b.N; i++ {
-		decoded = decoded[:0]
 		rle, _ := FromBuf(goldenRLE)
 		it, _ := rle.Iterator()
 		for it.HasNext() {
 			bit, _ := it.Next()
-			decoded = append(decoded, bit)
+			runtime.KeepAlive(bit)
 		}
-		runtime.KeepAlive(decoded)
 	}
 }
 
 func BenchmarkRunIterator(b *testing.B) {
-	runs := make([]Run, 0, 1000)
 	for i := 0; i < b.N; i++ {
-		runs = runs[:0]
 		rle, _ := FromBuf(goldenRLE)
 		rit, _ := rle.RunIterator()
 		for rit.HasNext() {
 			run, _ := rit.NextRun()
-			runs = append(runs, run)
+			runtime.KeepAlive(run)
 		}
-		runtime.KeepAlive(runs)
 	}
 }
 
 func BenchmarkRunsToBits(b *testing.B) {
-	decoded := make([]uint64, 0, 1000)
 
 	for i := 0; i < b.N; i++ {
-		decoded = decoded[:0]
 		rle, _ := FromBuf(goldenRLE)
 		rit, _ := rle.RunIterator()
 		it, _ := BitsFromRuns(rit)
 		for it.HasNext() {
 			bit, _ := it.Next()
-			decoded = append(decoded, bit)
+			runtime.KeepAlive(bit)
 		}
-		runtime.KeepAlive(decoded)
 	}
 }