{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE Rank2Types #-}

module Codec.Compression.Zlib.Monad (
  DeflateM,
  runDeflateM,
  ZlibDecoder (..),
  raise,
  DecompressionError (..),

  -- * Getting data from the input stream.
  nextBits,
  nextByte,
  nextWord16,
  nextWord32,
  nextBlock,
  nextCode,

  -- * Aligning
  advanceToByte,

  -- * Emitting data into the output window
  emitByte,
  emitBlock,
  emitPastChunk,

  -- * Getting and publishing output
  finalAdler,
  moveWindow,
  finalize,
) where

import Codec.Compression.Zlib.Adler32 (
  AdlerState,
  advanceAdler,
  advanceAdlerBlock,
  finalizeAdler,
  initialAdlerState,
 )
import Codec.Compression.Zlib.HuffmanTree (
  AdvanceResult (..),
  HuffmanTree,
  advanceTree,
 )
import Codec.Compression.Zlib.OutputWindow (
  OutputWindow,
  addByte,
  addChunk,
  addOldChunk,
  emitExcess,
  emptyWindow,
  finalizeWindow,
 )
import Control.Exception (Exception)
import Data.Bits (Bits (..))
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Data.Int (Int64)
import Data.Typeable (Typeable)
import Data.Word (Word16, Word32, Word8)
import GHC.ST (ST)
import Prelude.Compat
import Prelude ()

data DecompressionState s = DecompressionState
  { DecompressionState s -> Int
dcsNextBitNo :: !Int
  , DecompressionState s -> Word8
dcsCurByte :: !Word8
  , DecompressionState s -> AdlerState
dcsAdler32 :: !AdlerState
  , DecompressionState s -> ByteString
dcsInput :: !S.ByteString
  , DecompressionState s -> OutputWindow s
dcsOutput :: !(OutputWindow s)
  }

instance Show (DecompressionState s) where
  show :: DecompressionState s -> String
show DecompressionState s
dcs =
    String
"DecompressionState<nextBit=" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (DecompressionState s -> Int
forall s. DecompressionState s -> Int
dcsNextBitNo DecompressionState s
dcs) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
","
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"curByte="
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show (DecompressionState s -> Word8
forall s. DecompressionState s -> Word8
dcsCurByte DecompressionState s
dcs)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
",inputLen="
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (ByteString -> Int
S.length (DecompressionState s -> ByteString
forall s. DecompressionState s -> ByteString
dcsInput DecompressionState s
dcs))
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
">"

-- -----------------------------------------------------------------------------

data DecompressionError
  = HuffmanTreeError String
  | FormatError String
  | DecompressionError String
  | HeaderError String
  | ChecksumError String
  deriving (Typeable, DecompressionError -> DecompressionError -> Bool
(DecompressionError -> DecompressionError -> Bool)
-> (DecompressionError -> DecompressionError -> Bool)
-> Eq DecompressionError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DecompressionError -> DecompressionError -> Bool
$c/= :: DecompressionError -> DecompressionError -> Bool
== :: DecompressionError -> DecompressionError -> Bool
$c== :: DecompressionError -> DecompressionError -> Bool
Eq)

instance Show DecompressionError where
  show :: DecompressionError -> String
show DecompressionError
x =
    case DecompressionError
x of
      HuffmanTreeError String
s -> String
"Huffman tree manipulation error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s
      FormatError String
s -> String
"Block format error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s
      DecompressionError String
s -> String
"Decompression error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s
      HeaderError String
s -> String
"Header error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s
      ChecksumError String
s -> String
"Checksum error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s

instance Exception DecompressionError

-- -----------------------------------------------------------------------------

newtype DeflateM s a = DeflateM
  { DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM ::
      DecompressionState s ->
      (DecompressionState s -> a -> ST s (ZlibDecoder s)) ->
      ST s (ZlibDecoder s)
  }

instance Applicative (DeflateM s) where
  pure :: a -> DeflateM s a
pure a
x = (DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM (\DecompressionState s
s DecompressionState s -> a -> ST s (ZlibDecoder s)
k -> DecompressionState s -> a -> ST s (ZlibDecoder s)
k DecompressionState s
s a
x)

  DeflateM s (a -> b)
f <*> :: DeflateM s (a -> b) -> DeflateM s a -> DeflateM s b
<*> DeflateM s a
x = (DecompressionState s
 -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s b
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM ((DecompressionState s
  -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
  -> ST s (ZlibDecoder s))
 -> DeflateM s b)
-> (DecompressionState s
    -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
    -> ST s (ZlibDecoder s))
-> DeflateM s b
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s1 DecompressionState s -> b -> ST s (ZlibDecoder s)
k ->
    DeflateM s (a -> b)
-> DecompressionState s
-> (DecompressionState s -> (a -> b) -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s (a -> b)
f DecompressionState s
s1 ((DecompressionState s -> (a -> b) -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> (DecompressionState s -> (a -> b) -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s2 a -> b
g ->
      DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s a
x DecompressionState s
s2 ((DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s3 a
y -> DecompressionState s -> b -> ST s (ZlibDecoder s)
k DecompressionState s
s3 (a -> b
g a
y)

  DeflateM s a
m *> :: DeflateM s a -> DeflateM s b -> DeflateM s b
*> DeflateM s b
n = (DecompressionState s
 -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s b
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM ((DecompressionState s
  -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
  -> ST s (ZlibDecoder s))
 -> DeflateM s b)
-> (DecompressionState s
    -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
    -> ST s (ZlibDecoder s))
-> DeflateM s b
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s1 DecompressionState s -> b -> ST s (ZlibDecoder s)
k ->
    DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s a
m DecompressionState s
s1 ((DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s2 a
_ -> DeflateM s b
-> DecompressionState s
-> (DecompressionState s -> b -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s b
n DecompressionState s
s2 DecompressionState s -> b -> ST s (ZlibDecoder s)
k

  {-# INLINE pure #-}
  {-# INLINE (<*>) #-}
  {-# INLINE (*>) #-}

instance Functor (DeflateM s) where
  fmap :: (a -> b) -> DeflateM s a -> DeflateM s b
fmap a -> b
f DeflateM s a
m = (DecompressionState s
 -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s b
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM (\DecompressionState s
s DecompressionState s -> b -> ST s (ZlibDecoder s)
k -> DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s a
m DecompressionState s
s (\DecompressionState s
s' a
a -> DecompressionState s -> b -> ST s (ZlibDecoder s)
k DecompressionState s
s' (a -> b
f a
a)))
  {-# INLINE fmap #-}

instance Monad (DeflateM s) where
  {-# INLINE return #-}
  return :: a -> DeflateM s a
return = a -> DeflateM s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

  {-# INLINE (>>=) #-}
  DeflateM s a
m >>= :: DeflateM s a -> (a -> DeflateM s b) -> DeflateM s b
>>= a -> DeflateM s b
f = (DecompressionState s
 -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s b
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM ((DecompressionState s
  -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
  -> ST s (ZlibDecoder s))
 -> DeflateM s b)
-> (DecompressionState s
    -> (DecompressionState s -> b -> ST s (ZlibDecoder s))
    -> ST s (ZlibDecoder s))
-> DeflateM s b
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s1 DecompressionState s -> b -> ST s (ZlibDecoder s)
k ->
    DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s a
m DecompressionState s
s1 ((DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s2 a
a -> DeflateM s b
-> DecompressionState s
-> (DecompressionState s -> b -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM (a -> DeflateM s b
f a
a) DecompressionState s
s2 DecompressionState s -> b -> ST s (ZlibDecoder s)
k

  >> :: DeflateM s a -> DeflateM s b -> DeflateM s b
(>>) = DeflateM s a -> DeflateM s b -> DeflateM s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)
  {-# INLINE (>>) #-}

get :: DeflateM s (DecompressionState s)
get :: DeflateM s (DecompressionState s)
get = (DecompressionState s
 -> (DecompressionState s
     -> DecompressionState s -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s (DecompressionState s)
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM (\DecompressionState s
s DecompressionState s
-> DecompressionState s -> ST s (ZlibDecoder s)
k -> DecompressionState s
-> DecompressionState s -> ST s (ZlibDecoder s)
k DecompressionState s
s DecompressionState s
s)
{-# INLINE get #-}

set :: DecompressionState s -> DeflateM s ()
set :: DecompressionState s -> DeflateM s ()
set !DecompressionState s
s = (DecompressionState s
 -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s ()
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM (\DecompressionState s
_ DecompressionState s -> () -> ST s (ZlibDecoder s)
k -> DecompressionState s -> () -> ST s (ZlibDecoder s)
k DecompressionState s
s ())
{-# INLINE set #-}

raise :: DecompressionError -> DeflateM s a
raise :: DecompressionError -> DeflateM s a
raise DecompressionError
e = (DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM (\DecompressionState s
_ DecompressionState s -> a -> ST s (ZlibDecoder s)
_ -> ZlibDecoder s -> ST s (ZlibDecoder s)
forall (m :: * -> *) a. Monad m => a -> m a
return (DecompressionError -> ZlibDecoder s
forall s. DecompressionError -> ZlibDecoder s
DecompError DecompressionError
e))
{-# INLINE raise #-}

liftST :: ST s a -> DeflateM s a
liftST :: ST s a -> DeflateM s a
liftST ST s a
action = (DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM ((DecompressionState s
  -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
  -> ST s (ZlibDecoder s))
 -> DeflateM s a)
-> (DecompressionState s
    -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
    -> ST s (ZlibDecoder s))
-> DeflateM s a
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
s DecompressionState s -> a -> ST s (ZlibDecoder s)
k -> do
  a
res <- ST s a
action
  DecompressionState s -> a -> ST s (ZlibDecoder s)
k DecompressionState s
s a
res

-- -----------------------------------------------------------------------------

data ZlibDecoder s
  = NeedMore (S.ByteString -> ST s (ZlibDecoder s))
  | Chunk S.ByteString (ST s (ZlibDecoder s))
  | Done
  | DecompError DecompressionError

runDeflateM :: DeflateM s () -> ST s (ZlibDecoder s)
runDeflateM :: DeflateM s () -> ST s (ZlibDecoder s)
runDeflateM DeflateM s ()
m = do
  OutputWindow s
window <- ST s (OutputWindow s)
forall s. ST s (OutputWindow s)
emptyWindow
  let initialState :: DecompressionState s
initialState =
        DecompressionState :: forall s.
Int
-> Word8
-> AdlerState
-> ByteString
-> OutputWindow s
-> DecompressionState s
DecompressionState
          { dcsNextBitNo :: Int
dcsNextBitNo = Int
8
          , dcsCurByte :: Word8
dcsCurByte = Word8
0
          , dcsAdler32 :: AdlerState
dcsAdler32 = AdlerState
initialAdlerState
          , dcsInput :: ByteString
dcsInput = ByteString
S.empty
          , dcsOutput :: OutputWindow s
dcsOutput = OutputWindow s
window
          }
  DeflateM s ()
-> DecompressionState s
-> (DecompressionState s -> () -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
forall s a.
DeflateM s a
-> DecompressionState s
-> (DecompressionState s -> a -> ST s (ZlibDecoder s))
-> ST s (ZlibDecoder s)
unDeflateM DeflateM s ()
m DecompressionState s
initialState (\DecompressionState s
_ ()
_ -> ZlibDecoder s -> ST s (ZlibDecoder s)
forall (m :: * -> *) a. Monad m => a -> m a
return ZlibDecoder s
forall s. ZlibDecoder s
Done)
{-# INLINE runDeflateM #-}

-- -----------------------------------------------------------------------------

getNextChunk :: DeflateM s ()
getNextChunk :: DeflateM s ()
getNextChunk = (DecompressionState s
 -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s ()
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM ((DecompressionState s
  -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
  -> ST s (ZlibDecoder s))
 -> DeflateM s ())
-> (DecompressionState s
    -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
    -> ST s (ZlibDecoder s))
-> DeflateM s ()
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
st DecompressionState s -> () -> ST s (ZlibDecoder s)
k -> ZlibDecoder s -> ST s (ZlibDecoder s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((ByteString -> ST s (ZlibDecoder s)) -> ZlibDecoder s
forall s. (ByteString -> ST s (ZlibDecoder s)) -> ZlibDecoder s
NeedMore (DecompressionState s
-> (DecompressionState s -> () -> ST s (ZlibDecoder s))
-> ByteString
-> ST s (ZlibDecoder s)
forall s.
DecompressionState s
-> (DecompressionState s -> () -> ST s (ZlibDecoder s))
-> ByteString
-> ST s (ZlibDecoder s)
loadChunk DecompressionState s
st DecompressionState s -> () -> ST s (ZlibDecoder s)
k))
 where
  loadChunk ::
    DecompressionState s ->
    (DecompressionState s -> () -> ST s (ZlibDecoder s)) ->
    S.ByteString ->
    ST s (ZlibDecoder s)
  loadChunk :: DecompressionState s
-> (DecompressionState s -> () -> ST s (ZlibDecoder s))
-> ByteString
-> ST s (ZlibDecoder s)
loadChunk DecompressionState s
st DecompressionState s -> () -> ST s (ZlibDecoder s)
k ByteString
bstr =
    case ByteString -> Maybe (Word8, ByteString)
S.uncons ByteString
bstr of
      Maybe (Word8, ByteString)
Nothing -> ZlibDecoder s -> ST s (ZlibDecoder s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((ByteString -> ST s (ZlibDecoder s)) -> ZlibDecoder s
forall s. (ByteString -> ST s (ZlibDecoder s)) -> ZlibDecoder s
NeedMore (DecompressionState s
-> (DecompressionState s -> () -> ST s (ZlibDecoder s))
-> ByteString
-> ST s (ZlibDecoder s)
forall s.
DecompressionState s
-> (DecompressionState s -> () -> ST s (ZlibDecoder s))
-> ByteString
-> ST s (ZlibDecoder s)
loadChunk DecompressionState s
st DecompressionState s -> () -> ST s (ZlibDecoder s)
k))
      Just (Word8
nextb, ByteString
rest) ->
        DecompressionState s -> () -> ST s (ZlibDecoder s)
k DecompressionState s
st{dcsNextBitNo :: Int
dcsNextBitNo = Int
0, dcsCurByte :: Word8
dcsCurByte = Word8
nextb, dcsInput :: ByteString
dcsInput = ByteString
rest} ()

{-# SPECIALIZE nextBits :: Int -> DeflateM s Word8 #-}
{-# SPECIALIZE nextBits :: Int -> DeflateM s Int #-}
{-# SPECIALIZE nextBits :: Int -> DeflateM s Int64 #-}
{-# INLINE nextBits #-}
nextBits :: (Num a, Bits a) => Int -> DeflateM s a
nextBits :: Int -> DeflateM s a
nextBits Int
x = Int -> Int -> a -> DeflateM s a
forall a s. (Num a, Bits a) => Int -> Int -> a -> DeflateM s a
nextBits' Int
x Int
0 a
0

{-# SPECIALIZE nextBits' :: Int -> Int -> Word8 -> DeflateM s Word8 #-}
{-# SPECIALIZE nextBits' :: Int -> Int -> Int -> DeflateM s Int #-}
{-# SPECIALIZE nextBits' :: Int -> Int -> Int64 -> DeflateM s Int64 #-}
{-# INLINE nextBits' #-}
nextBits' :: (Num a, Bits a) => Int -> Int -> a -> DeflateM s a
nextBits' :: Int -> Int -> a -> DeflateM s a
nextBits' !Int
x' !Int
shiftNum !a
acc
  | Int
x' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = a -> DeflateM s a
forall (m :: * -> *) a. Monad m => a -> m a
return a
acc
  | Bool
otherwise = do
    DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
    case DecompressionState s -> Int
forall s. DecompressionState s -> Int
dcsNextBitNo DecompressionState s
dcs of
      Int
8 -> case ByteString -> Maybe (Word8, ByteString)
S.uncons (DecompressionState s -> ByteString
forall s. DecompressionState s -> ByteString
dcsInput DecompressionState s
dcs) of
        Maybe (Word8, ByteString)
Nothing -> do
          DeflateM s ()
forall s. DeflateM s ()
getNextChunk
          Int -> Int -> a -> DeflateM s a
forall a s. (Num a, Bits a) => Int -> Int -> a -> DeflateM s a
nextBits' Int
x' Int
shiftNum a
acc
        Just (Word8
nextb, ByteString
rest) -> do
          DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsNextBitNo :: Int
dcsNextBitNo = Int
0, dcsCurByte :: Word8
dcsCurByte = Word8
nextb, dcsInput :: ByteString
dcsInput = ByteString
rest}
          Int -> Int -> a -> DeflateM s a
forall a s. (Num a, Bits a) => Int -> Int -> a -> DeflateM s a
nextBits' Int
x' Int
shiftNum a
acc
      Int
nextBitNo -> do
        let !myBits :: Int
myBits = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
x' (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nextBitNo)
            !base :: Word8
base = DecompressionState s -> Word8
forall s. DecompressionState s -> Word8
dcsCurByte DecompressionState s
dcs Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
nextBitNo
            !mask :: Word8
mask = Word8 -> Word8
forall a. Bits a => a -> a
complement (Word8
0xFF Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
myBits)
            !res :: a
res = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
base Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
mask)
            !acc' :: a
acc' = a
acc a -> a -> a
forall a. Bits a => a -> a -> a
.|. (a
res a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
shiftNum)
        DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsNextBitNo :: Int
dcsNextBitNo = Int
nextBitNo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
myBits}
        Int -> Int -> a -> DeflateM s a
forall a s. (Num a, Bits a) => Int -> Int -> a -> DeflateM s a
nextBits' (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
myBits) (Int
shiftNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
myBits) a
acc'

nextByte :: DeflateM s Word8
nextByte :: DeflateM s Word8
nextByte = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  if
      | DecompressionState s -> Int
forall s. DecompressionState s -> Int
dcsNextBitNo DecompressionState s
dcs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> do
        DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsNextBitNo :: Int
dcsNextBitNo = Int
8}
        Word8 -> DeflateM s Word8
forall (m :: * -> *) a. Monad m => a -> m a
return (DecompressionState s -> Word8
forall s. DecompressionState s -> Word8
dcsCurByte DecompressionState s
dcs)
      | DecompressionState s -> Int
forall s. DecompressionState s -> Int
dcsNextBitNo DecompressionState s
dcs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
8 -> Int -> DeflateM s Word8
forall a s. (Num a, Bits a) => Int -> DeflateM s a
nextBits Int
8 -- we're not aligned. sigh.
      | Bool
otherwise -> case ByteString -> Maybe (Word8, ByteString)
S.uncons (DecompressionState s -> ByteString
forall s. DecompressionState s -> ByteString
dcsInput DecompressionState s
dcs) of
        Maybe (Word8, ByteString)
Nothing -> DeflateM s ()
forall s. DeflateM s ()
getNextChunk DeflateM s () -> DeflateM s Word8 -> DeflateM s Word8
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> DeflateM s Word8
forall s. DeflateM s Word8
nextByte
        Just (Word8
nextb, ByteString
rest) -> do
          DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set
            DecompressionState s
dcs
              { dcsNextBitNo :: Int
dcsNextBitNo = Int
8
              , dcsCurByte :: Word8
dcsCurByte = Word8
nextb
              , dcsInput :: ByteString
dcsInput = ByteString
rest
              }
          Word8 -> DeflateM s Word8
forall (m :: * -> *) a. Monad m => a -> m a
return Word8
nextb

nextWord16 :: DeflateM s Word16
nextWord16 :: DeflateM s Word16
nextWord16 = do
  Word16
low <- Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word16) -> DeflateM s Word8 -> DeflateM s Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word16
high <- Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word16) -> DeflateM s Word8 -> DeflateM s Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word16 -> DeflateM s Word16
forall (m :: * -> *) a. Monad m => a -> m a
return ((Word16
high Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftL` Int
8) Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16
low)

nextWord32 :: DeflateM s Word32
nextWord32 :: DeflateM s Word32
nextWord32 = do
  Word32
a <- Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word32) -> DeflateM s Word8 -> DeflateM s Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word32
b <- Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word32) -> DeflateM s Word8 -> DeflateM s Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word32
c <- Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word32) -> DeflateM s Word8 -> DeflateM s Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word32
d <- Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word32) -> DeflateM s Word8 -> DeflateM s Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` DeflateM s Word8
forall s. DeflateM s Word8
nextByte
  Word32 -> DeflateM s Word32
forall (m :: * -> *) a. Monad m => a -> m a
return ((Word32
a Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
24) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word32
b Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
16) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word32
c Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
8) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32
d)

nextBlock :: Integral a => a -> DeflateM s L.ByteString
nextBlock :: a -> DeflateM s ByteString
nextBlock a
amt = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  if
      | DecompressionState s -> Int
forall s. DecompressionState s -> Int
dcsNextBitNo DecompressionState s
dcs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> do
        let startByte :: Word8
startByte = DecompressionState s -> Word8
forall s. DecompressionState s -> Word8
dcsCurByte DecompressionState s
dcs
        DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsNextBitNo :: Int
dcsNextBitNo = Int
8}
        ByteString
rest <- a -> DeflateM s ByteString
forall a s. Integral a => a -> DeflateM s ByteString
nextBlock (a
amt a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
        ByteString -> DeflateM s ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (Word8 -> ByteString -> ByteString
L.cons Word8
startByte ByteString
rest)
      | DecompressionState s -> Int
forall s. DecompressionState s -> Int
dcsNextBitNo DecompressionState s
dcs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
8 ->
        Int -> ByteString -> DeflateM s ByteString
forall s. Int -> ByteString -> DeflateM s ByteString
getBlock (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
amt) (DecompressionState s -> ByteString
forall s. DecompressionState s -> ByteString
dcsInput DecompressionState s
dcs)
      | Bool
otherwise ->
        DecompressionError -> DeflateM s ByteString
forall s a. DecompressionError -> DeflateM s a
raise (String -> DecompressionError
FormatError String
"Can't get a block on a non-byte boundary.")
 where
  getBlock :: Int -> ByteString -> DeflateM s ByteString
getBlock Int
len ByteString
bstr
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ByteString -> Int
S.length ByteString
bstr = do
      let (ByteString
mine, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt Int
len ByteString
bstr
      DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
      DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsNextBitNo :: Int
dcsNextBitNo = Int
8, dcsInput :: ByteString
dcsInput = ByteString
rest}
      ByteString -> DeflateM s ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> ByteString
L.fromStrict ByteString
mine)
    | ByteString -> Bool
S.null ByteString
bstr = do
      DeflateM s ()
forall s. DeflateM s ()
getNextChunk
      DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
      let byte1 :: Word8
byte1 = DecompressionState s -> Word8
forall s. DecompressionState s -> Word8
dcsCurByte DecompressionState s
dcs
      ByteString
rest <- Int -> ByteString -> DeflateM s ByteString
getBlock (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (DecompressionState s -> ByteString
forall s. DecompressionState s -> ByteString
dcsInput DecompressionState s
dcs)
      ByteString -> DeflateM s ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (Word8 -> ByteString -> ByteString
L.cons Word8
byte1 ByteString
rest)
    | Bool
otherwise = do
      ByteString
rest <- Int -> ByteString -> DeflateM s ByteString
getBlock (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
S.length ByteString
bstr) ByteString
S.empty
      ByteString -> DeflateM s ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> ByteString
L.fromStrict ByteString
bstr ByteString -> ByteString -> ByteString
`L.append` ByteString
rest)

nextCode :: Show a => HuffmanTree a -> DeflateM s a
nextCode :: HuffmanTree a -> DeflateM s a
nextCode HuffmanTree a
tree = do
  Word8
b <- Int -> DeflateM s Word8
forall a s. (Num a, Bits a) => Int -> DeflateM s a
nextBits Int
1
  case Word8 -> HuffmanTree a -> AdvanceResult a
forall a. Word8 -> HuffmanTree a -> AdvanceResult a
advanceTree Word8
b HuffmanTree a
tree of
    AdvanceError String
str -> DecompressionError -> DeflateM s a
forall s a. DecompressionError -> DeflateM s a
raise (String -> DecompressionError
HuffmanTreeError String
str)
    NewTree HuffmanTree a
tree' -> HuffmanTree a -> DeflateM s a
forall a s. Show a => HuffmanTree a -> DeflateM s a
nextCode HuffmanTree a
tree'
    Result a
x -> a -> DeflateM s a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
{-# INLINE nextCode #-}

advanceToByte :: DeflateM s ()
advanceToByte :: DeflateM s ()
advanceToByte = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsNextBitNo :: Int
dcsNextBitNo = Int
8}

emitByte :: Word8 -> DeflateM s ()
emitByte :: Word8 -> DeflateM s ()
emitByte Word8
b = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  OutputWindow s
output' <- ST s (OutputWindow s) -> DeflateM s (OutputWindow s)
forall s a. ST s a -> DeflateM s a
liftST (OutputWindow s -> Word8 -> ST s (OutputWindow s)
forall s. OutputWindow s -> Word8 -> ST s (OutputWindow s)
addByte (DecompressionState s -> OutputWindow s
forall s. DecompressionState s -> OutputWindow s
dcsOutput DecompressionState s
dcs) Word8
b)
  let adler' :: AdlerState
adler' = AdlerState -> Word8 -> AdlerState
advanceAdler (DecompressionState s -> AdlerState
forall s. DecompressionState s -> AdlerState
dcsAdler32 DecompressionState s
dcs) Word8
b
  DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsOutput :: OutputWindow s
dcsOutput = OutputWindow s
output', dcsAdler32 :: AdlerState
dcsAdler32 = AdlerState
adler'}
{-# INLINE emitByte #-}

emitBlock :: L.ByteString -> DeflateM s ()
emitBlock :: ByteString -> DeflateM s ()
emitBlock ByteString
b = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  OutputWindow s
output' <- ST s (OutputWindow s) -> DeflateM s (OutputWindow s)
forall s a. ST s a -> DeflateM s a
liftST (OutputWindow s -> ByteString -> ST s (OutputWindow s)
forall s. OutputWindow s -> ByteString -> ST s (OutputWindow s)
addChunk (DecompressionState s -> OutputWindow s
forall s. DecompressionState s -> OutputWindow s
dcsOutput DecompressionState s
dcs) ByteString
b)
  let adler' :: AdlerState
adler' = (AdlerState -> ByteString -> AdlerState)
-> AdlerState -> ByteString -> AdlerState
forall a. (a -> ByteString -> a) -> a -> ByteString -> a
L.foldlChunks AdlerState -> ByteString -> AdlerState
advanceAdlerBlock (DecompressionState s -> AdlerState
forall s. DecompressionState s -> AdlerState
dcsAdler32 DecompressionState s
dcs) ByteString
b
  DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsOutput :: OutputWindow s
dcsOutput = OutputWindow s
output', dcsAdler32 :: AdlerState
dcsAdler32 = AdlerState
adler'}

emitPastChunk :: Int -> Int -> DeflateM s ()
emitPastChunk :: Int -> Int -> DeflateM s ()
emitPastChunk Int
dist Int
len = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  (OutputWindow s
output', ByteString
newChunk) <- ST s (OutputWindow s, ByteString)
-> DeflateM s (OutputWindow s, ByteString)
forall s a. ST s a -> DeflateM s a
liftST (OutputWindow s -> Int -> Int -> ST s (OutputWindow s, ByteString)
forall s.
OutputWindow s -> Int -> Int -> ST s (OutputWindow s, ByteString)
addOldChunk (DecompressionState s -> OutputWindow s
forall s. DecompressionState s -> OutputWindow s
dcsOutput DecompressionState s
dcs) Int
dist Int
len)
  DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set
    DecompressionState s
dcs
      { dcsOutput :: OutputWindow s
dcsOutput = OutputWindow s
output'
      , dcsAdler32 :: AdlerState
dcsAdler32 = AdlerState -> ByteString -> AdlerState
advanceAdlerBlock (DecompressionState s -> AdlerState
forall s. DecompressionState s -> AdlerState
dcsAdler32 DecompressionState s
dcs) ByteString
newChunk
      }
{-# INLINE emitPastChunk #-}

finalAdler :: DeflateM s Word32
finalAdler :: DeflateM s Word32
finalAdler = (AdlerState -> Word32
finalizeAdler (AdlerState -> Word32)
-> (DecompressionState s -> AdlerState)
-> DecompressionState s
-> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DecompressionState s -> AdlerState
forall s. DecompressionState s -> AdlerState
dcsAdler32) (DecompressionState s -> Word32)
-> DeflateM s (DecompressionState s) -> DeflateM s Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get

moveWindow :: DeflateM s ()
moveWindow :: DeflateM s ()
moveWindow = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  Maybe (ByteString, OutputWindow s)
possibleExcess <- ST s (Maybe (ByteString, OutputWindow s))
-> DeflateM s (Maybe (ByteString, OutputWindow s))
forall s a. ST s a -> DeflateM s a
liftST (OutputWindow s -> ST s (Maybe (ByteString, OutputWindow s))
forall s.
OutputWindow s -> ST s (Maybe (ByteString, OutputWindow s))
emitExcess (DecompressionState s -> OutputWindow s
forall s. DecompressionState s -> OutputWindow s
dcsOutput DecompressionState s
dcs))
  case Maybe (ByteString, OutputWindow s)
possibleExcess of
    Maybe (ByteString, OutputWindow s)
Nothing ->
      () -> DeflateM s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just (ByteString
builtChunk, OutputWindow s
output') -> do
      DecompressionState s -> DeflateM s ()
forall s. DecompressionState s -> DeflateM s ()
set DecompressionState s
dcs{dcsOutput :: OutputWindow s
dcsOutput = OutputWindow s
output'}
      ByteString -> DeflateM s ()
forall s. ByteString -> DeflateM s ()
publish ByteString
builtChunk

finalize :: DeflateM s ()
finalize :: DeflateM s ()
finalize = do
  DecompressionState s
dcs <- DeflateM s (DecompressionState s)
forall s. DeflateM s (DecompressionState s)
get
  ByteString
lastChunk <- ST s ByteString -> DeflateM s ByteString
forall s a. ST s a -> DeflateM s a
liftST (OutputWindow s -> ST s ByteString
forall s. OutputWindow s -> ST s ByteString
finalizeWindow (DecompressionState s -> OutputWindow s
forall s. DecompressionState s -> OutputWindow s
dcsOutput DecompressionState s
dcs))
  ByteString -> DeflateM s ()
forall s. ByteString -> DeflateM s ()
publish ByteString
lastChunk

{-# INLINE publish #-}
publish :: S.ByteString -> DeflateM s ()
publish :: ByteString -> DeflateM s ()
publish ByteString
bstr = (DecompressionState s
 -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s ()
forall s a.
(DecompressionState s
 -> (DecompressionState s -> a -> ST s (ZlibDecoder s))
 -> ST s (ZlibDecoder s))
-> DeflateM s a
DeflateM ((DecompressionState s
  -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
  -> ST s (ZlibDecoder s))
 -> DeflateM s ())
-> (DecompressionState s
    -> (DecompressionState s -> () -> ST s (ZlibDecoder s))
    -> ST s (ZlibDecoder s))
-> DeflateM s ()
forall a b. (a -> b) -> a -> b
$ \DecompressionState s
st DecompressionState s -> () -> ST s (ZlibDecoder s)
k ->
  ZlibDecoder s -> ST s (ZlibDecoder s)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> ST s (ZlibDecoder s) -> ZlibDecoder s
forall s. ByteString -> ST s (ZlibDecoder s) -> ZlibDecoder s
Chunk ByteString
bstr (DecompressionState s -> () -> ST s (ZlibDecoder s)
k DecompressionState s
st ()))