diff options
-rw-r--r-- | core/src/main/scala/org/apache/spark/util/hash/BitSet.scala | 82 |
1 files changed, 82 insertions, 0 deletions
diff --git a/core/src/main/scala/org/apache/spark/util/hash/BitSet.scala b/core/src/main/scala/org/apache/spark/util/hash/BitSet.scala index 0ec002b5d0..69b10566f3 100644 --- a/core/src/main/scala/org/apache/spark/util/hash/BitSet.scala +++ b/core/src/main/scala/org/apache/spark/util/hash/BitSet.scala @@ -28,6 +28,70 @@ class BitSet(numBits: Int) { private val numWords = words.length /** + * Compute the capacity (number of bits) that can be represented + * by this bitset. + */ + def capacity: Int = numWords * 64 + + + /** + * Set all the bits up to a given index + */ + def setUntil(bitIndex: Int) { + val wordIndex = bitIndex >> 6 // divide by 64 + var i = 0 + while(i < wordIndex) { words(i) = -1; i += 1 } + // Set the remaining bits + val mask = ~(-1L << (bitIndex & 0x3f)) + words(wordIndex) |= mask + } + + + /** + * Compute the bit-wise AND of the two sets returning the + * result. + */ + def &(other: BitSet): BitSet = { + val newBS = new BitSet(math.max(capacity, other.capacity)) + val smaller = math.min(numWords, other.numWords) + assert(newBS.numWords >= numWords) + assert(newBS.numWords >= other.numWords) + var ind = 0 + while( ind < smaller ) { + newBS.words(ind) = words(ind) & other.words(ind) + ind += 1 + } + newBS + } + + + /** + * Compute the bit-wise OR of the two sets returning the + * result. + */ + def |(other: BitSet): BitSet = { + val newBS = new BitSet(math.max(capacity, other.capacity)) + assert(newBS.numWords >= numWords) + assert(newBS.numWords >= other.numWords) + val smaller = math.min(numWords, other.numWords) + var ind = 0 + while( ind < smaller ) { + newBS.words(ind) = words(ind) | other.words(ind) + ind += 1 + } + while( ind < numWords ) { + newBS.words(ind) = words(ind) + ind += 1 + } + while( ind < other.numWords ) { + newBS.words(ind) = other.words(ind) + ind += 1 + } + newBS + } + + + /** * Sets the bit at the specified index to true. * @param index the bit index */ @@ -36,6 +100,7 @@ class BitSet(numBits: Int) { words(index >> 6) |= bitmask // div by 64 and mask } + /** * Return the value of the bit with the specified index. The value is true if the bit with * the index is currently set in this BitSet; otherwise, the result is false. @@ -48,6 +113,21 @@ class BitSet(numBits: Int) { (words(index >>> 6) & bitmask) != 0 // div by 64 and mask } + + /** + * Get an iterator over the set bits. + */ + def iterator = new Iterator[Int] { + var ind = nextSetBit(0) + override def hasNext: Boolean = ind >= 0 + override def next() = { + val tmp = ind + ind = nextSetBit(ind+1) + tmp + } + } + + /** Return the number of bits set to true in this BitSet. */ def cardinality(): Int = { var sum = 0 @@ -59,6 +139,7 @@ class BitSet(numBits: Int) { sum } + /** * Returns the index of the first bit that is set to true that occurs on or after the * specified starting index. If no such bit exists then -1 is returned. @@ -98,6 +179,7 @@ class BitSet(numBits: Int) { -1 } + /** Return the number of longs it would take to hold numBits. */ private def bit2words(numBits: Int) = ((numBits - 1) >>> 6) + 1 } |