1 files changed, 151 insertions, 33 deletions
diff --git a/vendor/github.com/willf/bitset/bitset.go b/vendor/github.com/willf/bitset/bitset.go
index 65ef6851d1..32044f5c83 100644
--- a/vendor/github.com/willf/bitset/bitset.go
+++ b/vendor/github.com/willf/bitset/bitset.go
@@ -58,6 +58,18 @@ const log2WordSize = uint(6)
 // allBits has every bit set
 const allBits uint64 = 0xffffffffffffffff
 
+// default binary BigEndian
+var binaryOrder binary.ByteOrder = binary.BigEndian
+
+// default json encoding base64.URLEncoding
+var base64Encoding = base64.URLEncoding
+
+// Base64StdEncoding Marshal/Unmarshal BitSet with base64.StdEncoding(Default: base64.URLEncoding)
+func Base64StdEncoding() { base64Encoding = base64.StdEncoding }
+
+// LittleEndian Marshal/Unmarshal Binary as Little Endian(Default: binary.BigEndian)
+func LittleEndian() { binaryOrder = binary.LittleEndian }
+
 // A BitSet is a set of bits. The zero value of a BitSet is an empty set of length 0.
 type BitSet struct {
 	length uint
@@ -180,6 +192,70 @@ func (b *BitSet) Flip(i uint) *BitSet {
 	return b
 }
 
+// Shrink shrinks BitSet to desired length in bits. It clears all bits > length
+// and reduces the size and length of the set.
+//
+// A new slice is allocated to store the new bits, so you may see an increase in
+// memory usage until the GC runs. Normally this should not be a problem, but if you
+// have an extremely large BitSet its important to understand that the old BitSet will
+// remain in memory until the GC frees it.
+func (b *BitSet) Shrink(length uint) *BitSet {
+	idx := wordsNeeded(length + 1)
+	if idx > len(b.set) {
+		return b
+	}
+	shrunk := make([]uint64, idx)
+	copy(shrunk, b.set[:idx])
+	b.set = shrunk
+	b.length = length + 1
+	b.set[idx-1] &= (allBits >> (uint64(64) - uint64(length&(wordSize-1)) - 1))
+	return b
+}
+
+// InsertAt takes an index which indicates where a bit should be
+// inserted. Then it shifts all the bits in the set to the left by 1, starting
+// from the given index position, and sets the index position to 0.
+//
+// Depending on the size of your BitSet, and where you are inserting the new entry,
+// this method could be extremely slow and in some cases might cause the entire BitSet
+// to be recopied.
+func (b *BitSet) InsertAt(idx uint) *BitSet {
+	insertAtElement := (idx >> log2WordSize)
+
+	// if length of set is a multiple of wordSize we need to allocate more space first
+	if b.isLenExactMultiple() {
+		b.set = append(b.set, uint64(0))
+	}
+
+	var i uint
+	for i = uint(len(b.set) - 1); i > insertAtElement; i-- {
+		// all elements above the position where we want to insert can simply by shifted
+		b.set[i] <<= 1
+
+		// we take the most significant bit of the previous element and set it as
+		// the least significant bit of the current element
+		b.set[i] |= (b.set[i-1] & 0x8000000000000000) >> 63
+	}
+
+	// generate a mask to extract the data that we need to shift left
+	// within the element where we insert a bit
+	dataMask := ^(uint64(1)<<uint64(idx&(wordSize-1)) - 1)
+
+	// extract that data that we'll shift
+	data := b.set[i] & dataMask
+
+	// set the positions of the data mask to 0 in the element where we insert
+	b.set[i] &= ^dataMask
+
+	// shift data mask to the left and insert its data to the slice element
+	b.set[i] |= data << 1
+
+	// add 1 to length of BitSet
+	b.length++
+
+	return b
+}
+
 // String creates a string representation of the Bitmap
 func (b *BitSet) String() string {
 	// follows code from https://github.com/RoaringBitmap/roaring
@@ -205,6 +281,43 @@ func (b *BitSet) String() string {
 	return buffer.String()
 }
 
+// DeleteAt deletes the bit at the given index position from
+// within the bitset
+// All the bits residing on the left of the deleted bit get
+// shifted right by 1
+// The running time of this operation may potentially be
+// relatively slow, O(length)
+func (b *BitSet) DeleteAt(i uint) *BitSet {
+	// the index of the slice element where we'll delete a bit
+	deleteAtElement := i >> log2WordSize
+
+	// generate a mask for the data that needs to be shifted right
+	// within that slice element that gets modified
+	dataMask := ^((uint64(1) << (i & (wordSize - 1))) - 1)
+
+	// extract the data that we'll shift right from the slice element
+	data := b.set[deleteAtElement] & dataMask
+
+	// set the masked area to 0 while leaving the rest as it is
+	b.set[deleteAtElement] &= ^dataMask
+
+	// shift the previously extracted data to the right and then
+	// set it in the previously masked area
+	b.set[deleteAtElement] |= (data >> 1) & dataMask
+
+	// loop over all the consecutive slice elements to copy each
+	// lowest bit into the highest position of the previous element,
+	// then shift the entire content to the right by 1
+	for i := int(deleteAtElement) + 1; i < len(b.set); i++ {
+		b.set[i-1] |= (b.set[i] & 1) << 63
+		b.set[i] >>= 1
+	}
+
+	b.length = b.length - 1
+
+	return b
+}
+
 // NextSet returns the next bit set from the specified index,
 // including possibly the current index
 // along with an error code (true = valid, false = no set bit found)
@@ -234,7 +347,7 @@ func (b *BitSet) NextSet(i uint) (uint, bool) {
 // including possibly the current index and up to cap(buffer).
 // If the returned slice has len zero, then no more set bits were found
 //
-//    buffer := make([]uint, 256)
+//    buffer := make([]uint, 256) // this should be reused
 //    j := uint(0)
 //    j, buffer = bitmap.NextSetMany(j, buffer)
 //    for ; len(buffer) > 0; j, buffer = bitmap.NextSetMany(j,buffer) {
@@ -245,39 +358,44 @@ func (b *BitSet) NextSet(i uint) (uint, bool) {
 //    }
 //
 func (b *BitSet) NextSetMany(i uint, buffer []uint) (uint, []uint) {
-	myanswer := buffer[:0]
-
+	myanswer := buffer
+	capacity := cap(buffer)
 	x := int(i >> log2WordSize)
-	if x >= len(b.set) {
-		return 0, myanswer
+	if x >= len(b.set) || capacity == 0 {
+		return 0, myanswer[:0]
+	}
+	skip := i & (wordSize - 1)
+	word := b.set[x] >> skip
+	myanswer = myanswer[:capacity]
+	size := int(0)
+	for word != 0 {
+		r := trailingZeroes64(word)
+		t := word & ((^word) + 1)
+		myanswer[size] = r + i
+		size++
+		if size == capacity {
+			goto End
+		}
+		word = word ^ t
 	}
-	w := b.set[x]
-	w = w >> (i & (wordSize - 1))
-	base := uint(x << 6)
-	capacity := cap(buffer)
-	for len(myanswer) < capacity {
-		for w != 0 {
-			t := w & ((^w) + 1)
-			r := trailingZeroes64(w)
-			myanswer = append(myanswer, r+base)
-			if len(myanswer) == capacity {
+	x++
+	for idx, word := range b.set[x:] {
+		for word != 0 {
+			r := trailingZeroes64(word)
+			t := word & ((^word) + 1)
+			myanswer[size] = r + (uint(x+idx) << 6)
+			size++
+			if size == capacity {
 				goto End
 			}
-			w = w ^ t
-		}
-		x += 1
-		if x == len(b.set) {
-			break
+			word = word ^ t
 		}
-		base += 64
-		w = b.set[x]
 	}
 End:
-	if len(myanswer) > 0 {
-		return myanswer[len(myanswer)-1], myanswer
-	} else {
-		return 0, myanswer
+	if size > 0 {
+		return myanswer[size-1], myanswer[:size]
 	}
+	return 0, myanswer[:0]
 }
 
 // NextClear returns the next clear bit from the specified index,
@@ -654,7 +772,7 @@ func (b *BitSet) DumpAsBits() string {
 	for ; i >= 0; i-- {
 		fmt.Fprintf(buffer, "%064b.", b.set[i])
 	}
-	return string(buffer.Bytes())
+	return buffer.String()
 }
 
 // BinaryStorageSize returns the binary storage requirements
@@ -667,13 +785,13 @@ func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
 	length := uint64(b.length)
 
 	// Write length
-	err := binary.Write(stream, binary.BigEndian, length)
+	err := binary.Write(stream, binaryOrder, length)
 	if err != nil {
 		return 0, err
 	}
 
 	// Write set
-	err = binary.Write(stream, binary.BigEndian, b.set)
+	err = binary.Write(stream, binaryOrder, b.set)
 	return int64(b.BinaryStorageSize()), err
 }
 
@@ -682,7 +800,7 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
 	var length uint64
 
 	// Read length first
-	err := binary.Read(stream, binary.BigEndian, &length)
+	err := binary.Read(stream, binaryOrder, &length)
 	if err != nil {
 		return 0, err
 	}
@@ -693,7 +811,7 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
 	}
 
 	// Read remaining bytes as set
-	err = binary.Read(stream, binary.BigEndian, newset.set)
+	err = binary.Read(stream, binaryOrder, newset.set)
 	if err != nil {
 		return 0, err
 	}
@@ -736,7 +854,7 @@ func (b *BitSet) MarshalJSON() ([]byte, error) {
 	}
 
 	// URLEncode all bytes
-	return json.Marshal(base64.URLEncoding.EncodeToString(buffer.Bytes()))
+	return json.Marshal(base64Encoding.EncodeToString(buffer.Bytes()))
 }
 
 // UnmarshalJSON unmarshals a BitSet from JSON created using MarshalJSON
@@ -749,7 +867,7 @@ func (b *BitSet) UnmarshalJSON(data []byte) error {
 	}
 
 	// URLDecode string
-	buf, err := base64.URLEncoding.DecodeString(s)
+	buf, err := base64Encoding.DecodeString(s)
 	if err != nil {
 		return err
 	}