zio · iravid · Aug 13, 2020 · Jul 29, 2020
diff --git a/streams-tests/jvm/src/test/scala/zio/stream/compression/CompressionSpec.scala b/streams-tests/jvm/src/test/scala/zio/stream/compression/CompressionSpec.scala
@@ -3,14 +3,20 @@ package zio.stream.compression
 import java.io.{ ByteArrayInputStream, ByteArrayOutputStream }
 import java.nio.charset.StandardCharsets
 import java.util.Arrays
-import java.util.zip.{ CRC32, Deflater, DeflaterInputStream, GZIPOutputStream, Inflater, InflaterInputStream }
+import java.util.zip.{
+  CRC32,
+  Deflater,
+  DeflaterInputStream,
+  GZIPInputStream,
+  GZIPOutputStream,
+  Inflater,
+  InflaterInputStream
+}
 
 import scala.annotation.tailrec
-
 import TestData._
-
 import zio._
-import zio.stream.ZTransducer.{ deflate, gunzip, inflate }
+import zio.stream.ZTransducer.{ deflate, gunzip, gzip, inflate }
 import zio.stream._
 import zio.test.Assertion._
 import zio.test._
@@ -82,39 +88,39 @@ object CompressionSpec extends DefaultRunnableSpec {
         ),
         testM("fail if input stream finished unexpected")(
           assertM(
-            gzippedStream(longText).take(800).transduce(inflate()).runCollect.run
+            jdkGzippedStream(longText).take(800).transduce(inflate()).runCollect.run
           )(fails(anything))
         )
       ),
       suite("gunzip")(
         testM("short stream")(
           assertM(
-            gzippedStream(shortText).transduce(gunzip(64)).runCollect
+            jdkGzippedStream(shortText).transduce(gunzip(64)).runCollect
           )(equalTo(Chunk.fromArray(shortText)))
         ),
         testM("stream of two gzipped inputs")(
           assertM(
-            (gzippedStream(shortText) ++ gzippedStream(otherShortText)).transduce(gunzip(64)).runCollect
+            (jdkGzippedStream(shortText) ++ jdkGzippedStream(otherShortText)).transduce(gunzip(64)).runCollect
           )(equalTo(Chunk.fromArray(shortText) ++ Chunk.fromArray(otherShortText)))
         ),
         testM("long input")(
           assertM(
-            gzippedStream(longText).transduce(gunzip(64)).runCollect
+            jdkGzippedStream(longText).transduce(gunzip(64)).runCollect
           )(equalTo(Chunk.fromArray(longText)))
         ),
         testM("long input, no SYNC_FLUSH")(
           assertM(
-            gzippedStream(longText, false).transduce(gunzip(64)).runCollect
+            jdkGzippedStream(longText, false).transduce(gunzip(64)).runCollect
           )(equalTo(Chunk.fromArray(longText)))
         ),
         testM("long input, buffer smaller than chunks")(
           assertM(
-            gzippedStream(longText).chunkN(500).transduce(gunzip(1)).runCollect
+            jdkGzippedStream(longText).chunkN(500).transduce(gunzip(1)).runCollect
           )(equalTo(Chunk.fromArray(longText)))
         ),
         testM("long input, chunks smaller then buffer")(
           assertM(
-            gzippedStream(longText).chunkN(1).transduce(gunzip(500)).runCollect
+            jdkGzippedStream(longText).chunkN(1).transduce(gunzip(500)).runCollect
           )(equalTo(Chunk.fromArray(longText)))
         ),
         testM("fail early if header is corrupted")(
@@ -124,7 +130,7 @@ object CompressionSpec extends DefaultRunnableSpec {
         ),
         testM("fail if input stream finished unexpected")(
           assertM(
-            gzippedStream(longText).take(80).transduce(gunzip()).runCollect.run
+            jdkGzippedStream(longText).take(80).transduce(gunzip()).runCollect.run
           )(fails(anything))
         ),
         testM("no output on very incomplete stream is not OK")(
@@ -136,7 +142,7 @@ object CompressionSpec extends DefaultRunnableSpec {
           checkM(Gen.listOfBounded(0, `1K`)(Gen.anyByte).zip(Gen.int(1, `1K`)).zip(Gen.int(1, `1K`))) {
             case ((chunk, n), bufferSize) =>
               assertM(for {
-                deflated <- ZIO.effectTotal(gzippedStream(chunk.toArray))
+                deflated <- ZIO.effectTotal(jdkGzippedStream(chunk.toArray))
                 out      <- deflated.chunkN(n).transduce(gunzip(bufferSize)).runCollect
               } yield out.toList)(equalTo(chunk))
           }
@@ -180,28 +186,78 @@ object CompressionSpec extends DefaultRunnableSpec {
         testM("deflate empty bytes, small buffer")(
           assertM(
             Stream.fromIterable(List.empty).chunkN(1).transduce(deflate(100, false)).runCollect.map(_.toList)
-          )(equalTo(deflateJDK(Array.empty, new Deflater(-1, false)).toList))
+          )(equalTo(jdkDeflate(Array.empty, new Deflater(-1, false)).toList))
         ),
         testM("deflates same as JDK")(
           assertM(Stream.fromIterable(longText).chunkN(128).transduce(deflate(256, false)).runCollect)(
-            equalTo(Chunk.fromArray(deflateJDK(longText, new Deflater(-1, false))))
+            equalTo(Chunk.fromArray(jdkDeflate(longText, new Deflater(-1, false))))
           )
         ),
         testM("deflates same as JDK, nowrap")(
           assertM(Stream.fromIterable(longText).chunkN(128).transduce(deflate(256, true)).runCollect)(
-            equalTo(Chunk.fromArray(deflateJDK(longText, new Deflater(-1, true))))
+            equalTo(Chunk.fromArray(jdkDeflate(longText, new Deflater(-1, true))))
           )
         ),
         testM("deflates same as JDK, small buffer")(
           assertM(Stream.fromIterable(longText).chunkN(64).transduce(deflate(1, false)).runCollect)(
-            equalTo(Chunk.fromArray(deflateJDK(longText, new Deflater(-1, false))))
+            equalTo(Chunk.fromArray(jdkDeflate(longText, new Deflater(-1, false))))
           )
         ),
         testM("deflates same as JDK, nowrap, small buffer ")(
           assertM(Stream.fromIterable(longText).chunkN(64).transduce(deflate(1, true)).runCollect)(
-            equalTo(Chunk.fromArray(deflateJDK(longText, new Deflater(-1, true))))
+            equalTo(Chunk.fromArray(jdkDeflate(longText, new Deflater(-1, true))))
           )
         )
+      ),
+      suite("gzip")(
+        testM("JDK gunzips what was gzipped")(
+          checkM(Gen.listOfBounded(0, `1K`)(Gen.anyByte).zip(Gen.int(1, `1K`)).zip(Gen.int(1, `1K`))) {
+            case ((input, n), bufferSize) =>
+              assertM(for {
+                gzipped  <- Stream.fromIterable(input).chunkN(n).transduce(gzip(bufferSize)).runCollect
+                inflated <- jdkGunzip(gzipped)
+              } yield inflated)(equalTo(input))
+          }
+        ),
+        testM("gzip empty bytes, small buffer")(
+          assertM(for {
+            gzipped      <- Stream.empty.transduce(gzip(1)).runCollect
+            jdkGunzipped <- jdkGunzip(gzipped)
+          } yield jdkGunzipped)(isEmpty)
+        ),
+        testM("gzip empty bytes")(
+          assertM(for {
+            gzipped      <- Stream.empty.transduce(gzip(`1K`)).runCollect
+            jdkGunzipped <- jdkGunzip(gzipped)
+          } yield jdkGunzipped)(isEmpty)
+        ),
+        testM("gzips, small chunks, small buffer")(
+          assertM(for {
+            gzipped      <- Stream.fromIterable(longText).chunkN(1).transduce(gzip(1)).runCollect
+            jdkGunzipped <- jdkGunzip(gzipped)
+          } yield jdkGunzipped)(equalTo(longText.toList))
+        ),
+        testM("gzips, small chunks, 1k buffer")(
+          assertM(for {
+            gzipped      <- Stream.fromIterable(longText).chunkN(1).transduce(gzip(`1K`)).runCollect
+            jdkGunzipped <- jdkGunzip(gzipped)
+          } yield jdkGunzipped)(equalTo(longText.toList))
+        ),
+        testM("chunks bigger than buffer")(
+          assertM(for {
+            gzipped      <- Stream.fromIterable(longText).chunkN(`1K`).transduce(gzip(64)).runCollect
+            jdkGunzipped <- jdkGunzip(gzipped)
+          } yield jdkGunzipped)(equalTo(longText.toList))
+        ),
+        testM("transducer is re-usable")(
+          assertM(for {
+            gzipper       <- ZIO.effectTotal(gzip(64))
+            gzipped1      <- Stream.fromIterable(longText).transduce(gzipper).runCollect
+            gzipped2      <- Stream.fromIterable(longText.reverse).transduce(gzipper).runCollect
+            jdkGunzipped1 <- jdkGunzip(gzipped1)
+            jdkGunzipped2 <- jdkGunzip(gzipped2)
+          } yield jdkGunzipped1 ++ jdkGunzipped2)(equalTo(longText.toList ++ longText.reverse.toList))
+        )
       )
     )
 
@@ -220,15 +276,15 @@ object TestData {
   def noWrapDeflatedStream(bytes: Array[Byte]) =
     deflatedWith(bytes, new Deflater(9, true))
 
-  def deflateJDK(bytes: Array[Byte], deflater: Deflater): Array[Byte] = {
+  def jdkDeflate(bytes: Array[Byte], deflater: Deflater): Array[Byte] = {
     val bigBuffer = new Array[Byte](1024 * 1024)
     val dif       = new DeflaterInputStream(new ByteArrayInputStream(bytes), deflater)
     val read      = dif.read(bigBuffer, 0, bigBuffer.length)
     Arrays.copyOf(bigBuffer, read)
   }
 
   def deflatedWith(bytes: Array[Byte], deflater: Deflater) = {
-    val arr = deflateJDK(bytes, deflater)
+    val arr = jdkDeflate(bytes, deflater)
     ZStream.fromIterable(arr)
   }
 
@@ -247,10 +303,10 @@ object TestData {
     inflate(Nil)
   }
 
-  def gzippedStream(bytes: Array[Byte], syncFlush: Boolean = true) =
-    ZStream.fromIterable(gzip(bytes, syncFlush))
+  def jdkGzippedStream(bytes: Array[Byte], syncFlush: Boolean = true) =
+    ZStream.fromIterable(jdkGzip(bytes, syncFlush))
 
-  def gzip(bytes: Array[Byte], syncFlush: Boolean = true): Array[Byte] = {
+  def jdkGzip(bytes: Array[Byte], syncFlush: Boolean = true): Array[Byte] = {
     val baos = new ByteArrayOutputStream(1024)
     val gzos = new GZIPOutputStream(baos, 1024, syncFlush)
     gzos.write(bytes)
@@ -259,6 +315,20 @@ object TestData {
     baos.toByteArray()
   }
 
+  def jdkGunzip(gzipped: Chunk[Byte]): Task[List[Byte]] = ZIO.effect {
+    val bigBuffer = new Array[Byte](1024 * 1024)
+    val bais      = new ByteArrayInputStream(gzipped.toArray)
+    val gzis      = new GZIPInputStream(bais)
+
+    @tailrec
+    def gunzip(acc: List[Byte]): List[Byte] = {
+      val read = gzis.read(bigBuffer, 0, bigBuffer.length)
+      if (read <= 0) acc
+      else gunzip(acc ++ bigBuffer.take(read).toList)
+    }
+    gunzip(Nil)
+  }
+
   val shortText          = "abcdefg1234567890".getBytes
   val otherShortText     = "AXXX\u0000XXXA".getBytes
   val longText           = Array.fill(1000)(shortText).flatten
@@ -268,7 +338,7 @@ object TestData {
 
   def makeStreamWithCustomHeader(flag: Int, headerTail: Array[Byte]) = {
     val headerHead = Array(31, 139, 8, flag, 0, 0, 0, 0, 0, 0).map(_.toByte)
-    ZStream.fromIterable(headerHead ++ headerTail ++ gzip(shortText).drop(10))
+    ZStream.fromIterable(headerHead ++ headerTail ++ jdkGzip(shortText).drop(10))
   }
 
   val headerWithExtra =
@@ -289,7 +359,7 @@ object TestData {
     val crc16Byte1 = (crc16 & 0xff).toByte
     val crc16Byte2 = (crc16 >> 8).toByte
     val header     = headerBytes ++ Array(crc16Byte1, crc16Byte2)
-    ZStream.fromIterable(header ++ gzip(shortText).drop(10))
+    ZStream.fromIterable(header ++ jdkGzip(shortText).drop(10))
   }
 
   val headerWithAll = {
@@ -305,7 +375,7 @@ object TestData {
     val crc16Byte1 = (crc16 & 0xff).toByte
     val crc16Byte2 = (crc16 >> 8).toByte
     val header     = headerUpToCrc ++ Array(crc16Byte1, crc16Byte2)
-    ZStream.fromIterable(header ++ gzip(shortText).drop(10))
+    ZStream.fromIterable(header ++ jdkGzip(shortText).drop(10))
   }
 
   def u8(b: Byte): Int = b & 0xff

diff --git a/streams/jvm/src/main/scala/zio/stream/compression/Deflate.scala b/streams/jvm/src/main/scala/zio/stream/compression/Deflate.scala
@@ -0,0 +1,56 @@
+package zio.stream.compression
+
+import java.{ util => ju }
+import java.util.zip.Deflater
+
+import zio.{ Chunk, ZIO, ZManaged }
+
+import scala.annotation.tailrec
+
+object Deflate {
+
+  def makeDeflater(
+    bufferSize: Int = 64 * 1024,
+    noWrap: Boolean = false,
+    level: CompressionLevel,
+    strategy: CompressionStrategy,
+    flushMode: FlushMode
+  ): ZManaged[Any, Nothing, Option[Chunk[Byte]] => ZIO[Any, Nothing, Chunk[Byte]]] =
+    ZManaged
+      .make(ZIO.effectTotal {
+        val deflater = new Deflater(level.jValue, noWrap)
+        deflater.setStrategy(strategy.jValue)
+        (deflater, new Array[Byte](bufferSize))
+      }) {
+        case (deflater, _) => ZIO.effectTotal(deflater.end())
+      }
+      .map {
+        case (deflater, buffer) => {
+          case Some(chunk) =>
+            ZIO.effectTotal {
+              deflater.setInput(chunk.toArray)
+              Deflate.pullOutput(deflater, buffer, flushMode)
+            }
+          case None =>
+            ZIO.effectTotal {
+              deflater.finish()
+              val out = Deflate.pullOutput(deflater, buffer, flushMode)
+              deflater.reset()
+              out
+            }
+        }
+      }
+
+  private[compression] def pullOutput(deflater: Deflater, buffer: Array[Byte], flushMode: FlushMode): Chunk[Byte] = {
+    @tailrec
+    def next(acc: Chunk[Byte]): Chunk[Byte] = {
+      val size    = deflater.deflate(buffer, 0, buffer.length, flushMode.jValue)
+      val current = Chunk.fromArray(ju.Arrays.copyOf(buffer, size))
+      if (current.isEmpty) acc
+      else next(acc ++ current)
+    }
+
+    next(Chunk.empty)
+  }
+
+}
diff --git a/streams/jvm/src/main/scala/zio/stream/compression/Gzipper.scala b/streams/jvm/src/main/scala/zio/stream/compression/Gzipper.scala
@@ -0,0 +1,71 @@
+package zio.stream.compression
+
+import java.util.zip.{ CRC32, Deflater }
+
+import zio.stream.compression.Gzipper._
+import zio.{ Chunk, ZIO }
+
+private[compression] class Gzipper(
+  bufferSize: Int,
+  level: CompressionLevel,
+  strategy: CompressionStrategy,
+  flushMode: FlushMode
+) {
+  private val crc             = new CRC32()
+  private val buffer          = new Array[Byte](bufferSize)
+  private var headerSent      = false
+  private var inputSize: Long = 0
+  private val deflater: Deflater = {
+    val deflater = new Deflater(level.jValue, true)
+    deflater.setStrategy(strategy.jValue)
+    deflater
+  }
+
+  def onNone: ZIO[Any, Nothing, Chunk[Byte]] = ZIO.effectTotal {
+    deflater.finish()
+    val restAndTrailer = Deflate.pullOutput(deflater, buffer, flushMode) ++ getTrailer
+    val lastChunk      = if (headerSent) restAndTrailer else header ++ restAndTrailer
+    deflater.reset()
+    crc.reset()
+    inputSize = 0
+    headerSent = false
+    lastChunk
+  }
+
+  def onChunk(chunk: Chunk[Byte]): ZIO[Any, Nothing, Chunk[Byte]] = ZIO.effectTotal {
+    val input = chunk.toArray
+    inputSize += input.length
+    crc.update(input)
+    deflater.setInput(input)
+    val deflated = Deflate.pullOutput(deflater, buffer, flushMode)
+    if (headerSent) deflated
+    else {
+      headerSent = true
+      header ++ deflated
+    }
+  }
+
+  def getTrailer: Chunk[Byte] = {
+    def byte(v: Long, n: Int) = ((v >> n * 8) & 0xff).toByte
+
+    val v = crc.getValue
+    val s = inputSize & 0xffff
+    Chunk(byte(v, 0), byte(v, 1), byte(v, 2), byte(v, 3), byte(s, 0), byte(s, 1), byte(s, 2), byte(s, 3))
+  }
+
+  def close(): Unit =
+    deflater.finish()
+}
+
+private[stream] object Gzipper {
+
+  // No XFL regardless maximum of the fastest compression used, no MTIME timestamp, unknown OS.
+  private val header: Chunk[Byte] = Chunk(31, 139, 8, 0, 0, 0, 0, 0, 0, 255).map(_.toByte)
+
+  def make(
+    bufferSize: Int,
+    level: CompressionLevel,
+    strategy: CompressionStrategy,
+    flushMode: FlushMode
+  ): ZIO[Any, Nothing, Gzipper] = ZIO.succeed(new Gzipper(bufferSize, level, strategy, flushMode))
+}