{-# LANGUAGE BangPatterns, RecordWildCards #-}

module Network.HPACK.Huffman.Encode (
  -- * Huffman encoding
    HuffmanEncoding
  , encode
  , 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 (plusPtr, minusPtr)
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 (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.
type HuffmanEncoding = WriteBuffer -> ByteString -> IO Int

-- | Huffman encoding.
encode :: HuffmanEncoding
encode :: HuffmanEncoding
encode 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 :: WriteBuffer -> Buffer
limit :: WriteBuffer -> Buffer
offset :: WriteBuffer -> IORef Buffer
oldoffset :: WriteBuffer -> IORef Buffer
oldoffset :: IORef Buffer
offset :: IORef Buffer
limit :: Buffer
start :: Buffer
..} ReadBuffer
rbuf = do
    Buffer
beg <- IORef Buffer -> IO Buffer
forall a. IORef a -> IO a
readIORef IORef Buffer
offset
    Buffer
end <- (Buffer, Word64, Int) -> IO Buffer
go (Buffer
beg,Word64
0,Int
initialOffset)
    IORef Buffer -> Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Buffer
offset Buffer
end
    let !len :: Int
len = Buffer
end Buffer -> Buffer -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Buffer
beg
    Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
len
  where
    go :: (Buffer, Word64, Int) -> IO Buffer
go (Buffer
dst,Word64
encoded,Int
off) = do
        !Int
i <- ReadBuffer -> IO Int
forall a. Readable a => a -> IO Int
readInt8 ReadBuffer
rbuf
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then
            Buffer -> (Word64, Int) -> IO (Buffer, Word64, Int)
forall a.
(Integral a, Bits a) =>
Buffer -> (a, Int) -> IO (Buffer, a, Int)
cpy Buffer
dst (Int -> (Word64, Int)
bond Int
i) IO (Buffer, Word64, Int)
-> ((Buffer, Word64, Int) -> IO Buffer) -> IO Buffer
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Buffer, Word64, Int) -> IO Buffer
go
          else if Int
off Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
initialOffset then
            Buffer -> IO Buffer
forall (m :: * -> *) a. Monad m => a -> m a
return Buffer
dst
          else do
            let (Word64
encoded1,Int
_) = Int -> (Word64, Int)
bond Int
idxEos
            Buffer -> Word64 -> IO Buffer
forall a b. (Integral a, Bits a) => Buffer -> a -> IO (Ptr b)
write Buffer
dst Word64
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 (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 (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 -> a -> IO (Ptr b)
write Buffer
p a
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. Exception e => e -> IO a
throwIO BufferOverrun
BufferOverrun
            let !w8 :: Word8
w8 = a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
w a -> Int -> a
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 (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 (m :: * -> *) a. Monad m => a -> m a
return (Buffer
p,a
w,Int
o)
          | Bool
otherwise = do
              Buffer
p' <- Buffer -> a -> IO Buffer
forall a b. (Integral a, Bits a) => Buffer -> a -> IO (Ptr b)
write Buffer
p a
w
              let !w' :: a
w' = a
w a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
8
                  !o' :: Int
o' = Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8
              Buffer -> (a, Int) -> IO (Buffer, a, Int)
cpy Buffer
p' (a
w',Int
o')

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
$ HuffmanEncoding
encode WriteBuffer
wbuf ByteString
bs