Zlib.hs

module Hypermedia.Datastar.Compression.Zlib
  ( gzip
  , gzipWith
  , deflate
  , deflateWith
  , defaultCompressionLevel
  )
where

import Control.Exception (throwIO)

import Data.ByteString qualified as BS
import Data.ByteString.Builder qualified as BSB
import Data.ByteString.Lazy qualified as BL

import Data.Streaming.Zlib
  ( PopperRes (..)
  , WindowBits (..)
  , feedDeflate
  , finishDeflate
  , flushDeflate
  , initDeflate
  )

import Hypermedia.Datastar.WAI (Compressor (..))

defaultCompressionLevel :: Int
defaultCompressionLevel = 6

-- | A gzip 'Compressor' (@Content-Encoding: gzip@) at 'defaultCompressionLevel'.
gzip :: Compressor
gzip = gzipWith defaultCompressionLevel

-- | A gzip 'Compressor' at an explicit zlib level (0–9).
gzipWith :: Int -> Compressor
gzipWith level = zlibCompressor "gzip" level (WindowBits 31)

-- | A zlib/deflate 'Compressor' (@Content-Encoding: deflate@) at 'defaultCompressionLevel'.
deflate :: Compressor
deflate = deflateWith defaultCompressionLevel

-- | A zlib/deflate 'Compressor' at an explicit zlib level (0–9).
deflateWith :: Int -> Compressor
deflateWith level = zlibCompressor "deflate" level (WindowBits 15)

-- WindowBits 31 selects gzip framing, 15 selects zlib/deflate framing.
zlibCompressor :: BS.ByteString -> Int -> WindowBits -> Compressor
zlibCompressor enc level wbits = Compressor enc wrap
  where 
    wrap rawWrite rawFlush = do
        def <- initDeflate level wbits

        let drain popper = do
              res <- popper
              case res of
                PRDone -> pure ()
                PRNext bs -> rawWrite (BSB.byteString bs) >> drain popper
                PRError e -> throwIO e

            write builder = do
              popper <- feedDeflate def $ BL.toStrict $ BSB.toLazyByteString builder
              drain popper

            flush = drain (flushDeflate def) >> rawFlush

            finish = drain (finishDeflate def) >> rawFlush

        pure (write, flush, finish)