{-# 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 (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 = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (Int
0,Int
idxEos) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bits]
huffmanTable

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