<commit_before>{-# LANGUAGE CPP, ForeignFunctionInterface #-}
{-# LANGUAGE EmptyDataDecls #-}
module Network.Wai.Zlib (compress) where

import Prelude hiding (length)
import Network.Wai hiding (status)

import Data.ByteString.Unsafe
import Data.ByteString.Lazy.Internal (defaultChunkSize)

import Foreign.C
import Foreign.Ptr
import Foreign.Marshal.Alloc
import Control.Exception
import Control.Monad

import Data.ByteString.Char8 hiding (putStrLn)

chunkSize :: Int
chunkSize = defaultChunkSize

compress :: Enumerator -> Enumerator
compress enum =
    Enumerator $ \iter acc ->
        bracket
            c_create_z_stream
            c_free_z_stream
            (compressInner iter acc enum)

compressInner :: (acc -> ByteString -> IO (Either acc acc))
              -> acc
              -> Enumerator
              -> ZStream
              -> IO (Either acc acc)
compressInner iter acc enum stream = allocaBytes chunkSize $ \outbuff -> do
    when (stream == nullPtr) $ fail "Error from zlib library"
    eacc <- runEnumerator enum (compressIter iter stream outbuff) acc
    case eacc of
        Left acc' -> return $ Left acc'
        Right acc' -> finishStream iter stream outbuff acc'

finishStream :: (acc -> ByteString -> IO (Either acc acc))
             -> ZStream
             -> Ptr CChar
             -> acc
             -> IO (Either acc acc)
finishStream iter stream outbuff acc = do
    c_set_avail_out stream outbuff $ fromIntegral chunkSize
    status <- c_call_deflate_finish stream
    newAvail <- fromIntegral `fmap` c_get_avail_out stream
    bs <- packCStringLen (outbuff, chunkSize - newAvail)
    eacc <- iter acc bs
    if status == 0 || newAvail == 0
        then case eacc of
                Left acc' -> return $ Left acc'
                Right acc' -> finishStream iter stream outbuff acc'
        else return eacc

compressIter :: (acc -> ByteString -> IO (Either acc acc))
             -> ZStream
             -> Ptr CChar
             -> acc
             -> ByteString
             -> IO (Either acc acc)
compressIter iter stream outbuff accInit bsInput = do
    unsafeUseAsCStringLen bsInput $ \(cstr, len) ->
        c_set_avail_in stream cstr $ fromIntegral len
    go accInit
  where
    go acc = do
        c_set_avail_out stream outbuff $ fromIntegral chunkSize
        c_call_deflate_noflush stream
        newAvail <- fromIntegral `fmap` c_get_avail_out stream
        bs <- packCStringLen (outbuff, chunkSize - newAvail)
        eacc <- iter acc bs<commit_msg>Switched to unsafe (ie, no-copy) pack functions for zlib<commit_after>{-# LANGUAGE CPP, ForeignFunctionInterface #-}
{-# LANGUAGE EmptyDataDecls #-}
module Network.Wai.Zlib (compress) where

import Prelude hiding (length)
import Network.Wai hiding (status)

import Data.ByteString.Unsafe
import Data.ByteString.Lazy.Internal (defaultChunkSize)

import Foreign.C
import Foreign.Ptr
import Foreign.Marshal.Alloc
import Control.Exception
import Control.Monad

import Data.ByteString.Char8 hiding (putStrLn)

chunkSize :: Int
chunkSize = defaultChunkSize

compress :: Enumerator -> Enumerator
compress enum =
    Enumerator $ \iter acc ->
        bracket
            c_create_z_stream
            c_free_z_stream
            (compressInner iter acc enum)

compressInner :: (acc -> ByteString -> IO (Either acc acc))
              -> acc
              -> Enumerator
              -> ZStream
              -> IO (Either acc acc)
compressInner iter acc enum stream = allocaBytes chunkSize $ \outbuff -> do
    when (stream == nullPtr) $ fail "Error from zlib library"
    eacc <- runEnumerator enum (compressIter iter stream outbuff) acc
    case eacc of
        Left acc' -> return $ Left acc'
        Right acc' -> finishStream iter stream outbuff acc'

finishStream :: (acc -> ByteString -> IO (Either acc acc))
             -> ZStream
             -> Ptr CChar
             -> acc
             -> IO (Either acc acc)
finishStream iter stream outbuff acc = do
    c_set_avail_out stream outbuff $ fromIntegral chunkSize
    status <- c_call_deflate_finish stream
    newAvail <- fromIntegral `fmap` c_get_avail_out stream
    bs <- unsafePackCStringLen (outbuff, chunkSize - newAvail)
    eacc <- iter acc bs
    if status == 0 || newAvail == 0
        then case eacc of
                Left acc' -> return $ Left acc'
                Right acc' -> finishStream iter stream outbuff acc'
        else return eacc

compressIter :: (acc -> ByteString -> IO (Either acc acc))
             -> ZStream
             -> Ptr CChar
             -> acc
             -> ByteString
             -> IO (Either acc acc)
compressIter iter stream outbuff accInit bsInput = do
    unsafeUseAsCStringLen bsInput $ \(cstr, len) ->
        c_set_avail_in stream cstr $ fromIntegral len
    go accInit
  where
    go acc = do
        c_set_avail_out stream outbuff $ fromIntegral chunkSize
        c_call_deflate_noflush stream
        newAvail <- fromIntegral `fmap` c_get_avail_out stream
        bs <- unsafePackCStringLen (outbuff, chunkSize - newAvail)
        eacc <- iter acc bs
{'fraction_non_alphanumeric': 0.07860952004740272, 'fraction_numerical': 0.0011850681414181316, 'mean_word_length': 2.3492063492063493, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 57, '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': '4967748', 'n_tokens_mistral': 1498, 'n_tokens_neox': 1445, 'n_words': 479}