Add shift operations to BitSet

README.md

@@ -41,6 +41,9 @@ The following operations are provided:
 - `Map`
     - [`zipByKey`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#zipByKey[W,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,\(MapDecorator.this.map.V,W\)\),That]\):That) / [`join`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#join[W,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,\(MapDecorator.this.map.V,W\)\),That]\):That) / [`zipByKeyWith`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#zipByKeyWith[W,X,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(f:\(MapDecorator.this.map.V,W\)=>X\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,X\),That]\):That)
     - [`mergeByKey`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#mergeByKey[W,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,\(Option[MapDecorator.this.map.V],Option[W]\)\),That]\):That) / [`fullOuterJoin`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#fullOuterJoin[W,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,\(Option[MapDecorator.this.map.V],Option[W]\)\),That]\):That) / [`mergeByKeyWith`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#mergeByKeyWith[W,X,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(f:PartialFunction[\(Option[MapDecorator.this.map.V],Option[W]\),X]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,X\),That]\):That) / [`leftOuterJoin`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#leftOuterJoin[W,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,\(MapDecorator.this.map.V,Option[W]\)\),That]\):That) / [`rightOuterJoin`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/MapDecorator.html#rightOuterJoin[W,That]\(other:scala.collection.Map[MapDecorator.this.map.K,W]\)\(implicitbf:scala.collection.BuildFrom[C,\(MapDecorator.this.map.K,\(Option[MapDecorator.this.map.V],W\)\),That]\):That)
+- `BitSet`
+    - [`<<`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/BitSetDecorator.html)
+    - [`>>`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/BitSetDecorator.html)
 
 ## Maintenance status
 

src/main/scala/scala/collection/decorators/BitSetDecorator.scala

@@ -0,0 +1,119 @@
+package scala.collection.decorators
+
+import scala.collection.{BitSet, BitSetOps}
+
+class BitSetDecorator[+C <: BitSet with BitSetOps[C]](protected val bs: C) {
+
+  import BitSetDecorator._
+  import BitSetOps._
+
+  /**
+    * Bitwise left shift of this BitSet by given the shift distance.
+    * The shift distance may be negative, in which case this method performs a right shift.
+    * @param shiftBy shift distance, in bits
+    * @return a new BitSet whose value is a bitwise shift left of this BitSet by given shift distance (`shiftBy`)
+    */
+  def <<(shiftBy: Int): C = {
+
+    val shiftedBits = if (bs.nwords == 0 || bs.nwords == 1 && bs.word(0) == 0) Array.emptyLongArray
+    else if (shiftBy > 0) shiftLeft(shiftBy)
+    else if (shiftBy == 0) bs.toBitMask
+    else shiftRight(-shiftBy)
+
+    bs.fromBitMaskNoCopy(shiftedBits)
+  }
+
+  /**
+    * Bitwise right shift of this BitSet by the given shift distance.
+    * The shift distance may be negative, in which case this method performs a left shift.
+    * @param shiftBy shift distance, in bits
+    * @return a new BitSet whose value is a bitwise shift right of this BitSet by given shift distance (`shiftBy`)
+    */
+  def >>(shiftBy: Int): C = {
+
+    val shiftedBits = if (bs.nwords == 0 || bs.nwords == 1 && bs.word(0) == 0) Array.emptyLongArray
+    else if (shiftBy > 0) shiftRight(shiftBy)
+    else if (shiftBy == 0) bs.toBitMask
+    else shiftLeft(-shiftBy)
+
+    bs.fromBitMaskNoCopy(shiftedBits)
+  }
+
+  private def shiftLeft(shiftBy: Int): Array[Long] = {
+
+    val bitOffset = shiftBy & WordMask
+    val wordOffset = shiftBy >>> LogWL
+
+    if (bitOffset == 0) {
+      val newSize = bs.nwords + wordOffset
+      require(newSize <= MaxSize)
+      val newBits = Array.ofDim[Long](newSize)
+      var i = wordOffset
+      while (i < newSize) {
+        newBits(i) = bs.word(i - wordOffset)
+        i += 1
+      }
+      newBits
+    } else {
+      val revBitOffset = WordLength - bitOffset
+      val extraBits = bs.word(bs.nwords - 1) >>> revBitOffset
+      val extraWordCount = if (extraBits == 0) 0 else 1
+      val newSize = bs.nwords + wordOffset + extraWordCount
+      require(newSize <= MaxSize)
+      val newBits = Array.ofDim[Long](newSize)
+      var previous = 0L
+      var i = 0
+      while (i < bs.nwords) {
+        val current = bs.word(i)
+        newBits(i + wordOffset) = (previous >>> revBitOffset) | (current << bitOffset)
+        previous = current
+        i += 1
+      }
+      if (extraWordCount != 0) newBits(newSize - 1) = extraBits
+      newBits
+    }
+  }
+
+  private def shiftRight(shiftBy: Int): Array[Long] = {
+
+    val bitOffset = shiftBy & WordMask
+
+    if (bitOffset == 0) {
+      val wordOffset = shiftBy >>> LogWL
+      val newSize = bs.nwords - wordOffset
+      if (newSize > 0) {
+        val newBits = Array.ofDim[Long](newSize)
+        var i = 0
+        while (i < newSize) {
+          newBits(i) = bs.word(i + wordOffset)
+          i += 1
+        }
+        newBits
+      } else Array.emptyLongArray
+    } else {
+      val wordOffset = (shiftBy >>> LogWL) + 1
+      val extraBits = bs.word(bs.nwords - 1) >>> bitOffset
+      val extraWordCount = if (extraBits == 0) 0 else 1
+      val newSize = bs.nwords - wordOffset + extraWordCount
+      if (newSize > 0) {
+        val revBitOffset = WordLength - bitOffset
+        val newBits = Array.ofDim[Long](newSize)
+        var previous = bs.word(wordOffset - 1)
+        var i = wordOffset
+        while (i < bs.nwords) {
+          val current = bs.word(i)
+          newBits(i - wordOffset) = (previous >>> bitOffset) | (current << revBitOffset)
+          previous = current
+          i += 1
+        }
+        if (extraWordCount != 0) newBits(newSize - 1) = extraBits
+        newBits
+      } else Array.emptyLongArray
+    }
+  }
+
+}
+
+object BitSetDecorator {
+  private[collection] final val WordMask = BitSetOps.WordLength - 1
+}

src/main/scala/scala/collection/decorators/package.scala

@@ -17,4 +17,7 @@ package object decorators {
   implicit def MapDecorator[C](coll: C)(implicit map: IsMap[C]): MapDecorator[C, map.type] =
     new MapDecorator(coll)(map)
 
+  implicit def bitSetDecorator[C <: BitSet with BitSetOps[C]](bs: C): BitSetDecorator[C] =
+    new BitSetDecorator(bs)
+
 }

src/test/scala/scala/collection/decorators/BitSetDecoratorTest.scala

@@ -0,0 +1,76 @@
+package scala.collection.decorators
+
+import org.junit.{Assert, Test}
+
+import scala.collection.BitSet
+
+class BitSetDecoratorTest {
+
+  import Assert.{assertEquals, assertSame}
+  import BitSet.empty
+
+  @Test
+  def shiftEmptyLeft(): Unit = {
+    for (shiftBy <- 0 to 128) {
+      assertSame(empty, empty << shiftBy)
+    }
+  }
+
+  @Test
+  def shiftLowestBitLeft(): Unit = {
+    for (shiftBy <- 0 to 128) {
+      assertEquals(BitSet(shiftBy), BitSet(0) << shiftBy)
+    }
+  }
+
+  @Test
+  def shiftNegativeLeft(): Unit = {
+    assertEquals(BitSet(0), BitSet(1) << -1)
+  }
+
+  @Test
+  def largeShiftLeft(): Unit = {
+    val bs = BitSet(0 to 300 by 5: _*)
+    for (shiftBy <- 0 to 128) {
+      assertEquals(bs.map(_ + shiftBy), bs << shiftBy)
+    }
+  }
+
+  @Test
+  def shiftEmptyRight(): Unit = {
+    for (shiftBy <- 0 to 128) {
+      assertSame(empty, empty >> shiftBy)
+    }
+  }
+
+  @Test
+  def shiftLowestBitRight(): Unit = {
+    assertEquals(BitSet(0), BitSet(0) >> 0)
+    for (shiftBy <- 1 to 128) {
+      assertSame(empty, BitSet(0) >> shiftBy)
+    }
+  }
+
+  @Test
+  def shiftToLowestBitRight(): Unit = {
+    for (shiftBy <- 0 to 128) {
+      assertEquals(BitSet(0), BitSet(shiftBy) >> shiftBy)
+    }
+  }
+
+  @Test
+  def shiftNegativeRight(): Unit = {
+    assertEquals(BitSet(1), BitSet(0) >> -1)
+  }
+
+  @Test
+  def largeShiftRight(): Unit = {
+    val bs = BitSet(0 to 300 by 5: _*)
+    for (shiftBy <- 0 to 128) {
+      assertEquals(bs.collect {
+        case b if b >= shiftBy => b - shiftBy
+      }, bs >> shiftBy)
+    }
+  }
+
+}