    {-# LANGUAGE RecordWildCards #-}

    module Network.HPACK.Huffman.Encode (
        -- * Huffman encoding
        encodeH,
        encodeHuffman,
    ) where

    import Control.Exception (throwIO)
    import Data.Array.Base (unsafeAt)
    import Data.Array.IArray (listArray)
    import Data.Array.Unboxed (UArray)
    import Data.IORef
    import Foreign.Ptr (minusPtr, plusPtr)
    import Foreign.Storable (poke)
    import Network.ByteOrder hiding (copy)

    import Imports
    import Network.HPACK.Huffman.Params (idxEos)
    import Network.HPACK.Huffman.Table

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

    huffmanLength :: UArray Int Int
    huffmanLength :: UArray Int Int
    huffmanLength = (Int, Int) -> [Int] -> UArray Int Int
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    (i, i) -> [e] -> a i e
    listArray (Int
    0, Int
    idxEos) ([Int] -> UArray Int Int) -> [Int] -> UArray Int Int
    forall a b. (a -> b) -> a -> b
    $ ([B] -> Int) -> [[B]] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map [B] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [[B]]
    huffmanTable

    huffmanCode :: UArray Int Word64
    huffmanCode :: UArray Int Word64
    huffmanCode = (Int, Int) -> [Word64] -> UArray Int Word64
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    (i, i) -> [e] -> a i e
    listArray (Int
    0, Int
    idxEos) [Word64]
    huffmanTable'

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

    -- | Huffman encoding.
    encodeH
        :: WriteBuffer
        -> ByteString
        -- ^ Target
        -> IO Int
        -- ^ The length of the encoded string.
    encodeH :: WriteBuffer -> ByteString -> IO Int
    encodeH WriteBuffer
    dst ByteString
    bs = ByteString -> (ReadBuffer -> IO Int) -> IO Int
    forall a. ByteString -> (ReadBuffer -> IO a) -> IO a
    withReadBuffer ByteString
    bs ((ReadBuffer -> IO Int) -> IO Int)
    -> (ReadBuffer -> IO Int) -> IO Int
    forall a b. (a -> b) -> a -> b
    $ WriteBuffer -> ReadBuffer -> IO Int
    enc WriteBuffer
    dst

    -- The maximum length of Huffman code is 30.
    -- 40 is enough as a work space.
    initialOffset :: Int
    initialOffset :: Int
    initialOffset = Int
    40

    shiftForWrite :: Int
    shiftForWrite :: Int
    shiftForWrite = Int
    32

    enc :: WriteBuffer -> ReadBuffer -> IO Int
    enc :: WriteBuffer -> ReadBuffer -> IO Int
    enc WriteBuffer{Buffer
    IORef Buffer
    start :: Buffer
    limit :: Buffer
    offset :: IORef Buffer
    oldoffset :: IORef Buffer
    oldoffset :: WriteBuffer -> IORef Buffer
    offset :: WriteBuffer -> IORef Buffer
    limit :: WriteBuffer -> Buffer
    start :: WriteBuffer -> Buffer
    ..} ReadBuffer
    rbuf = do
        beg <- IORef Buffer -> IO Buffer
    forall a. IORef a -> IO a
    readIORef IORef Buffer
    offset
        end <- go (beg, 0, initialOffset)
        writeIORef offset end
        let len = Buffer
    end Buffer -> Buffer -> Int
    forall a b. Ptr a -> Ptr b -> Int
    `minusPtr` Buffer
    beg
        return len
      where
        go :: (Buffer, Word64, Int) -> IO Buffer
    go (Buffer
    dst, Word64
    encoded, Int
    off) = do
            i <- ReadBuffer -> IO Int
    forall a. Readable a => a -> IO Int
    readInt8 ReadBuffer
    rbuf
            if i >= 0
                then cpy dst (bond i) >>= go
                else
                    if off == initialOffset
                        then return dst
                        else do
                            let (encoded1, _) = bond idxEos
                            write dst encoded1
          where
            {-# INLINE bond #-}
            bond :: Int -> (Word64, Int)
    bond Int
    i = (Word64
    encoded', Int
    off')
              where
                len :: Int
    len = UArray Int Int
    huffmanLength UArray Int Int -> Int -> Int
    forall i. Ix i => UArray i Int -> Int -> Int
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> Int -> e
    `unsafeAt` Int
    i
                code :: Word64
    code = UArray Int Word64
    huffmanCode UArray Int Word64 -> Int -> Word64
    forall i. Ix i => UArray i Word64 -> Int -> Word64
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> Int -> e
    `unsafeAt` Int
    i
                scode :: Word64
    scode = Word64
    code Word64 -> Int -> Word64
    forall a. Bits a => a -> Int -> a
    `shiftL` (Int
    off Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    len)
                encoded' :: Word64
    encoded' = Word64
    encoded Word64 -> Word64 -> Word64
    forall a. Bits a => a -> a -> a
    .|. Word64
    scode
                off' :: Int
    off' = Int
    off Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    len
            {-# INLINE write #-}
            write :: Buffer -> p -> IO (Ptr b)
    write Buffer
    p p
    w = do
                Bool -> IO () -> IO ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Buffer
    p Buffer -> Buffer -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Buffer
    limit) (IO () -> IO ()) -> IO () -> IO ()
    forall a b. (a -> b) -> a -> b
    $ BufferOverrun -> IO ()
    forall e a. (HasCallStack, Exception e) => e -> IO a
    throwIO BufferOverrun
    BufferOverrun
                let w8 :: Word8
    w8 = p -> Word8
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (p
    w p -> Int -> p
    forall a. Bits a => a -> Int -> a
    `shiftR` Int
    shiftForWrite) :: Word8
                Buffer -> Word8 -> IO ()
    forall a. Storable a => Ptr a -> a -> IO ()
    poke Buffer
    p Word8
    w8
                let p' :: Ptr b
    p' = Buffer
    p Buffer -> Int -> Ptr b
    forall a b. Ptr a -> Int -> Ptr b
    `plusPtr` Int
    1
                Ptr b -> IO (Ptr b)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Ptr b
    forall {b}. Ptr b
    p'
            {-# INLINE cpy #-}
            cpy :: Buffer -> (a, Int) -> IO (Buffer, a, Int)
    cpy Buffer
    p (a
    w, Int
    o)
                | Int
    o Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    shiftForWrite = (Buffer, a, Int) -> IO (Buffer, a, Int)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Buffer
    p, a
    w, Int
    o)
                | Bool
    otherwise = do
                    p' <- Buffer -> a -> IO Buffer
    forall {p} {b}. (Integral p, Bits p) => Buffer -> p -> IO (Ptr b)
    write Buffer
    p a
    w
                    let w' = a
    w a -> Int -> a
    forall a. Bits a => a -> Int -> a
    `shiftL` Int
    8
                        o' = Int
    o Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    8
                    cpy p' (w', o')

    -- | Huffman encoding with a temporary buffer whose size is 4096.
    encodeHuffman :: ByteString -> IO ByteString
    encodeHuffman :: ByteString -> IO ByteString
    encodeHuffman ByteString
    bs = Int -> (WriteBuffer -> IO ()) -> IO ByteString
    withWriteBuffer Int
    4096 ((WriteBuffer -> IO ()) -> IO ByteString)
    -> (WriteBuffer -> IO ()) -> IO ByteString
    forall a b. (a -> b) -> a -> b
    $ \WriteBuffer
    wbuf ->
        IO Int -> IO ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (IO Int -> IO ()) -> IO Int -> IO ()
    forall a b. (a -> b) -> a -> b
    $ WriteBuffer -> ByteString -> IO Int
    encodeH WriteBuffer
    wbuf ByteString
    bs