content: {-# 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 = 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) $ 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
  