{-# LANGUAGE BangPatterns, RecordWildCards #-}
module Network.HPACK.Huffman.Encode (
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 (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 (0,Int
idxEos) [Word64]
huffmanTable'
type HuffmanEncoding = WriteBuffer -> ByteString -> IO Int
encode :: HuffmanEncoding
encode :: HuffmanEncoding
encode dst :: WriteBuffer
dst bs :: 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
initialOffset :: Int
initialOffset :: Int
initialOffset = 40
shiftForWrite :: Int
shiftForWrite :: Int
shiftForWrite = 32
enc :: WriteBuffer -> ReadBuffer -> IO Int
enc :: WriteBuffer -> ReadBuffer -> IO Int
enc WriteBuffer{..} rbuf :: 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,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 (dst :: Buffer
dst,encoded :: Word64
encoded,off :: 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
>= 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 (encoded1 :: Word64
encoded1,_) = 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 i :: 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 p :: Buffer
p w :: 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` 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 p :: Buffer
p (w :: a
w,o :: 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` 8
!o' :: Int
o' = Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 8
Buffer -> (a, Int) -> IO (Buffer, a, Int)
cpy Buffer
p' (a
w',Int
o')
encodeHuffman :: ByteString -> IO ByteString
encodeHuffman :: ByteString -> IO ByteString
encodeHuffman bs :: ByteString
bs = Int -> (WriteBuffer -> IO ()) -> IO ByteString
withWriteBuffer 4096 ((WriteBuffer -> IO ()) -> IO ByteString)
-> (WriteBuffer -> IO ()) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \wbuf :: 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