{-# LANGUAGE MultiWayIf #-}

module Codec.Compression.Zlib (
  DecompressionError (..),
  ZlibDecoder (NeedMore, Chunk, Done, DecompError),
  decompress,
  decompressIncremental,
) where

import Codec.Compression.Zlib.Deflate (inflate)
import Codec.Compression.Zlib.Monad (
  DecompressionError (..),
  DeflateM,
  ZlibDecoder (..),
  nextByte,
  raise,
  runDeflateM,
 )
import Control.Monad (replicateM_, unless, when)
import Data.Bits (shiftL, shiftR, testBit, (.&.), (.|.))
import qualified Data.ByteString as S
import Data.ByteString.Builder (Builder, byteString, toLazyByteString)
import qualified Data.ByteString.Lazy as L
import Data.Word (Word16)
import GHC.ST (ST, runST)
import Prelude.Compat
import Prelude ()

decompressIncremental :: ST s (ZlibDecoder s)
decompressIncremental :: ST s (ZlibDecoder s)
decompressIncremental = DeflateM s () -> ST s (ZlibDecoder s)
forall s. DeflateM s () -> ST s (ZlibDecoder s)
runDeflateM DeflateM s ()
forall s. DeflateM s ()
inflateWithHeaders

decompress :: L.ByteString -> Either DecompressionError L.ByteString
decompress :: ByteString -> Either DecompressionError ByteString
decompress ByteString
ifile = (forall s. ST s (Either DecompressionError ByteString))
-> Either DecompressionError ByteString
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Either DecompressionError ByteString))
 -> Either DecompressionError ByteString)
-> (forall s. ST s (Either DecompressionError ByteString))
-> Either DecompressionError ByteString
forall a b. (a -> b) -> a -> b
$ do
  ZlibDecoder s
base <- ST s (ZlibDecoder s)
forall s. ST s (ZlibDecoder s)
decompressIncremental
  ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
forall s.
ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
run ZlibDecoder s
base (ByteString -> [ByteString]
L.toChunks ByteString
ifile) Builder
forall a. Monoid a => a
mempty
 where
  run :: ZlibDecoder s -> [S.ByteString] -> Builder -> ST s (Either DecompressionError L.ByteString)
  run :: ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
run (NeedMore ByteString -> ST s (ZlibDecoder s)
_) [] Builder
_ =
    Either DecompressionError ByteString
-> ST s (Either DecompressionError ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (DecompressionError -> Either DecompressionError ByteString
forall a b. a -> Either a b
Left (String -> DecompressionError
DecompressionError String
"Ran out of data mid-decompression 2."))
  run (NeedMore ByteString -> ST s (ZlibDecoder s)
f) (ByteString
first : [ByteString]
rest) Builder
acc = do
    ZlibDecoder s
nextState <- ByteString -> ST s (ZlibDecoder s)
f ByteString
first
    ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
forall s.
ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
run ZlibDecoder s
nextState [ByteString]
rest Builder
acc
  run (Chunk ByteString
c ST s (ZlibDecoder s)
m) [ByteString]
ls Builder
acc = do
    ZlibDecoder s
nextState <- ST s (ZlibDecoder s)
m
    ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
forall s.
ZlibDecoder s
-> [ByteString]
-> Builder
-> ST s (Either DecompressionError ByteString)
run ZlibDecoder s
nextState [ByteString]
ls (Builder
acc Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
byteString ByteString
c)
  run ZlibDecoder s
Done [] Builder
acc =
    Either DecompressionError ByteString
-> ST s (Either DecompressionError ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> Either DecompressionError ByteString
forall a b. b -> Either a b
Right (Builder -> ByteString
toLazyByteString Builder
acc))
  run ZlibDecoder s
Done (ByteString
_ : [ByteString]
_) Builder
_ =
    Either DecompressionError ByteString
-> ST s (Either DecompressionError ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (DecompressionError -> Either DecompressionError ByteString
forall a b. a -> Either a b
Left (String -> DecompressionError
DecompressionError String
"Finished with data remaining."))
  run (DecompError DecompressionError
e) [ByteString]
_ Builder
_ =
    Either DecompressionError ByteString
-> ST s (Either DecompressionError ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (DecompressionError -> Either DecompressionError ByteString
forall a b. a -> Either a b
Left DecompressionError
e)

inflateWithHeaders :: DeflateM s ()
inflateWithHeaders :: DeflateM s ()
inflateWithHeaders = do
  Word8
cmf <- DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word8
flg <- DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  let both :: Word16
both = Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
cmf Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftL` Int
8 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
flg
      cm :: Word8
cm = Word8
cmf Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0f
      cinfo :: Word8
cinfo = Word8
cmf Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
4
      fdict :: Bool
fdict = Word8 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word8
flg Int
5
  --       flevel = flg `shiftR` 6
  Bool -> DeflateM s () -> DeflateM s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ((Word16
both :: Word16) Word16 -> Word16 -> Word16
forall a. Integral a => a -> a -> a
`mod` Word16
31 Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
0) (DeflateM s () -> DeflateM s ()) -> DeflateM s () -> DeflateM s ()
forall a b. (a -> b) -> a -> b
$
    DecompressionError -> DeflateM s ()
forall s a. DecompressionError -> DeflateM s a
raise (String -> DecompressionError
HeaderError String
"Header checksum failed")
  Bool -> DeflateM s () -> DeflateM s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word8
cm Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
8) (DeflateM s () -> DeflateM s ()) -> DeflateM s () -> DeflateM s ()
forall a b. (a -> b) -> a -> b
$
    DecompressionError -> DeflateM s ()
forall s a. DecompressionError -> DeflateM s a
raise (String -> DecompressionError
HeaderError (String
"Bad compression method: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
cm))
  Bool -> DeflateM s () -> DeflateM s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word8
cinfo Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
7) (DeflateM s () -> DeflateM s ()) -> DeflateM s () -> DeflateM s ()
forall a b. (a -> b) -> a -> b
$
    DecompressionError -> DeflateM s ()
forall s a. DecompressionError -> DeflateM s a
raise (String -> DecompressionError
HeaderError (String
"Window size too big: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
cinfo))
  Bool -> DeflateM s () -> DeflateM s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
fdict (DeflateM s () -> DeflateM s ()) -> DeflateM s () -> DeflateM s ()
forall a b. (a -> b) -> a -> b
$ Int -> DeflateM s Word8 -> DeflateM s ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
4 DeflateM s Word8
forall s. DeflateM s Word8
nextByte -- just skip them for now (FIXME)
  DeflateM s ()
forall s. DeflateM s ()
inflate