<commit_before>// Package roaring is an implementation of Roaring Bitmaps in Go.
// They provide fast compressed bitmap data structures (also called bitset).
// They are ideally suited to represent sets of integers over
// relatively small ranges.
// See http://roaringbitmap.org for details.
package roaring

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
)

// RoaringBitmap represents a compressed bitmap where you can add integers.
type RoaringBitmap struct {
	highlowcontainer roaringArray
}

// Write out a serialized version of this bitmap to stream
func (b *RoaringBitmap) WriteTo(stream io.Writer) (int, error) {
	return b.highlowcontainer.writeTo(stream)
}

// Read a serialized version of this bitmap from stream
func (b *RoaringBitmap) ReadFrom(stream io.Reader) (int, error) {
	return b.highlowcontainer.readFrom(stream)
}

// NewRoaringBitmap creates a new empty RoaringBitmap
func NewRoaringBitmap() *RoaringBitmap {
	return &RoaringBitmap{*newRoaringArray()}
}

// Clear removes all content from the RoaringBitmap and frees the memory
func (rb *RoaringBitmap) Clear() {
	rb.highlowcontainer = *newRoaringArray()
}

// ToArray creates a new slice containing all of the integers stored in the RoaringBitmap in sorted order
func (rb *RoaringBitmap) ToArray() []uint32 {
	array := make([]uint32, rb.GetCardinality())
	pos := 0
	pos2 := 0

	for pos < rb.highlowcontainer.size() {
		hs := toIntUnsigned(rb.highlowcontainer.getKeyAtIndex(pos)) << 16
		c := rb.highlowcontainer.getContainerAtIndex(pos)
		pos++
		c.fillLeastSignificant16bits(array, pos2, hs)
		pos2 += c.getCardinality()
	}
	return array
}

// GetSizeInBytes estimates the memory usage of the RoaringBitmap. Note that this
// might differ slightly from the amount of bytes required for persistent storage
func (rb *RoaringBitmap) GetSizeInBytes() uint64 {
	size := uint64(8)
	for i := 0; i < rb.highlowcontainer.size(); i++ {
		c := rb.highlowcontainer.getContainerAtIndex(i)
		size += uint64(2) + uint64(c.getSizeInBytes())
	}
	return size
}

// GetSerializedSizeInBytes computes the serialized size in bytes  the RoaringBitmap. It should correspond to the
// number of bytes written when invoking WriteTo
func (rb *RoaringBitmap) GetSerializedSizeInBytes() uint64 {
	return rb.highlowcontainer.serializedSizeInBytes()
}

// IntIterable allows you to iterate over the values in a RoaringBitmap
type IntIterable interface {
	HasNext() bool
	Next() uint32
}

type intIterator struct {
	pos              int
	hs               uint32
	iter             shortIterable
	highlowcontainer *roaringArray
}

// HasNext returns true if there are more integers to iterate over
func (ii *intIterator) HasNext() bool {
	return ii.pos < ii.highlowcontainer.size()
}

func (ii *intIterator) init() {
	if ii.highlowcontainer.size() > ii.pos {
		ii.iter = ii.highlowcontainer.getContainerAtIndex(ii.pos).getShortIterator()
		ii.hs = toIntUnsigned(ii.highlowcontainer.getKeyAtIndex(ii.pos)) << 16
	}
}

// Next returns the next integer
func (ii *intIterator) Next() uint32 {
	x := toIntUnsigned(ii.iter.next()) | ii.hs
	if !ii.iter.hasNext() {
		ii.pos = ii.pos + 1
		ii.init()
	}
	return x
}

func newIntIterator(a *RoaringBitmap) *intIterator {
	p := new(intIterator)
	p.pos = 0
	p.highlowcontainer = &a.highlowcontainer
	p.init()
	return p
}

// String creates a string representation of the RoaringBitmap
func (rb *RoaringBitmap) String() string {
	// inspired by https://github.com/fzandona/goroar/blob/master/roaringbitmap.go
	var buffer bytes.Buffer
	start := []byte("{")
	buffer.Write(start)
	i := rb.Iterator()
	for i.HasNext() {
		buffer.WriteString(strconv.Itoa(int(i.Next())))
		if i.HasNext() { // todo: optimize
			buffer.WriteString(",")
		}
	}
	buffer.WriteString("}")
	return buffer.String()
}

// Iterator creates a new IntIterable to iterate over the integers contained in the bitmap, in sorted order
func (rb *RoaringBitmap) Iterator() IntIterable {
	return newIntIterator(rb)
}

// Clone creates a copy of the RoaringBitmap
func (rb *RoaringBitmap) Clone() *RoaringBitmap {
	ptr := new(RoaringBitmap)
	ptr.highlowcontainer = *rb.highlowcontainer.clone()
	return ptr
}

// Contains returns true if the integer is contained in the bitmap
func (rb *RoaringBitmap) Contains(x uint32) bool {
	hb := highbits(x)
	c := rb.highlowcontainer.getContainer(hb)
	return c != nil && c.contains(lowbits(x))
}

// Contains returns true if the integer is contained in the bitmap (this is a convenience method, the parameter is casted to uint32 and Contains is called)
func (rb *RoaringBitmap) ContainsInt(x int) bool {
	return rb.Contains(uint32(x))
}

// Equals returns true if the two bitmaps contain the same integers
func (rb *RoaringBitmap) Equals(o interface{}) bool {
	srb, ok := o.(*RoaringBitmap)
	if ok {
		return srb.highlowcontainer.equals(rb.highlowcontainer)
	}
	return false
}

// Add the integer x to the bitmap
func (rb *RoaringBitmap) Add(x uint32) {
	hb := highbits(x)
	i := rb.highlowcontainer.getIndex(hb)
	if i >= 0 {
		rb.highlowcontainer.setContainerAtIndex(i, rb.highlowcontainer.getContainerAtIndex(i).add(lowbits(x)))
	} else {
		newac := newArrayContainer()
		rb.highlowcontainer.insertNewKeyValueAt(-i-1, hb, newac.add(lowbits(x)))
	}
}

// Add the integer x to the bitmap (convenience method: the parameter is casted to uint32 and we call Add)
func (rb *RoaringBitmap) AddInt(x int) {
	rb.Add(uint32(x))
}

// Remove the integer x from the bitmap
func (rb *RoaringBitmap) Remove(x uint32) {
	hb := highbits(x)
	i := rb.highlowcontainer.getIndex(hb)
	if i >= 0 {
		rb.highlowcontainer.setContainerAtIndex(i, rb.highlowcontainer.getContainerAtIndex(i).remove(lowbits(x)))
		if rb.highlowcontainer.getContainerAtIndex(i).getCardinality() == 0 {
			rb.highlowcontainer.removeAtIndex(i)
		}
	}
}

// IsEmpty returns true if the RoaringBitmap is empty (it is faster than doing (GetCardinality() == 0))
func (rb *RoaringBitmap) IsEmpty() bool {
	return rb.highlowcontainer.size() == 0
}

// GetCardinality returns the number of integers contained in the bitmap
func (rb *RoaringBitmap) GetCardinality() uint64 {
	size := uint64(0)
	for i := 0; i < rb.highlowcontainer.size(); i++ {
		size += uint64(rb.highlowcontainer.getContainerAtIndex(i).getCardinality())
	}
	return size
}

// Rank returns the number of integers that are smaller or equal to x (Rank(infinity) would be GetCardinality())<commit_msg>Added CheckedAdd and CheckedRemove.<commit_after>// Package roaring is an implementation of Roaring Bitmaps in Go.
// They provide fast compressed bitmap data structures (also called bitset).
// They are ideally suited to represent sets of integers over
// relatively small ranges.
// See http://roaringbitmap.org for details.
package roaring

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
)

// RoaringBitmap represents a compressed bitmap where you can add integers.
type RoaringBitmap struct {
	highlowcontainer roaringArray
}

// Write out a serialized version of this bitmap to stream
func (b *RoaringBitmap) WriteTo(stream io.Writer) (int, error) {
	return b.highlowcontainer.writeTo(stream)
}

// Read a serialized version of this bitmap from stream
func (b *RoaringBitmap) ReadFrom(stream io.Reader) (int, error) {
	return b.highlowcontainer.readFrom(stream)
}

// NewRoaringBitmap creates a new empty RoaringBitmap
func NewRoaringBitmap() *RoaringBitmap {
	return &RoaringBitmap{*newRoaringArray()}
}

// Clear removes all content from the RoaringBitmap and frees the memory
func (rb *RoaringBitmap) Clear() {
	rb.highlowcontainer = *newRoaringArray()
}

// ToArray creates a new slice containing all of the integers stored in the RoaringBitmap in sorted order
func (rb *RoaringBitmap) ToArray() []uint32 {
	array := make([]uint32, rb.GetCardinality())
	pos := 0
	pos2 := 0

	for pos < rb.highlowcontainer.size() {
		hs := toIntUnsigned(rb.highlowcontainer.getKeyAtIndex(pos)) << 16
		c := rb.highlowcontainer.getContainerAtIndex(pos)
		pos++
		c.fillLeastSignificant16bits(array, pos2, hs)
		pos2 += c.getCardinality()
	}
	return array
}

// GetSizeInBytes estimates the memory usage of the RoaringBitmap. Note that this
// might differ slightly from the amount of bytes required for persistent storage
func (rb *RoaringBitmap) GetSizeInBytes() uint64 {
	size := uint64(8)
	for i := 0; i < rb.highlowcontainer.size(); i++ {
		c := rb.highlowcontainer.getContainerAtIndex(i)
		size += uint64(2) + uint64(c.getSizeInBytes())
	}
	return size
}

// GetSerializedSizeInBytes computes the serialized size in bytes  the RoaringBitmap. It should correspond to the
// number of bytes written when invoking WriteTo
func (rb *RoaringBitmap) GetSerializedSizeInBytes() uint64 {
	return rb.highlowcontainer.serializedSizeInBytes()
}

// IntIterable allows you to iterate over the values in a RoaringBitmap
type IntIterable interface {
	HasNext() bool
	Next() uint32
}

type intIterator struct {
	pos              int
	hs               uint32
	iter             shortIterable
	highlowcontainer *roaringArray
}

// HasNext returns true if there are more integers to iterate over
func (ii *intIterator) HasNext() bool {
	return ii.pos < ii.highlowcontainer.size()
}

func (ii *intIterator) init() {
	if ii.highlowcontainer.size() > ii.pos {
		ii.iter = ii.highlowcontainer.getContainerAtIndex(ii.pos).getShortIterator()
		ii.hs = toIntUnsigned(ii.highlowcontainer.getKeyAtIndex(ii.pos)) << 16
	}
}

// Next returns the next integer
func (ii *intIterator) Next() uint32 {
	x := toIntUnsigned(ii.iter.next()) | ii.hs
	if !ii.iter.hasNext() {
		ii.pos = ii.pos + 1
		ii.init()
	}
	return x
}

func newIntIterator(a *RoaringBitmap) *intIterator {
	p := new(intIterator)
	p.pos = 0
	p.highlowcontainer = &a.highlowcontainer
	p.init()
	return p
}

// String creates a string representation of the RoaringBitmap
func (rb *RoaringBitmap) String() string {
	// inspired by https://github.com/fzandona/goroar/blob/master/roaringbitmap.go
	var buffer bytes.Buffer
	start := []byte("{")
	buffer.Write(start)
	i := rb.Iterator()
	for i.HasNext() {
		buffer.WriteString(strconv.Itoa(int(i.Next())))
		if i.HasNext() { // todo: optimize
			buffer.WriteString(",")
		}
	}
	buffer.WriteString("}")
	return buffer.String()
}

// Iterator creates a new IntIterable to iterate over the integers contained in the bitmap, in sorted order
func (rb *RoaringBitmap) Iterator() IntIterable {
	return newIntIterator(rb)
}

// Clone creates a copy of the RoaringBitmap
func (rb *RoaringBitmap) Clone() *RoaringBitmap {
	ptr := new(RoaringBitmap)
	ptr.highlowcontainer = *rb.highlowcontainer.clone()
	return ptr
}

// Contains returns true if the integer is contained in the bitmap
func (rb *RoaringBitmap) Contains(x uint32) bool {
	hb := highbits(x)
	c := rb.highlowcontainer.getContainer(hb)
	return c != nil && c.contains(lowbits(x))
}

// Contains returns true if the integer is contained in the bitmap (this is a convenience method, the parameter is casted to uint32 and Contains is called)
func (rb *RoaringBitmap) ContainsInt(x int) bool {
	return rb.Contains(uint32(x))
}

// Equals returns true if the two bitmaps contain the same integers
func (rb *RoaringBitmap) Equals(o interface{}) bool {
	srb, ok := o.(*RoaringBitmap)
	if ok {
		return srb.highlowcontainer.equals(rb.highlowcontainer)
	}
	return false
}

// Add the integer x to the bitmap
func (rb *RoaringBitmap) Add(x uint32) {
	hb := highbits(x)
	i := rb.highlowcontainer.getIndex(hb)
	if i >= 0 {
		rb.highlowcontainer.setContainerAtIndex(i, rb.highlowcontainer.getContainerAtIndex(i).add(lowbits(x)))
	} else {
		newac := newArrayContainer()
		rb.highlowcontainer.insertNewKeyValueAt(-i-1, hb, newac.add(lowbits(x)))
	}
}


// Add the integer x to the bitmap and return true  if it was added (false if the integer was already present)
func (rb *RoaringBitmap) CheckedAdd(x uint32) bool {
	// TODO: add unit tests for this method
	hb := highbits(x)
	i := rb.highlowcontainer.getIndex(hb)
	if i >= 0 {
		C := rb.highlowcontainer.getContainerAtIndex(i)
		oldcard := C.getCardinality()
		C = C.add(lowbits(x))
		rb.highlowcontainer.setContainerAtIndex(i, C)
		return C.getCardinality() > oldcard
	} else {
		newac := newArrayContainer()
		rb.highlowcontainer.insertNewKeyValueAt(-i-1, hb, newac.add(lowbits(x)))
		return true
	}
}

// Add the integer x to the bitmap (convenience method: the parameter is casted to uint32 and we call Add)
func (rb *RoaringBitmap) AddInt(x int) {
	rb.Add(uint32(x))
}

// Remove the integer x from the bitmap
func (rb *RoaringBitmap) Remove(x uint32) {
	hb := highbits(x)
	i := rb.highlowcontainer.getIndex(hb)
	if i >= 0 {
		rb.highlowcontainer.setContainerAtIndex(i, rb.highlowcontainer.getContainerAtIndex(i).remove(lowbits(x)))
		if rb.highlowcontainer.getContainerAtIndex(i).getCardinality() == 0 {
			rb.highlowcontainer.removeAtIndex(i)
		}
	}
}

// Remove the integer x from the bitmap and return true if the integer was effectively remove (and false if the integer was not present)
func (rb *RoaringBitmap) CheckedRemove(x uint32) bool {
	// TODO: add unit tests for this method
	hb := highbits(x)
	i := rb.highlowcontainer.getIndex(hb)
	if i >= 0 {
		C := rb.highlowcontainer.getContainerAtIndex(i)
		oldcard := C.getCardinality()
		C = C.remove(lowbits(x))
		rb.highlowcontainer.setContainerAtIndex(i, C)
		if rb.highlowcontainer.getContainerAtIndex(i).getCardinality() == 0 {
			rb.highlowcontainer.removeAtIndex(i)
			return true
		}
		return C.getCardinality() > oldcard
	} else {
		return false
	}
}


// IsEmpty returns true if the RoaringBitmap is empty (it is faster than doing (GetCardinality() == 0))
func (rb *RoaringBitmap) IsEmpty() bool {
	return rb.highlowcontainer.size() == 0
}

// GetCardinality returns the number of integers contained in the bitmap
func (rb *RoaringBitmap) GetCardinality() uint64 {
	size := uint64(0)
	for i := 0; i < rb.highlowcontainer.size(); i++ {
		size += uint64(rb.highlowcontainer.getContainerAtIndex(i).getCardinality())
	}
	return size
}

// Rank returns the number of integers that are smaller or equal to x (Rank(infinity) would be GetCardinality())
{'fraction_non_alphanumeric': 0.11108738969541702, 'fraction_numerical': 0.009678337603188158, 'mean_word_length': 4.392555640828856, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '282597', 'n_tokens_mistral': 4742, 'n_tokens_neox': 4514, 'n_words': 1408}