Encode.hs

{-# LANGUAGE RecordWildCards #-}

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

import qualified Control.Exception as E
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 = listArray (0, idxEos) $ map length huffmanTable

huffmanCode :: UArray Int Word64
huffmanCode = listArray (0, idxEos) huffmanTable'

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

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

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

shiftForWrite :: Int
shiftForWrite = 32

enc :: WriteBuffer -> ReadBuffer -> IO Int
enc WriteBuffer{..} rbuf = do
    beg <- readIORef offset
    end <- go (beg, 0, initialOffset)
    writeIORef offset end
    let len = end `minusPtr` beg
    return len
  where
    go (dst, encoded, off) = do
        i <- readInt8 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 i = (encoded', off')
          where
            len = huffmanLength `unsafeAt` i
            code = huffmanCode `unsafeAt` i
            scode = code `shiftL` (off - len)
            encoded' = encoded .|. scode
            off' = off - len
        {-# INLINE write #-}
        write p w = do
            when (p >= limit) $ E.throwIO BufferOverrun
            let w8 = fromIntegral (w `shiftR` shiftForWrite) :: Word8
            poke p w8
            let p' = p `plusPtr` 1
            return p'
        {-# INLINE cpy #-}
        cpy p (w, o)
            | o > shiftForWrite = return (p, w, o)
            | otherwise = do
                p' <- write p w
                let w' = w `shiftL` 8
                    o' = o + 8
                cpy p' (w', o')

-- | Huffman encoding with a temporary buffer whose size is 4096.
encodeHuffman :: ByteString -> IO ByteString
encodeHuffman bs = withWriteBuffer 4096 $ \wbuf ->
    void $ encodeH wbuf bs