roaringarray.go

package roaring

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"

	snappy "github.com/glycerine/go-unsnap-stream"
	"github.com/tinylib/msgp/msgp"
)

//go:generate msgp -unexported

type container interface {
	clone() container
	and(container) container
	iand(container) container // i stands for inplace
	andNot(container) container
	iandNot(container) container // i stands for inplace
	getCardinality() int
	// rank returns the number of integers that are
	// smaller or equal to x. rank(infinity) would be getCardinality().
	rank(uint16) int

	iadd(x uint16) bool                   // inplace, returns true if x was new.
	iaddReturnMinimized(uint16) container // may change return type to minimize storage.

	//addRange(start, final int) container  // range is [firstOfRange,lastOfRange) (unused)
	iaddRange(start, final int) container // i stands for inplace, range is [firstOfRange,lastOfRange)

	iremove(x uint16) bool                   // inplace, returns true if x was present.
	iremoveReturnMinimized(uint16) container // may change return type to minimize storage.

	not(start, final int) container               // range is [firstOfRange,lastOfRange)
	inot(firstOfRange, lastOfRange int) container // i stands for inplace, range is [firstOfRange,lastOfRange)
	xor(r container) container
	getShortIterator() shortIterable
	contains(i uint16) bool

	// equals is now logical equals; it does not require the
	// same underlying container types, but compares across
	// any of the implementations.
	equals(r container) bool

	fillLeastSignificant16bits(array []uint32, i int, mask uint32)
	or(r container) container
	ior(r container) container   // i stands for inplace
	intersects(r container) bool // whether the two containers intersect
	lazyOR(r container) container
	lazyIOR(r container) container
	getSizeInBytes() int
	//removeRange(start, final int) container  // range is [firstOfRange,lastOfRange) (unused)
	iremoveRange(start, final int) container // i stands for inplace, range is [firstOfRange,lastOfRange)
	selectInt(x uint16) int                  // selectInt returns the xth integer in the container
	serializedSizeInBytes() int
	readFrom(io.Reader) (int, error)
	writeTo(io.Writer) (int, error)

	numberOfRuns() int
	toEfficientContainer() container
	String() string
	containerType() contype
}

type contype uint8

const (
	bitmapContype contype = iota
	arrayContype
	run16Contype
	run32Contype
)

// careful: range is [firstOfRange,lastOfRange]
func rangeOfOnes(start, last int) container {
	if start > MaxUint16 {
		panic("rangeOfOnes called with start > MaxUint16")
	}
	if last > MaxUint16 {
		panic("rangeOfOnes called with last > MaxUint16")
	}
	if start < 0 {
		panic("rangeOfOnes called with start < 0")
	}
	if last < 0 {
		panic("rangeOfOnes called with last < 0")
	}
	return newRunContainer16Range(uint16(start), uint16(last))
}

type roaringArray struct {
	keys            []uint16
	containers      []container `msg:"-"` // don't try to serialize directly.
	needCopyOnWrite []bool
	copyOnWrite     bool

	// conserz is used at serialization time
	// to serialize containers. Otherwise empty.
	conserz []containerSerz
}

// containerSerz facilitates serializing container (tricky to
// serialize because it is an interface) by providing a
// light wrapper with a type identifier.
type containerSerz struct {
	t contype  `msg:"t"` // type
	r msgp.Raw `msg:"r"` // Raw msgpack of the actual container type
}

func newRoaringArray() *roaringArray {
	return &roaringArray{}
}

// runOptimize compresses the element containers to minimize space consumed.
// Q: how does this interact with copyOnWrite and needCopyOnWrite?
// A: since we aren't changing the logical content, just the representation,
//    we don't both to check the needCopyOnWrite bits. We replace
//    (possibly all) elements of ra.containers in-place with space
//    optimized versions.
func (ra *roaringArray) runOptimize() {
	for i := range ra.containers {
		ra.containers[i] = ra.containers[i].toEfficientContainer()
	}
}

func (ra *roaringArray) appendContainer(key uint16, value container, mustCopyOnWrite bool) {
	ra.keys = append(ra.keys, key)
	ra.containers = append(ra.containers, value)
	ra.needCopyOnWrite = append(ra.needCopyOnWrite, mustCopyOnWrite)
}

func (ra *roaringArray) appendWithoutCopy(sa roaringArray, startingindex int) {
	ra.appendContainer(sa.keys[startingindex], sa.containers[startingindex], false)
}

func (ra *roaringArray) appendCopy(sa roaringArray, startingindex int) {
	ra.appendContainer(sa.keys[startingindex], sa.containers[startingindex], true)
	sa.setNeedsCopyOnWrite(startingindex)
}

func (ra *roaringArray) appendWithoutCopyMany(sa roaringArray, startingindex, end int) {
	for i := startingindex; i < end; i++ {
		ra.appendWithoutCopy(sa, i)
	}
}

func (ra *roaringArray) appendCopyMany(sa roaringArray, startingindex, end int) {
	for i := startingindex; i < end; i++ {
		ra.appendCopy(sa, i)
	}
}

func (ra *roaringArray) appendCopiesUntil(sa roaringArray, stoppingKey uint16) {
	for i := 0; i < sa.size(); i++ {
		if sa.keys[i] >= stoppingKey {
			break
		}
		ra.appendContainer(sa.keys[i], sa.containers[i], true)
		sa.setNeedsCopyOnWrite(i)
	}
}

func (ra *roaringArray) appendCopiesAfter(sa roaringArray, beforeStart uint16) {
	startLocation := sa.getIndex(beforeStart)
	if startLocation >= 0 {
		startLocation++
	} else {
		startLocation = -startLocation - 1
	}

	for i := startLocation; i < sa.size(); i++ {
		ra.appendContainer(sa.keys[i], sa.containers[i], true)
		sa.setNeedsCopyOnWrite(i)
	}
}

func (ra *roaringArray) removeIndexRange(begin, end int) {
	if end <= begin {
		return
	}

	r := end - begin

	copy(ra.keys[begin:], ra.keys[end:])
	copy(ra.containers[begin:], ra.containers[end:])
	copy(ra.needCopyOnWrite[begin:], ra.needCopyOnWrite[end:])

	ra.resize(len(ra.keys) - r)
}

func (ra *roaringArray) resize(newsize int) {
	for k := newsize; k < len(ra.containers); k++ {
		ra.containers[k] = nil
	}

	ra.keys = ra.keys[:newsize]
	ra.containers = ra.containers[:newsize]
	ra.needCopyOnWrite = ra.needCopyOnWrite[:newsize]
}

func (ra *roaringArray) clear() {
	*ra = roaringArray{}
}

func (ra *roaringArray) clone() *roaringArray {

	// shallow copy, slices will have the same backing arrays.
	sa := *ra

	// this is where copyOnWrite is used.
	if ra.copyOnWrite {
		ra.markAllAsNeedingCopyOnWrite()
		// sa.needCopyOnWrite is shared
	} else {
		// make a full copy

		sa.keys = make([]uint16, len(ra.keys))
		copy(sa.keys, ra.keys)

		sa.containers = make([]container, len(ra.containers))
		for i := range sa.containers {
			sa.containers[i] = ra.containers[i].clone()
		}

		sa.needCopyOnWrite = make([]bool, len(ra.needCopyOnWrite))
	}
	return &sa
}

func (ra *roaringArray) containsKey(x uint16) bool {
	return (ra.binarySearch(0, int64(len(ra.keys)), x) >= 0)
}

func (ra *roaringArray) getContainer(x uint16) container {
	i := ra.binarySearch(0, int64(len(ra.keys)), x)
	if i < 0 {
		return nil
	}
	return ra.containers[i]
}

func (ra *roaringArray) getWritableContainerContainer(x uint16) container {
	i := ra.binarySearch(0, int64(len(ra.keys)), x)
	if i < 0 {
		return nil
	}
	if ra.needCopyOnWrite[i] {
		ra.containers[i] = ra.containers[i].clone()
		ra.needCopyOnWrite[i] = false
	}
	return ra.containers[i]
}

func (ra *roaringArray) getContainerAtIndex(i int) container {
	return ra.containers[i]
}

func (ra *roaringArray) getWritableContainerAtIndex(i int) container {
	if ra.needCopyOnWrite[i] {
		ra.containers[i] = ra.containers[i].clone()
		ra.needCopyOnWrite[i] = false
	}
	return ra.containers[i]
}

func (ra *roaringArray) getIndex(x uint16) int {
	// before the binary search, we optimize for frequent cases
	size := len(ra.keys)
	if (size == 0) || (ra.keys[size-1] == x) {
		return size - 1
	}
	return int(ra.binarySearch(0, int64(size), x))
}

func (ra *roaringArray) getKeyAtIndex(i int) uint16 {
	return ra.keys[i]
}

func (ra *roaringArray) insertNewKeyValueAt(i int, key uint16, value container) {
	ra.keys = append(ra.keys, 0)
	ra.containers = append(ra.containers, nil)

	copy(ra.keys[i+1:], ra.keys[i:])
	copy(ra.containers[i+1:], ra.containers[i:])

	ra.keys[i] = key
	ra.containers[i] = value

	ra.needCopyOnWrite = append(ra.needCopyOnWrite, false)
	copy(ra.needCopyOnWrite[i+1:], ra.needCopyOnWrite[i:])
	ra.needCopyOnWrite[i] = false
}

func (ra *roaringArray) remove(key uint16) bool {
	i := ra.binarySearch(0, int64(len(ra.keys)), key)
	if i >= 0 { // if a new key
		ra.removeAtIndex(i)
		return true
	}
	return false
}

func (ra *roaringArray) removeAtIndex(i int) {
	copy(ra.keys[i:], ra.keys[i+1:])
	copy(ra.containers[i:], ra.containers[i+1:])

	copy(ra.needCopyOnWrite[i:], ra.needCopyOnWrite[i+1:])

	ra.resize(len(ra.keys) - 1)
}

func (ra *roaringArray) setContainerAtIndex(i int, c container) {
	ra.containers[i] = c
}

func (ra *roaringArray) replaceKeyAndContainerAtIndex(i int, key uint16, c container, mustCopyOnWrite bool) {
	ra.keys[i] = key
	ra.containers[i] = c
	ra.needCopyOnWrite[i] = mustCopyOnWrite
}

func (ra *roaringArray) size() int {
	return len(ra.keys)
}

func (ra *roaringArray) binarySearch(begin, end int64, ikey uint16) int {
	low := begin
	high := end - 1
	for low+16 <= high {
		middleIndex := low + (high-low)/2 // avoid overflow
		middleValue := ra.keys[middleIndex]

		if middleValue < ikey {
			low = middleIndex + 1
		} else if middleValue > ikey {
			high = middleIndex - 1
		} else {
			return int(middleIndex)
		}
	}
	for ; low <= high; low++ {
		val := ra.keys[low]
		if val >= ikey {
			if val == ikey {
				return int(low)
			}
			break
		}
	}
	return -int(low + 1)
}

func (ra *roaringArray) equals(o interface{}) bool {
	srb, ok := o.(roaringArray)
	if ok {

		if srb.size() != ra.size() {
			return false
		}
		for i, k := range ra.keys {
			if k != srb.keys[i] {
				return false
			}
		}

		for i, c := range ra.containers {
			if !c.equals(srb.containers[i]) {
				return false
			}
		}
		return true
	}
	return false
}

// this can be expensive; don't call it unnecessarily in production
func (ra *roaringArray) serializedSizeInBytes() uint64 {
	n, _ := ra.writeToHelper(nil, true)
	return uint64(n)
}

//
// spec: https://github.com/RoaringBitmap/RoaringFormatSpec
//
func (ra *roaringArray) writeTo(out io.Writer) (int64, error) {
	return ra.writeToHelper(out, false)
}

// if fake then `out io.Writer` won't be touched, and the number of bytes
// that would have been written and nil will be returned
func (ra *roaringArray) writeToHelper(out io.Writer, fake bool) (int64, error) {
	stream := &bytes.Buffer{}

	//p("roaringArray.writeTo starting")
	numKeys := len(ra.keys)
	isRunSizeInBytes := (numKeys + 7) / 8

	// compute isRun bitmap
	var ir []byte
	isRun := newBitmapContainer()
	for i, c := range ra.containers {
		switch c.(type) {
		case *runContainer16:
			isRun.iadd(uint16(i))
		}
	}
	ir = uint64SliceAsByteSlice(isRun.bitmap)[:isRunSizeInBytes]

	const cookieSize = 4
	descriptiveHeaderSize := 4 * len(ra.keys)
	preambleSize := cookieSize + isRunSizeInBytes + descriptiveHeaderSize

	buf := make([]byte, preambleSize+4*len(ra.keys))
	binary.LittleEndian.PutUint16(buf[0:], uint16(serialCookie))
	if len(ra.keys) > MaxUint16 {
		panic("should be impossible to have this many keys")
	}
	binary.LittleEndian.PutUint16(buf[2:], uint16(len(ra.keys)-1))
	nw := 4
	//p("after writing serialCookie and number of keys -1, now nw: %v", nw)

	// isRun bitset
	nw += copy(buf[nw:], ir)
	//p("after writing isRun array of len(%v), now nw: %v, this will be start of descriptive header", len(ir), nw)

	// descriptive header
	for i, key := range ra.keys {
		binary.LittleEndian.PutUint16(buf[nw:], uint16(key))
		nw += 2
		c := ra.containers[i]
		binary.LittleEndian.PutUint16(buf[nw:], uint16(c.getCardinality()-1))
		nw += 2
	}
	//p("after writing descriptiveHeader, now nw: %v", nw)

	startOffset := int64(preambleSize + 4*len(ra.keys))
	if len(ra.keys) >= noOffsetThreshold {
		//p("writing offset header at nw: %v", nw)
		// offset header
		for _, c := range ra.containers {
			binary.LittleEndian.PutUint32(buf[nw:], uint32(startOffset))
			nw += 4
			switch rc := c.(type) {
			case *runContainer16:
				startOffset += 2 + int64(len(rc.iv))*4
			default:
				startOffset += int64(getSizeInBytesFromCardinality(c.getCardinality()))
			}
		}
	} else {
		//p("skiping offset header at nw %v", nw)
	}

	if !fake {
		nsw, err := stream.Write(buf[:nw])
		if err != nil {
			return 0, err
		}
		if nsw != nw {
			panic("short write!")
		}
	}
	for i, c := range ra.containers {
		_ = i
		//p("writeTo writing %v-th container with key %v (card: %v) at fileloc: %v / nw: %v", i, ra.keys[i], c.getCardinality(), len(stream.Bytes()), nw)
		writ, err := c.writeTo(stream)
		if err != nil {
			return 0, err
		}
		nw += writ
	}

	by := stream.Bytes()

	if fake {
		return int64(nw), nil
	}

	n, err := out.Write(by)
	if err == nil && n < len(by) {
		err = io.ErrShortWrite
	}
	return int64(n), err
}

func (ra *roaringArray) readFrom(stream io.Reader) (int64, error) {

	pos := 0

	var cookie uint32
	err := binary.Read(stream, binary.LittleEndian, &cookie)

	if err != nil {
		return 0, fmt.Errorf("error in roaringArray.readFrom: could not read initial cookie: %s", err)
	}
	pos += 4

	if cookie != serialCookieNoRunContainer && cookie&0x0000FFFF != serialCookie {
		return 0, fmt.Errorf("error in roaringArray.readFrom: did not find expected serialCookie in header")
	}

	var size uint32
	haveRunContainers := false
	var isRun container = newRunContainer16()
	if cookie&0x0000FFFF == serialCookie {
		haveRunContainers = true
		size = (cookie >> 16) + 1
		//p("stream contains run containers. total number of containers == size = %v", size)
		bytesToRead := (int(size) + 7) / 8
		//p("bytesToRead = %v", bytesToRead)
		by := make([]byte, bytesToRead)
		nr, err := io.ReadAtLeast(stream, by, bytesToRead)
		if err != nil {
			return 8 + int64(nr), fmt.Errorf("error in readFrom: could not read the "+
				"runContainer bit flags of length %v bytes: %v", bytesToRead, err)
		}
		pos += nr
		//p("by = %#v", by)

		bc := newBitmapContainer()
		asBy := uint64SliceAsByteSlice(bc.bitmap)
		for i := range by {
			asBy[i] = by[i]
		}
		bc.computeCardinality()
		//p("bc = %s", bc)
		isRun = bc
		//p("read in isRun bitmap container ok: '%s'", bc.String())
		//p("after read of isRun bitmap, now at pos = %v", pos)
	} else {
		err = binary.Read(stream, binary.LittleEndian, &size)
		if err != nil {
			return 0, fmt.Errorf("error in roaringArray.readFrom: when reading size, got: %s", err)
		}
		pos += 4
		//p("!serialCookie, so after skip of isRun bitmap, and read of size, now at pos = %v", pos)
	}

	// descriptive header
	keycard := make([]uint16, 2*size, 2*size)
	err = binary.Read(stream, binary.LittleEndian, keycard)
	if err != nil {
		return 0, err
	}
	//p("got descriptive header, starting at pos %v:", pos)
	pos += 2 * 2 * int(size)
	//p("descriptive header ends at pos %v:", pos)

	k := 0
	for i := 0; i < len(keycard); i += 2 {
		if rleVerbose {
			// debug stuff here:
			yesRun := false
			if isRun.contains(uint16(k)) {
				yesRun = true
			}
			fmt.Printf("\n    descriptive header: key %x:  card %v   isRun: %v\n", keycard[i], int(keycard[i+1])+1, yesRun)
		}
		k++
	}

	// offset header
	if !haveRunContainers || size >= noOffsetThreshold {
		//p("reading offset header of size %v at pos %v", size, pos)
		offsets := make([]uint32, size, size)
		err = binary.Read(stream, binary.LittleEndian, offsets)
		if err != nil {
			return 0, err
		}
		pos += 4 * int(size)
		//p("read offset header of count %v. offsets: %v. pos after is %v", size, offsets, pos)
	} else {
		//p("skipping/should be no offset header, because size %v < noOffsetThreshold.  pos is now %v", size, pos)
	}
	//offset := int64(4 + 4 + 8*size) // probably wrong now
	for i := uint32(0); i < size; i++ {
		key := int(keycard[2*i])
		card := int(keycard[2*i+1]) + 1
		if haveRunContainers && isRun.contains(uint16(i)) {
			//p("at key %v, filepos %v, reading a run container", key, pos)
			nb := newRunContainer16()
			nr, err := nb.readFrom(stream)
			if err != nil {
				return 0, err
			}
			//p("at key %v, filepos %v, read a run container: %s", key, pos, nb)
			pos += nr
			//offset += 2 + int64(len(nb.iv))*4
			ra.appendContainer(uint16(key), nb, false)
		} else {
			//offset += int64(getSizeInBytesFromCardinality(card))
			if card > arrayDefaultMaxSize {
				//p("at key %v, filepos %v, reading a bitmap container", key, pos)

				nb := newBitmapContainer()
				nr, err := nb.readFrom(stream)
				if err != nil {
					return 0, err
				}
				nb.cardinality = int(card)
				//p("at key %v, filepos %v, read a bitmap container = %s", key, pos, nb)
				pos += nr
				ra.appendContainer(keycard[2*i], nb, false)
			} else {
				//p("at key %v, filepos %v, reading an array container with card %v", key, pos, card)

				nb := newArrayContainerSize(int(card))
				nr, err := nb.readFrom(stream)
				if err != nil {
					return 0, err
				}
				//p("at key %v, filepos %v, read an array container (card %v) = %s",key, pos, card, nb)
				pos += nr
				ra.appendContainer(keycard[2*i], nb, false)
			}
		}
	}
	return int64(pos), nil
}

func (ra *roaringArray) writeToMsgpack(stream io.Writer) error {

	ra.conserz = make([]containerSerz, len(ra.containers))
	for i, v := range ra.containers {
		switch cn := v.(type) {
		case *bitmapContainer:
			bts, err := cn.MarshalMsg(nil)
			if err != nil {
				return err
			}
			ra.conserz[i].t = bitmapContype
			ra.conserz[i].r = bts
		case *arrayContainer:
			bts, err := cn.MarshalMsg(nil)
			if err != nil {
				return err
			}
			ra.conserz[i].t = arrayContype
			ra.conserz[i].r = bts
		case *runContainer16:
			bts, err := cn.MarshalMsg(nil)
			if err != nil {
				return err
			}
			ra.conserz[i].t = run16Contype
			ra.conserz[i].r = bts
		default:
			fmt.Errorf("Unrecognized container implementation: %T", cn)
		}
	}
	w := snappy.NewWriter(stream)
	err := msgp.Encode(w, ra)
	ra.conserz = nil
	return err
}

func (ra *roaringArray) readFromMsgpack(stream io.Reader) error {
	r := snappy.NewReader(stream)
	err := msgp.Decode(r, ra)
	if err != nil {
		return err
	}

	if len(ra.containers) != len(ra.keys) {
		ra.containers = make([]container, len(ra.keys))
	}

	for i, v := range ra.conserz {
		switch v.t {
		case bitmapContype:
			c := &bitmapContainer{}
			_, err = c.UnmarshalMsg(v.r)
			if err != nil {
				return err
			}
			ra.containers[i] = c
		case arrayContype:
			c := &arrayContainer{}
			_, err = c.UnmarshalMsg(v.r)
			if err != nil {
				return err
			}
			ra.containers[i] = c
		case run16Contype:
			c := &runContainer16{}
			_, err = c.UnmarshalMsg(v.r)
			if err != nil {
				return err
			}
			ra.containers[i] = c
		default:
			return fmt.Errorf("unrecognized contype serialization code: '%v'", v.t)
		}
	}
	ra.conserz = nil
	return nil
}

func (ra *roaringArray) advanceUntil(min uint16, pos int) int {
	lower := pos + 1

	if lower >= len(ra.keys) || ra.keys[lower] >= min {
		return lower
	}

	spansize := 1

	for lower+spansize < len(ra.keys) && ra.keys[lower+spansize] < min {
		spansize *= 2
	}
	var upper int
	if lower+spansize < len(ra.keys) {
		upper = lower + spansize
	} else {
		upper = len(ra.keys) - 1
	}

	if ra.keys[upper] == min {
		return upper
	}

	if ra.keys[upper] < min {
		// means
		// array
		// has no
		// item
		// >= min
		// pos = array.length;
		return len(ra.keys)
	}

	// we know that the next-smallest span was too small
	lower += (spansize / 2)

	mid := 0
	for lower+1 != upper {
		mid = (lower + upper) / 2
		if ra.keys[mid] == min {
			return mid
		} else if ra.keys[mid] < min {
			lower = mid
		} else {
			upper = mid
		}
	}
	return upper
}

func (ra *roaringArray) markAllAsNeedingCopyOnWrite() {
	for i := range ra.needCopyOnWrite {
		ra.needCopyOnWrite[i] = true
	}
}

func (ra *roaringArray) needsCopyOnWrite(i int) bool {
	return ra.needCopyOnWrite[i]
}

func (ra *roaringArray) setNeedsCopyOnWrite(i int) {
	ra.needCopyOnWrite[i] = true
}