{-# 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 :: 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
    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 a. a -> IO a
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 a b. IO a -> (a -> IO b) -> IO b
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 a. a -> IO a
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 {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 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. Exception e => e -> IO a
E.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
                Buffer
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
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')

-- | 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