Adding logical AND/OR, setUntil, and iterators to the BitSet.

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
 }