featurebase/index/bitmap.go

package index

// #cgo  CFLAGS:-mpopcnt

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/binary"
	"encoding/gob"

	log "github.com/cihub/seelog"
	"github.com/yasushi-saito/rbtree"
)

const CounterMask = uint64(0xffffffffffffffff)

var CounterKey = int64(-1)

// Bitmap represents a bitmap broken up into Chunks.
// Internally it is represented as a red-black tree of chunks.
type Bitmap struct {
	tree   *rbtree.Tree
	bcount uint64
}

// NewBitmap returns a new instance of Bitmap.
func NewBitmap() *Bitmap {
	return &Bitmap{
		tree: rbtree.NewTree(rbtreeItemCompare),
	}
}

// Chunk returns the chunk within the bitmap.
// Returns nil if the chunk key does not exist.
func (b *Bitmap) Chunk(c *Chunk) *Chunk {
	if n := b.tree.Get(c); n != nil {
		return n.(*Chunk)
	}
	return nil
}

// AddChunk adds c to the bitmap.
func (b *Bitmap) AddChunk(c *Chunk) { b.tree.Insert(c) }

// ChunkIterator returns an iterator for looping over the bitmap's chunks.
func (b *Bitmap) ChunkIterator() *ChunkIterator {
	return &ChunkIterator{b.tree.Min()}
}

// Clone returns a copy of b.
func (b *Bitmap) Clone() *Bitmap {
	itr := b.ChunkIterator()
	other := NewBitmap()

	for {
		if itr.Limit() {
			break
		}

		node := itr.Item()
		other.AddChunk(&Chunk{
			Key:   node.Key,
			Value: node.Value.copy(),
		})

		itr = itr.Next()
	}
	return other
}

// IntersectionCount returns the number of itersections between b and other.
func (b *Bitmap) IntersectionCount(other *Bitmap) uint64 {
	itr0 := b.ChunkIterator()
	itr1 := other.ChunkIterator()

	results := uint64(0)
	for {
		if itr1.Limit() || itr0.Limit() {
			break
		} else if itr0.Item().Key < itr1.Item().Key {
			itr0 = itr0.Next()
		} else if itr0.Item().Key > itr1.Item().Key {
			itr1 = itr1.Next()
		} else if itr0.Item().Key == itr1.Item().Key {
			results += itr0.Item().Value.andcount(itr1.Item().Value)
			itr0 = itr0.Next()
			itr1 = itr1.Next()
		}
	}
	return results
}

// Intersection returns the itersection of b and other.
func (b *Bitmap) Intersection(other *Bitmap) *Bitmap {
	itr0 := b.ChunkIterator()
	itr1 := other.ChunkIterator()

	output := NewBitmap()
	for {
		if itr1.Limit() || itr0.Limit() {
			break
		} else if itr0.Item().Key < itr1.Item().Key {
			itr0 = itr0.Next()
		} else if itr0.Item().Key > itr1.Item().Key {
			itr1 = itr1.Next()
		} else if itr0.Item().Key == itr1.Item().Key {
			output.AddChunk(&Chunk{
				Key:   itr0.Item().Key,
				Value: itr0.Item().Value.intersection(itr1.Item().Value),
			})
			itr0 = itr0.Next()
			itr1 = itr1.Next()
		}
	}
	return output
}

// Invert returns a bitwise inversion of b.
func (b *Bitmap) Invert() *Bitmap {
	other := NewBitmap()
	for i := b.ChunkIterator(); !i.Limit(); i = i.Next() {
		other.AddChunk(&Chunk{
			Key:   i.Item().Key,
			Value: i.Item().Value.invert(),
		})
	}
	return other

}

// Union returns the bitwise union of b and other.
func (b *Bitmap) Union(other *Bitmap) *Bitmap {
	itr0 := b.ChunkIterator()
	itr1 := other.ChunkIterator()

	output := NewBitmap()
	eof := uint64(0xdeadbeef)

	for {
		if itr0.Limit() && itr1.Limit() {
			break
		} else if itr0.Limit() {
			if eof == itr1.Item().Key {
				break
			}
			output.AddChunk(&Chunk{itr1.Item().Key, itr1.Item().Value})
			eof = itr1.Item().Key
			itr1 = itr1.Next()
		} else if itr1.Limit() {
			if eof == itr0.Item().Key {
				break
			}
			output.AddChunk(&Chunk{itr0.Item().Key, itr0.Item().Value})
			eof = itr0.Item().Key
			itr0 = itr0.Next()
		} else if itr0.Item().Key < itr1.Item().Key {
			output.AddChunk(&Chunk{itr0.Item().Key, itr0.Item().Value})
			eof = itr0.Item().Key
			itr0 = itr0.Next()
		} else if itr0.Item().Key > itr1.Item().Key {
			output.AddChunk(&Chunk{itr1.Item().Key, itr1.Item().Value})
			eof = itr1.Item().Key
			itr1 = itr1.Next()
		} else if itr0.Item().Key == itr1.Item().Key {
			output.AddChunk(&Chunk{
				Key:   itr0.Item().Key,
				Value: itr0.Item().Value.union(itr1.Item().Value),
			})
			eof = itr0.Item().Key
			itr0 = itr0.Next()
			itr1 = itr1.Next()
		} else {
			panic("unreachable")
		}
	}
	return output
}

// Difference returns the diff of b and other.
func (b *Bitmap) Difference(other *Bitmap) *Bitmap {
	itr0 := b.ChunkIterator()
	itr1 := other.ChunkIterator()

	output := NewBitmap()
	for {
		if itr0.Limit() && itr1.Limit() {
			break
		} else if itr0.Limit() {
			break
		} else if itr1.Limit() {
			output.AddChunk(&Chunk{itr0.Item().Key, itr0.Item().Value})
			itr0 = itr0.Next()
		} else if itr0.Item().Key < itr1.Item().Key {
			output.AddChunk(&Chunk{itr0.Item().Key, itr0.Item().Value})
			itr0 = itr0.Next()
		} else if itr0.Item().Key > itr1.Item().Key {
			itr1 = itr1.Next()
		} else if itr0.Item().Key == itr1.Item().Key {
			chunk := &Chunk{
				Key:   itr0.Item().Key,
				Value: itr0.Item().Value.difference(itr1.Item().Value),
			}

			// Could not add if all zero
			if chunk.Value.bitcount() > 0 {
				output.AddChunk(chunk)
			}

			itr0 = itr0.Next()
			itr1 = itr1.Next()
		} else {
			panic("unreachable")
		}
	}
	return output
}

// ToRawCompressString returns a compressed, hex-encoded string of b.
func (b *Bitmap) ToRawCompressString() (string, int) {
	var bt bytes.Buffer
	buf := gzip.NewWriter(&bt)
	binary.Write(buf, binary.LittleEndian, uint64(b.tree.Len()))
	max_slice := 0
	for i := b.tree.Min(); !i.Limit(); i = i.Next() {
		obj := i.Item().(*Chunk)
		max_slice = int(obj.Key)
		binary.Write(buf, binary.LittleEndian, obj.Key)
		for _, v := range obj.Value {
			binary.Write(buf, binary.LittleEndian, v)
		}
	}
	buf.Flush()
	//buf.Close()
	max_slice = max_slice / 32
	return base64.StdEncoding.EncodeToString(bt.Bytes()), max_slice
}

// ToBytes returns a gob-encoded byte slice of b.
func (b *Bitmap) ToBytes() []byte {
	var buf bytes.Buffer
	enc := gob.NewEncoder(&buf)
	enc.Encode(b.tree.Len())
	for i := b.tree.Min(); !i.Limit(); i = i.Next() {
		obj := i.Item().(*Chunk)
		err := enc.Encode(obj)
		if err != nil {
			log.Warn(err)
		}
	}
	return buf.Bytes()
}

// FromBytes decodes a gob-encoded byte slice into b.
func (b *Bitmap) FromBytes(raw []byte) {
	buf := bytes.NewBuffer(raw)
	dec := gob.NewDecoder(buf)

	var size int
	dec.Decode(&size)
	b.tree = rbtree.NewTree(rbtreeItemCompare)
	for i := 0; i < size; i++ {
		var chunk Chunk
		dec.Decode(&chunk)
		b.AddChunk(&chunk)
	}
	b.SetCount(b.BitCount())
}

// Bits returns the bits in b as a slice of ints.
func (b *Bitmap) Bits() []uint64 {
	result := make([]uint64, b.Count())

	x := 0
	for i := b.ChunkIterator(); !i.Limit(); i = i.Next() {
		item := i.Item()
		chunk := item.Key
		for bi, block := range item.Value {
			for bit := uint(0); bit < 64; bit++ {
				if (block & (1 << bit)) != 0 {
					idx := chunk << 11
					idx = idx | uint64((uint(bi)<<6)|bit)
					result[x] = idx
					x++
				}
			}
		}

	}
	return result
}

// SetBit sets the i-th bit of the bitmap.
func (b *Bitmap) SetBit(i uint64) (bool, *Chunk, Address) {
	address := deref(i)

	chunk := b.Chunk(&Chunk{address.ChunkKey, make(Blocks, 32)})
	if chunk == nil {
		chunk = &Chunk{address.ChunkKey, make(Blocks, 32)}
		b.AddChunk(chunk)
	}

	changed := chunk.Value.setBit(address.BlockIndex, address.Bit)
	if changed {
		b.bcount++
	}

	return changed, chunk, address
}

// ClearBit clears the i-th bit of the bitmap.
func (b *Bitmap) ClearBit(i uint64) (bool, *Chunk, Address) {
	address := deref(i)

	chunk := b.Chunk(&Chunk{address.ChunkKey, make(Blocks, 32)})
	if chunk == nil {
		return false, nil, address
	}

	changed := chunk.Value.clearBit(address.BlockIndex, address.Bit)
	if changed && b.bcount > 0 {
		b.bcount--
	}

	return changed, chunk, address
}

// Len returns the number of chunks in b.
func (b *Bitmap) Len() int { return b.tree.Len() }

// SetCount sets the number of set bits in the bitmap.
func (b *Bitmap) SetCount(c uint64) { b.bcount = c }

// Count returns the number of set bits in the bitmap.
func (b *Bitmap) Count() uint64 { return b.bcount }

// BitCount calculates the number of set bits in the bitmap from raw chunk data.
func (b *Bitmap) BitCount() uint64 {
	var n uint64
	for i := b.ChunkIterator(); !i.Limit(); i = i.Next() {
		n += i.Item().Value.bitcount()
	}
	return n
}

// Chunk represents a set of blocks in a Bitmap.
type Chunk struct {
	Key   uint64
	Value Blocks
}

// ChunkIterator represents an object for iterating over chunks in a bitmap.
type ChunkIterator struct {
	itr rbtree.Iterator
}

// Limit return true when the iterator is at the end of iteration.
func (r *ChunkIterator) Limit() bool {
	return r.itr.Limit()
}

// Next moves the iterator to the next chunk.
func (r *ChunkIterator) Next() *ChunkIterator {
	r.itr = r.itr.Next()
	return r
}

// Item returns the current item that the iterator is pointing at.
func (r *ChunkIterator) Item() *Chunk {
	if r.itr.Item() != nil {
		return r.itr.Item().(*Chunk)
	}
	return nil
}

func rbtreeItemCompare(a, b rbtree.Item) int {
	aKey, bKey := a.(*Chunk).Key, b.(*Chunk).Key
	if aKey < bKey {
		return -1
	} else if aKey > bKey {
		return 1
	}
	return 0
}

// Address represents a location for a given chunk/block/bit.
type Address struct {
	ChunkKey   uint64
	BlockIndex uint8
	Bit        uint8
}

func deref(pos uint64) Address {
	chunkKey := pos >> 11              // div by 2048
	offset := pos & 0x7FF              // mod by 2048
	blockIndex := uint8(offset >> 6)   // div by 64
	bit_offset := uint8(offset & 0x3F) // mod by 64
	return Address{chunkKey, blockIndex, bit_offset}
}