<commit_before>{-# LANGUAGE ForeignFunctionInterface, BangPatterns #-}

module Network.HPACK.Huffman.Encode (
  -- * Huffman encoding
    Encoder
  , toEncoder
  , HuffmanEncoding
  , encode
  ) where

import Control.Applicative ((<$>))
import Control.Monad (void, when)
import Data.Array
import Data.Bits ((.|.))
import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString(..), create)
import Data.Word (Word8)
import Foreign.C.Types
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (peek, poke)
import Network.HPACK.Huffman.Bit
import Network.HPACK.Huffman.Params
import System.IO.Unsafe (unsafePerformIO)

data Shifted = Shifted !Int  -- Total bytes
                       !Int  -- How many bits in the last byte
                       ByteString -- Up to 5 bytes
                       deriving Show

-- fixme
-- |
--
-- >>> toShifted [T,T,T,T] 0
-- Shifted 1 4 "\240"
-- >>> toShifted [T,T,T,T] 4
-- Shifted 1 0 "\SI"
-- >>> toShifted [T,T,T,T] 5
-- Shifted 2 1 "\a\128"

toShifted :: Bits -> Int -> Shifted
toShifted bits n = Shifted total r bs
  where
    shifted = replicate n F ++ bits
    len = length shifted
    (q,r) = len `divMod` 8
    total
      | r == 0    = q
      | otherwise = q + 1
    bs = BS.pack $ map fromBits $ group8 shifted
    group8 xs
      | null zs   = pad ys : []
      | otherwise = ys : group8 zs
      where
        (ys,zs) = splitAt 8 xs
    pad xs = take 8 $ xs ++ repeat F

type ShiftedArray = Array Int Shifted

toShiftedArray :: Bits -> ShiftedArray
toShiftedArray bits = listArray (0,7) $ map (toShifted bits) [0..7]

-- | Type for Huffman encoding.
newtype Encoder = Encoder (Array Int ShiftedArray) deriving Show

-- | Creating 'Encoder'.
toEncoder :: [Bits] -> Encoder
toEncoder bss = Encoder $ listArray (0,idxEos) $ map toShiftedArray bss

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

-- | Huffman encoding.
encode :: Encoder -> HuffmanEncoding
encode (Encoder aoa) (PS fptr off len) = unsafePerformIO $ withForeignPtr fptr $ \ptr -> do
    let beg = ptr `plusPtr` off
        end = beg `plusPtr` len
    size <- accumSize beg end 0 0
    create size (\dst -> go dst 0 beg end)
  where
    accumSize :: Ptr Word8 -> Ptr Word8 -> Int -> Int -> IO Int
    accumSize src lim n acc
      | src == lim = return acc
      | otherwise  = do<commit_msg>refactoring.<commit_after>{-# LANGUAGE ForeignFunctionInterface, BangPatterns #-}

module Network.HPACK.Huffman.Encode (
  -- * Huffman encoding
    Encoder
  , toEncoder
  , HuffmanEncoding
  , encode
  ) where

import Control.Applicative ((<$>))
import Control.Monad (void, when)
import Data.Array
import Data.Bits ((.|.))
import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString(..), create)
import Data.Word (Word8)
import Foreign.C.Types
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (peek, poke)
import Network.HPACK.Huffman.Bit
import Network.HPACK.Huffman.Params
import System.IO.Unsafe (unsafePerformIO)

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

-- | Type for Huffman encoding.
newtype Encoder = Encoder (Array Int ShiftedArray) deriving Show

type ShiftedArray = Array Int Shifted

data Shifted = Shifted !Int  -- Total bytes
                       !Int  -- How many bits in the last byte
                       ByteString -- Up to 5 bytes
                       deriving Show

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

-- | Creating 'Encoder'.
toEncoder :: [Bits] -> Encoder
toEncoder bss = Encoder $ listArray (0,idxEos) $ map toShiftedArray bss

-- fixme
-- |
--
-- >>> toShifted [T,T,T,T] 0
-- Shifted 1 4 "\240"
-- >>> toShifted [T,T,T,T] 4
-- Shifted 1 0 "\SI"
-- >>> toShifted [T,T,T,T] 5
-- Shifted 2 1 "\a\128"

toShifted :: Bits -> Int -> Shifted
toShifted bits n = Shifted total r bs
  where
    shifted = replicate n F ++ bits
    len = length shifted
    (q,r) = len `divMod` 8
    total
      | r == 0    = q
      | otherwise = q + 1
    bs = BS.pack $ map fromBits $ group8 shifted
    group8 xs
      | null zs   = pad ys : []
      | otherwise = ys : group8 zs
      where
        (ys,zs) = splitAt 8 xs
    pad xs = take 8 $ xs ++ repeat F


toShiftedArray :: Bits -> ShiftedArray
toShiftedArray bits = listArray (0,7) $ map (toShifted bits) [0..7]

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

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

-- | Huffman encoding.
encode :: Encoder -> HuffmanEncoding
encode (Encoder aoa) (PS fptr off len) = unsafePerformIO $ withForeignPtr fptr $ \ptr -> do
    let beg = ptr `plusPtr` off
        end = beg `plusPtr` len
    size <- accumSize beg end 0 0
    create size (\dst -> go dst 0 beg end)
  where
    accumSize :: Ptr Word8 -> Ptr Word8 -> Int -> Int -> IO Int
    accumSize src lim n acc
      | src == lim = return acc
      | otherwise  = do
{'fraction_non_alphanumeric': 0.15044070512820512, 'fraction_numerical': 0.014022435897435898, 'mean_word_length': 2.2740983606557377, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 47, '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': '2111825', 'n_tokens_mistral': 1677, 'n_tokens_neox': 1567, 'n_words': 529}