bitset fixes

3 files changed, 23 insertions, 6 deletions

diff --git a/sources/scala/collection/BitSet.scala b/sources/scala/collection/BitSet.scala
index 7612335e6d..1fbb32be37 100644
--- a/sources/scala/collection/BitSet.scala
+++ b/sources/scala/collection/BitSet.scala
@@ -25,10 +25,22 @@ abstract class BitSet with Function1[Int,Boolean] {
   def apply(i: Int):Boolean;
 
   /** returns an iterator over the truth values of all bits */
-  final def booleanElements:Iterator[Boolean] = new Iterator[Boolean] {
+  final def booleanElements: Iterator[Boolean] = new Iterator[Boolean] {
     var i = 0;
     def hasNext: Boolean = i < size;
     def next: Boolean = { i = i + 1; apply(i-1) }
   }
 
+  /** returns the subset of [0..size] whose elements are indices of bits set to v */
+  final def toSet( v:Boolean ) = {
+    var res = new immutable.TreeSet[Int]();
+    var j = 0;
+    while( j < size ) {
+      if( v == apply(j) )
+        res = res + j;
+      j = j + 1;
+    }
+    res
+  }
+
 }
diff --git a/sources/scala/collection/immutable/BitSet.scala b/sources/scala/collection/immutable/BitSet.scala
index 3bc0cb3c8e..194cca7e60 100644
--- a/sources/scala/collection/immutable/BitSet.scala
+++ b/sources/scala/collection/immutable/BitSet.scala
@@ -18,7 +18,8 @@ package scala.collection.immutable ;
  */
 class BitSet(n:Int, ba: Array[Byte], copy:Boolean) extends scala.collection.BitSet {
 
-  val size   = n;
+  final def size   = n;
+
   val array:Array[Byte]  =
     if( copy ) {
       val arr = new Array[Byte](ba.length);
diff --git a/sources/scala/collection/mutable/ResizableBitSet.scala b/sources/scala/collection/mutable/ResizableBitSet.scala
index a8379bd6c6..a92184b441 100644
--- a/sources/scala/collection/mutable/ResizableBitSet.scala
+++ b/sources/scala/collection/mutable/ResizableBitSet.scala
@@ -18,12 +18,14 @@ class ResizableBitSet(initSize: Int) extends scala.collection.BitSet {
   /** default constructor, initial size of 16 bits */
   def this() = this( 16 );
 
+  final def byteSize(size:Int) = { (size >>> 3) + (if( (size & 0x07)!= 0 ) 1 else 0) };
+
   class ByteArray with ResizableArray[Byte] {
-    override protected val initialSize: Int = initSize >>> 3;
+    override protected val initialSize: Int = byteSize( initSize );
     override protected var array: Array[Byte] = new Array[Byte](initialSize);
 
     /** size of this bitset in nbits */
-    def ensureBits(nbits: Int): Unit = ensureSize(nbits >>> 3);
+    def ensureBits(nbits: Int): Unit = ensureSize( byteSize( nbits ));
 
     final def and(j: Int, mask:Int): Unit = {
       array.update( j, (array(j) & mask).asInstanceOf[Byte] );
@@ -45,23 +47,25 @@ class ResizableBitSet(initSize: Int) extends scala.collection.BitSet {
   protected val internal = new ByteArray();
 
   /** size of this bitset in nbytes */
-  var size: Int = 0;
+  var size: Int = initSize;
 
   /** size of this bitset in nbits */
   def ensureSize(nbits: Int): Unit = {
     internal.ensureBits( nbits );
-    size = nbits;
+    if( size < nbits ) size = nbits;
   }
 
   final def set(i: Int, b: Boolean): Unit = if( b ) set(i) else clear(i);
 
   final def set(i: Int): Unit = {
+    ensureSize(i+1);
     val j         = (i >>> 3);
     val mask      = (1 << (i & 0x07));
     internal.or(j, mask);
   }
 
   def clear(i: Int): Unit = {
+    ensureSize(i+1);
     val j    = (i >>> 3);
     val mask = (1 << (i & 0x07));
     internal.and(j, ~mask);