Skip to content

Commit ad3e1c4

Browse files
committed
Add shift operations to BitSet
1 parent 9ccba2c commit ad3e1c4

File tree

4 files changed

+201
-0
lines changed

4 files changed

+201
-0
lines changed

README.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ The following operations are provided:
4141
- `Map`
4242
- [`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)
4343
- [`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)
44+
- `BitSet`
45+
- [`<<`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/BitSetDecorator.html)
46+
- [`>>`](https://static.javadoc.io/org.scala-lang.modules/scala-collection-contrib_2.13/0.1.0/scala/collection/decorators/BitSetDecorator.html)
4447

4548
## Maintenance status
4649

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
package scala.collection.decorators
2+
3+
import scala.collection.{BitSet, BitSetOps}
4+
5+
class BitSetDecorator[+C <: BitSet with BitSetOps[C]](protected val bs: C) {
6+
7+
import BitSetDecorator._
8+
import BitSetOps._
9+
10+
/**
11+
* Bitwise left shift of this BitSet by given the shift distance.
12+
* The shift distance may be negative, in which case this method performs a right shift.
13+
* @param shiftBy shift distance, in bits
14+
* @return a new BitSet whose value is a bitwise shift left of this BitSet by given shift distance (`shiftBy`)
15+
*/
16+
def <<(shiftBy: Int): C = {
17+
18+
val shiftedBits = if (bs.nwords == 0 || bs.nwords == 1 && bs.word(0) == 0) Array.emptyLongArray
19+
else if (shiftBy > 0) shiftLeft(shiftBy)
20+
else if (shiftBy == 0) bs.toBitMask
21+
else shiftRight(-shiftBy)
22+
23+
bs.fromBitMaskNoCopy(shiftedBits)
24+
}
25+
26+
/**
27+
* Bitwise right shift of this BitSet by the given shift distance.
28+
* The shift distance may be negative, in which case this method performs a left shift.
29+
* @param shiftBy shift distance, in bits
30+
* @return a new BitSet whose value is a bitwise shift right of this BitSet by given shift distance (`shiftBy`)
31+
*/
32+
def >>(shiftBy: Int): C = {
33+
34+
val shiftedBits = if (bs.nwords == 0 || bs.nwords == 1 && bs.word(0) == 0) Array.emptyLongArray
35+
else if (shiftBy > 0) shiftRight(shiftBy)
36+
else if (shiftBy == 0) bs.toBitMask
37+
else shiftLeft(-shiftBy)
38+
39+
bs.fromBitMaskNoCopy(shiftedBits)
40+
}
41+
42+
private def shiftLeft(shiftBy: Int): Array[Long] = {
43+
44+
val bitOffset = shiftBy & WordMask
45+
val wordOffset = shiftBy >>> LogWL
46+
47+
if (bitOffset == 0) {
48+
val newSize = bs.nwords + wordOffset
49+
require(newSize <= MaxSize)
50+
val newBits = Array.ofDim[Long](newSize)
51+
var i = wordOffset
52+
while (i < newSize) {
53+
newBits(i) = bs.word(i - wordOffset)
54+
i += 1
55+
}
56+
newBits
57+
} else {
58+
val revBitOffset = WordLength - bitOffset
59+
val extraBits = bs.word(bs.nwords - 1) >>> revBitOffset
60+
val extraWordCount = if (extraBits == 0) 0 else 1
61+
val newSize = bs.nwords + wordOffset + extraWordCount
62+
require(newSize <= MaxSize)
63+
val newBits = Array.ofDim[Long](newSize)
64+
var previous = 0L
65+
var i = 0
66+
while (i < bs.nwords) {
67+
val current = bs.word(i)
68+
newBits(i + wordOffset) = (previous >>> revBitOffset) | (current << bitOffset)
69+
previous = current
70+
i += 1
71+
}
72+
if (extraWordCount != 0) newBits(newSize - 1) = extraBits
73+
newBits
74+
}
75+
}
76+
77+
private def shiftRight(shiftBy: Int): Array[Long] = {
78+
79+
val bitOffset = shiftBy & WordMask
80+
81+
if (bitOffset == 0) {
82+
val wordOffset = shiftBy >>> LogWL
83+
val newSize = bs.nwords - wordOffset
84+
if (newSize > 0) {
85+
val newBits = Array.ofDim[Long](newSize)
86+
var i = 0
87+
while (i < newSize) {
88+
newBits(i) = bs.word(i + wordOffset)
89+
i += 1
90+
}
91+
newBits
92+
} else Array.emptyLongArray
93+
} else {
94+
val wordOffset = (shiftBy >>> LogWL) + 1
95+
val extraBits = bs.word(bs.nwords - 1) >>> bitOffset
96+
val extraWordCount = if (extraBits == 0) 0 else 1
97+
val newSize = bs.nwords - wordOffset + extraWordCount
98+
if (newSize > 0) {
99+
val revBitOffset = WordLength - bitOffset
100+
val newBits = Array.ofDim[Long](newSize)
101+
var previous = bs.word(wordOffset - 1)
102+
var i = wordOffset
103+
while (i < bs.nwords) {
104+
val current = bs.word(i)
105+
newBits(i - wordOffset) = (previous >>> bitOffset) | (current << revBitOffset)
106+
previous = current
107+
i += 1
108+
}
109+
if (extraWordCount != 0) newBits(newSize - 1) = extraBits
110+
newBits
111+
} else Array.emptyLongArray
112+
}
113+
}
114+
115+
}
116+
117+
object BitSetDecorator {
118+
private[collection] final val WordMask = BitSetOps.WordLength - 1
119+
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,7 @@ package object decorators {
1717
implicit def MapDecorator[C](coll: C)(implicit map: IsMap[C]): MapDecorator[C, map.type] =
1818
new MapDecorator(coll)(map)
1919

20+
implicit def bitSetDecorator[C <: BitSet with BitSetOps[C]](bs: C): BitSetDecorator[C] =
21+
new BitSetDecorator(bs)
22+
2023
}
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package scala.collection.decorators
2+
3+
import org.junit.{Assert, Test}
4+
5+
import scala.collection.BitSet
6+
7+
class BitSetDecoratorTest {
8+
9+
import Assert.{assertEquals, assertSame}
10+
import BitSet.empty
11+
12+
@Test
13+
def shiftEmptyLeft(): Unit = {
14+
for (shiftBy <- 0 to 128) {
15+
assertSame(empty, empty << shiftBy)
16+
}
17+
}
18+
19+
@Test
20+
def shiftLowestBitLeft(): Unit = {
21+
for (shiftBy <- 0 to 128) {
22+
assertEquals(BitSet(shiftBy), BitSet(0) << shiftBy)
23+
}
24+
}
25+
26+
@Test
27+
def shiftNegativeLeft(): Unit = {
28+
assertEquals(BitSet(0), BitSet(1) << -1)
29+
}
30+
31+
@Test
32+
def largeShiftLeft(): Unit = {
33+
val bs = BitSet(0 to 300 by 5: _*)
34+
for (shiftBy <- 0 to 128) {
35+
assertEquals(bs.map(_ + shiftBy), bs << shiftBy)
36+
}
37+
}
38+
39+
@Test
40+
def shiftEmptyRight(): Unit = {
41+
for (shiftBy <- 0 to 128) {
42+
assertSame(empty, empty >> shiftBy)
43+
}
44+
}
45+
46+
@Test
47+
def shiftLowestBitRight(): Unit = {
48+
assertEquals(BitSet(0), BitSet(0) >> 0)
49+
for (shiftBy <- 1 to 128) {
50+
assertSame(empty, BitSet(0) >> shiftBy)
51+
}
52+
}
53+
54+
@Test
55+
def shiftToLowestBitRight(): Unit = {
56+
for (shiftBy <- 0 to 128) {
57+
assertEquals(BitSet(0), BitSet(shiftBy) >> shiftBy)
58+
}
59+
}
60+
61+
@Test
62+
def shiftNegativeRight(): Unit = {
63+
assertEquals(BitSet(1), BitSet(0) >> -1)
64+
}
65+
66+
@Test
67+
def largeShiftRight(): Unit = {
68+
val bs = BitSet(0 to 300 by 5: _*)
69+
for (shiftBy <- 0 to 128) {
70+
assertEquals(bs.collect {
71+
case b if b >= shiftBy => b - shiftBy
72+
}, bs >> shiftBy)
73+
}
74+
}
75+
76+
}

0 commit comments

Comments
 (0)