Back up Chunk[Boolean] by Chunk.BitChunks (#3419)

zio · Apr 22, 2020 · 428ceb4 · 428ceb4
1 parent 4995015
commit 428ceb4
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 33 deletions.
diff --git a/core/shared/src/main/scala-2.11-2.12/ChunkBuilder.scala b/core/shared/src/main/scala-2.11-2.12/ChunkBuilder.scala
@@ -19,15 +19,17 @@ package zio
 import scala.collection.mutable.{ ArrayBuilder, Builder }
 import scala.{
   Boolean => SBoolean,
-  Short => SShort,
-  Int => SInt,
+  Byte => SByte,
+  Char => SChar,
   Double => SDouble,
   Float => SFloat,
+  Int => SInt,
   Long => SLong,
-  Char => SChar,
-  Byte => SByte
+  Short => SShort
 }
 
+import zio.Chunk.BitChunk
+
 /**
  * A `ChunkBuilder[A]` can build a `Chunk[A]` given elements of type `A`.
  * `ChunkBuilder` is a mutable data structure that is implemented to
@@ -79,28 +81,60 @@ private[zio] object ChunkBuilder {
    * values.
    */
   final class Boolean extends ChunkBuilder[SBoolean] { self =>
-    private val arrayBuilder: ArrayBuilder[SBoolean] = {
-      new ArrayBuilder.ofBoolean
+
+    private val arrayBuilder: ArrayBuilder[SByte] = {
+      new ArrayBuilder.ofByte
     }
-    def +=(a: SBoolean): this.type = {
-      arrayBuilder += a
+    private var lastByte: SByte   = 0.toByte
+    private var maxBitIndex: SInt = 0
+
+    def +=(b: SBoolean): this.type = {
+      if (b) {
+        if (maxBitIndex == 8) {
+          arrayBuilder += lastByte
+          lastByte = (1 << 7).toByte
+          maxBitIndex = 1
+        } else {
+          val bitIndex = 7 - maxBitIndex
+          lastByte = (lastByte | (1 << bitIndex)).toByte
+          maxBitIndex = maxBitIndex + 1
+        }
+      } else {
+        if (maxBitIndex == 8) {
+          arrayBuilder += lastByte
+          lastByte = 0.toByte
+          maxBitIndex = 1
+        } else {
+          maxBitIndex = maxBitIndex + 1
+        }
+      }
       this
     }
-    def clear(): Unit =
+    def clear(): Unit = {
       arrayBuilder.clear()
-    def result(): Chunk[SBoolean] =
-      Chunk.fromArray(arrayBuilder.result())
+      maxBitIndex = 0
+      lastByte = 0.toByte
+    }
+    def result(): Chunk[SBoolean] = {
+      val bytes: Chunk[SByte] = Chunk.fromArray(arrayBuilder.result() :+ lastByte)
+      BitChunk(bytes, 0, 8 * (bytes.length - 1) + maxBitIndex)
+    }
     override def ++=(as: TraversableOnce[SBoolean]): this.type = {
-      arrayBuilder ++= as
+      as.foreach(+= _)
       this
     }
+    override def sizeHint(n: SInt): Unit = {
+      val hint = if (n == 0) 0 else n / 8 + 1
+      arrayBuilder.sizeHint(hint)
+    }
     override def equals(that: Any): SBoolean =
       that match {
-        case that: Boolean => self.arrayBuilder == that.arrayBuilder
-        case _             => false
+        case that: Boolean =>
+          self.arrayBuilder.equals(that.arrayBuilder) &&
+            self.maxBitIndex == that.maxBitIndex &&
+            self.lastByte == that.lastByte
+        case _ => false
       }
-    override def sizeHint(n: SInt): Unit =
-      arrayBuilder.sizeHint(n)
     override def toString: String =
       "ChunkBuilder.Boolean"
   }

diff --git a/core/shared/src/main/scala-2.13+/ChunkBuilder.scala b/core/shared/src/main/scala-2.13+/ChunkBuilder.scala
@@ -19,15 +19,17 @@ package zio
 import scala.collection.mutable.{ ArrayBuilder, Builder }
 import scala.{
   Boolean => SBoolean,
-  Short => SShort,
-  Int => SInt,
+  Byte => SByte,
+  Char => SChar,
   Double => SDouble,
   Float => SFloat,
+  Int => SInt,
   Long => SLong,
-  Char => SChar,
-  Byte => SByte
+  Short => SShort
 }
 
+import zio.Chunk.BitChunk
+
 /**
  * A `ChunkBuilder[A]` can build a `Chunk[A]` given elements of type `A`.
  * `ChunkBuilder` is a mutable data structure that is implemented to
@@ -79,28 +81,60 @@ private[zio] object ChunkBuilder {
    * values.
    */
   final class Boolean extends ChunkBuilder[SBoolean] { self =>
-    private val arrayBuilder: ArrayBuilder[SBoolean] = {
-      new ArrayBuilder.ofBoolean
+
+    private val arrayBuilder: ArrayBuilder[SByte] = {
+      new ArrayBuilder.ofByte
     }
-    def addOne(a: SBoolean): this.type = {
-      arrayBuilder += a
+    private var lastByte: SByte   = 0.toByte
+    private var maxBitIndex: SInt = 0
+
+    def addOne(b: SBoolean): this.type = {
+      if (b) {
+        if (maxBitIndex == 8) {
+          arrayBuilder += lastByte
+          lastByte = (1 << 7).toByte
+          maxBitIndex = 1
+        } else {
+          val bitIndex = 7 - maxBitIndex
+          lastByte = (lastByte | (1 << bitIndex)).toByte
+          maxBitIndex = maxBitIndex + 1
+        }
+      } else {
+        if (maxBitIndex == 8) {
+          arrayBuilder += lastByte
+          lastByte = 0.toByte
+          maxBitIndex = 1
+        } else {
+          maxBitIndex = maxBitIndex + 1
+        }
+      }
       this
     }
-    def clear(): Unit =
+    def clear(): Unit = {
       arrayBuilder.clear()
-    def result(): Chunk[SBoolean] =
-      Chunk.fromArray(arrayBuilder.result())
+      maxBitIndex = 0
+      lastByte = 0.toByte
+    }
+    def result(): Chunk[SBoolean] = {
+      val bytes: Chunk[SByte] = Chunk.fromArray(arrayBuilder.result() :+ lastByte)
+      BitChunk(bytes, 0, 8 * (bytes.length - 1) + maxBitIndex)
+    }
     override def addAll(as: IterableOnce[SBoolean]): this.type = {
-      arrayBuilder ++= as
+      as.iterator.foreach(addOne _)
       this
     }
+    override def sizeHint(n: SInt): Unit = {
+      val hint = if (n == 0) 0 else n / 8 + 1
+      arrayBuilder.sizeHint(hint)
+    }
     override def equals(that: Any): SBoolean =
       that match {
-        case that: Boolean => self.arrayBuilder == that.arrayBuilder
-        case _             => false
+        case that: Boolean =>
+          self.arrayBuilder.equals(that.arrayBuilder) &&
+          self.maxBitIndex == that.maxBitIndex &&
+          self.lastByte == that.lastByte
+        case _ => false
       }
-    override def sizeHint(n: SInt): Unit =
-      arrayBuilder.sizeHint(n)
     override def toString: String =
       "ChunkBuilder.Boolean"
   }

diff --git a/core/shared/src/main/scala/zio/Chunk.scala b/core/shared/src/main/scala/zio/Chunk.scala
@@ -1224,7 +1224,8 @@ object Chunk {
     }
   }
 
-  private final case class BitChunk(bytes: Chunk[Byte], minBitIndex: Int, maxBitIndex: Int) extends Chunk[Boolean] {
+  private[zio] final case class BitChunk(bytes: Chunk[Byte], minBitIndex: Int, maxBitIndex: Int)
+      extends Chunk[Boolean] {
     self =>
 
     override val length: Int =