<commit_before>{-# LANGUAGE BangPatterns #-}

module Network.HPACK.Huffman.Code (
  -- * Huffman encoding
    Encoder
  , toEncoder
  , HuffmanEncoding
  , encode
  -- * Huffman decoding
  , Decoder
  , toDecoder
  , HuffmanDecoding
  , decode
  ) where

import Control.Applicative ((<$>))
import Control.Arrow (second)
import Data.Array (Array, (!), listArray)
import Data.Bits ((.&.), shiftR)
import Data.ByteString.Internal (ByteString(..))
import Data.List (partition)
import Data.Word (Word8)
import Foreign.ForeignPtr
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (peek)
import Network.HPACK.Builder
import Network.HPACK.Builder.Word8
import Network.HPACK.Huffman.Bit
import Network.HPACK.Types (DecodeError(..))
import System.IO.Unsafe (unsafePerformIO)

----------------------------------------------------------------

-- | Huffman encoding.
type HuffmanEncoding = ByteString -> ByteString

-- | Huffman decoding.
type HuffmanDecoding = ByteString -> Either DecodeError ByteString

----------------------------------------------------------------

idxEos :: Int
idxEos = 256

----------------------------------------------------------------

-- | Type for Huffman encoding.
newtype Encoder = Encoder (Array Int (Int,Bits))

-- | Creating 'Encoder'.
toEncoder :: [Bits] -> Encoder
toEncoder bss = Encoder $ listArray (0,idxEos) (map toEnt bss)
  where
    toEnt bs = (len, bs)
      where
        !len = length bs

-- | Huffman encoding.
encode :: Encoder -> HuffmanEncoding
encode encoder (PS fptr off len) = fromBitsToByteString nBytes (run b)
  where
    (nbits,b0) = unsafePerformIO $ withForeignPtr fptr $ \ptr ->
        loop (ptr `plusPtr` off) 0 0 empty
    (nBytes,b) = eos nbits b0
    loop :: Ptr Word8 -> Int -> Int -> Builder Bits -> IO (Int,Builder Bits)
    loop !ptr !cnt !nBits !builder
      | cnt == len = return (nBits,builder)
      | otherwise  = do
          i <- fromIntegral <$> peek ptr
          let (bits,bs) = enc encoder i
              builder' = builder << bs
          loop (ptr `plusPtr` 1) (cnt + 1) (nBits + bits) builder'
    eos nBits !builder
      | r == 0    = (q,builder)
      | otherwise = (q+1, builder << bools')
      where
--        (q,r) = nBits `divMod` 8
        q = nBits `shiftR` 3
        r = nBits .&. 0x7
        (_,bools) = enc encoder idxEos
        bools' = take (8 - r) bools

enc :: Encoder -> Int -> (Int,Bits)
enc (Encoder ary) i = ary ! i

----------------------------------------------------------------

-- | Type for Huffman decoding.
data Decoder = Tip (Maybe Int) {-# UNPACK #-} !Int
             | Bin (Maybe Int) Decoder Decoder
             deriving Show

-- | Creating 'Decoder'.
toDecoder :: [Bits] -> Decoder
toDecoder bs = mark 1 eos $ build $ zip [0..idxEos] bs
  where
    eos = bs !! idxEos

build :: [(Int,Bits)] -> Decoder
build [(v,[])] = Tip Nothing v
build xs       = Bin Nothing (build fs) (build ts)
  where
    (fs',ts') = partition ((==) F . head . snd) xs
    fs = map (second tail) fs'
    ts = map (second tail) ts'

mark :: Int -> Bits -> Decoder -> Decoder
mark i []     (Tip Nothing v)   = Tip (Just i) v
mark i (F:bs) (Bin Nothing l r) = Bin (Just i) (mark (i+1) bs l) r
mark i (T:bs) (Bin Nothing l r) = Bin (Just i) l (mark (i+1) bs r)
mark _ _      _                 = error "mark"

-- | Huffman decoding.
decode :: Decoder -> HuffmanDecoding
decode decoder bs = decodeBits decoder src w8empty
  where
    src = toBitSource bs<commit_msg>adding printTree.<commit_after>{-# LANGUAGE BangPatterns #-}

module Network.HPACK.Huffman.Code (
  -- * Huffman encoding
    Encoder
  , toEncoder
  , HuffmanEncoding
  , encode
  -- * Huffman decoding
  , Decoder
  , toDecoder
  , HuffmanDecoding
  , decode
  , printTree
  ) where

import Control.Applicative ((<$>))
import Control.Arrow (second)
import Data.Array (Array, (!), listArray)
import Data.Bits ((.&.), shiftR)
import Data.ByteString.Internal (ByteString(..))
import Data.List (partition)
import Data.Word (Word8)
import Foreign.ForeignPtr
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (peek)
import Network.HPACK.Builder
import Network.HPACK.Builder.Word8
import Network.HPACK.Huffman.Bit
import Network.HPACK.Types (DecodeError(..))
import System.IO.Unsafe (unsafePerformIO)

----------------------------------------------------------------

-- | Huffman encoding.
type HuffmanEncoding = ByteString -> ByteString

-- | Huffman decoding.
type HuffmanDecoding = ByteString -> Either DecodeError ByteString

----------------------------------------------------------------

idxEos :: Int
idxEos = 256

----------------------------------------------------------------

-- | Type for Huffman encoding.
newtype Encoder = Encoder (Array Int (Int,Bits))

-- | Creating 'Encoder'.
toEncoder :: [Bits] -> Encoder
toEncoder bss = Encoder $ listArray (0,idxEos) (map toEnt bss)
  where
    toEnt bs = (len, bs)
      where
        !len = length bs

-- | Huffman encoding.
encode :: Encoder -> HuffmanEncoding
encode encoder (PS fptr off len) = fromBitsToByteString nBytes (run b)
  where
    (nbits,b0) = unsafePerformIO $ withForeignPtr fptr $ \ptr ->
        loop (ptr `plusPtr` off) 0 0 empty
    (nBytes,b) = eos nbits b0
    loop :: Ptr Word8 -> Int -> Int -> Builder Bits -> IO (Int,Builder Bits)
    loop !ptr !cnt !nBits !builder
      | cnt == len = return (nBits,builder)
      | otherwise  = do
          i <- fromIntegral <$> peek ptr
          let (bits,bs) = enc encoder i
              builder' = builder << bs
          loop (ptr `plusPtr` 1) (cnt + 1) (nBits + bits) builder'
    eos nBits !builder
      | r == 0    = (q,builder)
      | otherwise = (q+1, builder << bools')
      where
--        (q,r) = nBits `divMod` 8
        q = nBits `shiftR` 3
        r = nBits .&. 0x7
        (_,bools) = enc encoder idxEos
        bools' = take (8 - r) bools

enc :: Encoder -> Int -> (Int,Bits)
enc (Encoder ary) i = ary ! i

----------------------------------------------------------------

-- | Type for Huffman decoding.
data Decoder = Tip (Maybe Int) {-# UNPACK #-} !Int
             | Bin (Maybe Int) Decoder Decoder
             deriving Show

showTree :: Decoder -> String
showTree = showTree' ""

showTree' :: String -> Decoder -> String
showTree' _    (Tip _ i)   = show i ++ "\n"
showTree' pref (Bin _ l r) = "\n"
                             ++ pref ++ "+ " ++ showTree' pref' l
                             ++ pref ++ "+ " ++ showTree' pref' r
  where
    pref' = "  " ++ pref

printTree :: Decoder -> IO ()
printTree = putStr . showTree

-- | Creating 'Decoder'.
toDecoder :: [Bits] -> Decoder
toDecoder bs = mark 1 eos $ build $ zip [0..idxEos] bs
  where
    eos = bs !! idxEos

build :: [(Int,Bits)] -> Decoder
build [(v,[])] = Tip Nothing v
build xs       = Bin Nothing (build fs) (build ts)
  where
    (fs',ts') = partition ((==) F . head . snd) xs
    fs = map (second tail) fs'
    ts = map (second tail) ts'

mark :: Int -> Bits -> Decoder -> Decoder
mark i []     (Tip Nothing v)   = Tip (Just i) v
mark i (F:bs) (Bin Nothing l r) = Bin (Just i) (mark (i+1) bs l) r
mark i (T:bs) (Bin Nothing l r) = Bin (Just i) l (mark (i+1) bs r)
mark _ _      _                 = error "mark"

-- | Huffman decoding.
decode :: Decoder -> HuffmanDecoding
decode decoder bs = decodeBits decoder src w8empty
  where
    src = toBitSource bs
{'fraction_non_alphanumeric': 0.19383440185513573, 'fraction_numerical': 0.006820351930159597, 'mean_word_length': 2.036024844720497, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7031848', 'n_tokens_mistral': 2413, 'n_tokens_neox': 2306, 'n_words': 753}