<commit_before>{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-cse -fno-full-laziness #-} -- recommended by GHC manual

module Codec.ShonkyCrypt.ShonkyCryptFFI where

import Control.Applicative ((<$>), (<*>))
import Control.Monad (liftM)
import Data.ByteString (ByteString)
import Data.Word (Word8)
import Foreign.C.String (CStringLen, CString)
import Foreign.C.Types (CChar(..), CULong(..), CDouble(..))
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc (mallocBytes, free)
import Foreign.Marshal.Utils (with, fromBool, toBool)
import Foreign.Ptr
import Foreign.Storable (Storable(..))
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.ByteString.Internal as BSI
import qualified Data.ByteString.Unsafe as BSU

#include "shonky-crypt.h"

data ShonkyCryptKey = ShonkyCryptKey
    { sckKeyStart :: !Word8
    , sckKeyInc :: !Word8
    , sckOnlyAlnum :: !Bool
    } deriving Show

instance Storable ShonkyCryptKey where
    sizeOf _ = {#sizeof shonky_crypt_key_t #}
    alignment _ = 4
    peek p = ShonkyCryptKey
               <$> liftM fromIntegral ({#get shonky_crypt_key_t->key_start #} p)
               <*> liftM fromIntegral ({#get shonky_crypt_key_t->key_inc #} p)
               <*> liftM toBool ({#get shonky_crypt_key_t->only_alnum #} p)
    poke p x =
      do {#set shonky_crypt_key_t.key_start #} p (fromIntegral $ sckKeyStart x)
         {#set shonky_crypt_key_t.key_inc #} p (fromIntegral $ sckKeyInc x)
         {#set shonky_crypt_key_t.only_alnum #} p (fromBool $ sckOnlyAlnum x)

{#pointer shonky_crypt_key_t as ShonkyCryptKeyPtr -> ShonkyCryptKey #}

fromMallocedStorable :: Storable a => Ptr a -> IO a
fromMallocedStorable p =
    do key <- peek p
       free p
       return key

newShonkyCryptContextPointer :: Ptr ShonkyCryptContext -> IO ShonkyCryptContext
newShonkyCryptContextPointer p =
    do fp <- newForeignPtr scReleaseContextPtr p
       return $ ShonkyCryptContext fp

unsafePackMallocCStringLen :: CStringLen -> IO ByteString
unsafePackMallocCStringLen (cstr, len) = do
    fp <- newForeignPtr BSI.c_free_finalizer (castPtr cstr)
    return $! BSI.PS fp 0 len

withShonkyCryptContext :: ShonkyCryptContext -> (Ptr ShonkyCryptContext -> IO b) -> IO b
{#pointer shonky_crypt_context_t as ShonkyCryptContext foreign newtype #}

foreign import ccall "shonky-crypt.h &sc_release_context"
  scReleaseContextPtr :: FunPtr (Ptr ShonkyCryptContext -> IO ())

withTrickC2HS :: Storable a => a -> (Ptr a -> IO b) -> IO b
withTrickC2HS = with

{#fun pure unsafe sc_alloc_context_with_key as
    ^ { withTrickC2HS* `ShonkyCryptKey' }
    -> `ShonkyCryptContext' newShonkyCryptContextPointer* #}

withByteStringLen :: ByteString -> ((CString, CULong) -> IO a) -> IO a
withByteStringLen str f = BSU.unsafeUseAsCStringLen str (\(cstr, len) ->
    f (cstr, fromIntegral len))

{#fun pure unsafe sc_entropy as
    ^ { withByteStringLen *`ByteString'& } -> `Double' #}

{#fun pure unsafe sc_copy_context as
    ^ { withShonkyCryptContext* `ShonkyCryptContext' }
    -> `ShonkyCryptContext' newShonkyCryptContextPointer* #}

type InPlaceEnDeCryptFun =
     Ptr ShonkyCryptContext -> Ptr CChar -> Ptr CChar -> CULong -> IO ()
type NewEnDeCryptFun =
     Ptr ShonkyCryptKey -> Ptr CChar -> CULong -> IO CString

scEnDeCryptInplace :: InPlaceEnDeCryptFun
                   -> ShonkyCryptContext
                   -> ByteString
                   -> (ByteString, ShonkyCryptContext)
scEnDeCryptInplace f ctx input =
    let !newContext = scCopyContext ctx
     in unsafePerformIO $
        BSU.unsafeUseAsCStringLen input $ \(inputBytes, inputLen) ->
        withShonkyCryptContext newContext $ \newContext' ->
        do
           outBuffer <- mallocBytes inputLen
           f newContext' inputBytes outBuffer (fromIntegral inputLen)
           out <- unsafePackMallocCStringLen (outBuffer, inputLen)
           return (out, newContext)

encryptS :: ShonkyCryptContext -> ByteString -> (ByteString, ShonkyCryptContext)
encryptS = scEnDeCryptInplace {#call unsafe sc_encrypt_inplace #}
decryptS :: ShonkyCryptContext -> ByteString -> (ByteString, ShonkyCryptContext)
decryptS = scEnDeCryptInplace {#call unsafe sc_decrypt_inplace #}

scEnDecryptNew :: NewEnDeCryptFun -> ShonkyCryptKey -> ByteString -> ByteString
scEnDecryptNew f key input =
    unsafePerformIO $
    BSU.unsafeUseAsCStringLen input $ \(inputBytes, inputLen) ->
    with key $ \key' ->
    do outputC <- f key' inputBytes (fromIntegral inputLen)
       unsafePackMallocCStringLen (outputC, inputLen)
<commit_msg>MINOR: Improve scEnDeCryptInplace Implementation<commit_after>{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-cse -fno-full-laziness #-} -- recommended by GHC manual

module Codec.ShonkyCrypt.ShonkyCryptFFI where

import Control.Applicative ((<$>), (<*>))
import Control.Monad (liftM)
import Data.ByteString (ByteString)
import Data.Word (Word8)
import Foreign.C.String (CStringLen, CString)
import Foreign.C.Types (CChar(..), CULong(..), CDouble(..))
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc (mallocBytes, free)
import Foreign.Marshal.Utils (with, fromBool, toBool)
import Foreign.Ptr
import Foreign.Storable (Storable(..))
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.ByteString.Internal as BSI
import qualified Data.ByteString.Unsafe as BSU

#include "shonky-crypt.h"

data ShonkyCryptKey = ShonkyCryptKey
    { sckKeyStart :: !Word8
    , sckKeyInc :: !Word8
    , sckOnlyAlnum :: !Bool
    } deriving Show

instance Storable ShonkyCryptKey where
    sizeOf _ = {#sizeof shonky_crypt_key_t #}
    alignment _ = 4
    peek p = ShonkyCryptKey
               <$> liftM fromIntegral ({#get shonky_crypt_key_t->key_start #} p)
               <*> liftM fromIntegral ({#get shonky_crypt_key_t->key_inc #} p)
               <*> liftM toBool ({#get shonky_crypt_key_t->only_alnum #} p)
    poke p x =
      do {#set shonky_crypt_key_t.key_start #} p (fromIntegral $ sckKeyStart x)
         {#set shonky_crypt_key_t.key_inc #} p (fromIntegral $ sckKeyInc x)
         {#set shonky_crypt_key_t.only_alnum #} p (fromBool $ sckOnlyAlnum x)

{#pointer shonky_crypt_key_t as ShonkyCryptKeyPtr -> ShonkyCryptKey #}

fromMallocedStorable :: Storable a => Ptr a -> IO a
fromMallocedStorable p =
    do key <- peek p
       free p
       return key

newShonkyCryptContextPointer :: Ptr ShonkyCryptContext -> IO ShonkyCryptContext
newShonkyCryptContextPointer p =
    do fp <- newForeignPtr scReleaseContextPtr p
       return $ ShonkyCryptContext fp

unsafePackMallocCStringLen :: CStringLen -> IO ByteString
unsafePackMallocCStringLen (cstr, len) = do
    fp <- newForeignPtr BSI.c_free_finalizer (castPtr cstr)
    return $! BSI.PS fp 0 len

withShonkyCryptContext :: ShonkyCryptContext -> (Ptr ShonkyCryptContext -> IO b) -> IO b
{#pointer shonky_crypt_context_t as ShonkyCryptContext foreign newtype #}

foreign import ccall "shonky-crypt.h &sc_release_context"
  scReleaseContextPtr :: FunPtr (Ptr ShonkyCryptContext -> IO ())

withTrickC2HS :: Storable a => a -> (Ptr a -> IO b) -> IO b
withTrickC2HS = with

{#fun pure unsafe sc_alloc_context_with_key as
    ^ { withTrickC2HS* `ShonkyCryptKey' }
    -> `ShonkyCryptContext' newShonkyCryptContextPointer* #}

withByteStringLen :: ByteString -> ((CString, CULong) -> IO a) -> IO a
withByteStringLen str f = BSU.unsafeUseAsCStringLen str (\(cstr, len) ->
    f (cstr, fromIntegral len))

{#fun pure unsafe sc_entropy as
    ^ { withByteStringLen *`ByteString'& } -> `Double' #}

{#fun pure unsafe sc_copy_context as
    ^ { withShonkyCryptContext* `ShonkyCryptContext' }
    -> `ShonkyCryptContext' newShonkyCryptContextPointer* #}

type InPlaceEnDeCryptFun =
     Ptr ShonkyCryptContext -> Ptr CChar -> Ptr CChar -> CULong -> IO ()
type NewEnDeCryptFun =
     Ptr ShonkyCryptKey -> Ptr CChar -> CULong -> IO CString

scEnDeCryptInplace :: InPlaceEnDeCryptFun
                   -> ShonkyCryptContext
                   -> ByteString
                   -> (ByteString, ShonkyCryptContext)
scEnDeCryptInplace f ctx input =
     unsafePerformIO $
     BSU.unsafeUseAsCStringLen input $ \(inputBytes, inputLen) ->
     withShonkyCryptContext ctx $ \ctx' ->
     do
        newContext' <- {#call sc_copy_context #} ctx'
        outBuffer <- mallocBytes inputLen
        f newContext' inputBytes outBuffer (fromIntegral inputLen)
        out <- unsafePackMallocCStringLen (outBuffer, inputLen)
        newContext <- newShonkyCryptContextPointer newContext'
        return (out, newContext)

encryptS :: ShonkyCryptContext -> ByteString -> (ByteString, ShonkyCryptContext)
encryptS = scEnDeCryptInplace {#call unsafe sc_encrypt_inplace #}
decryptS :: ShonkyCryptContext -> ByteString -> (ByteString, ShonkyCryptContext)
decryptS = scEnDeCryptInplace {#call unsafe sc_decrypt_inplace #}

scEnDecryptNew :: NewEnDeCryptFun -> ShonkyCryptKey -> ByteString -> ByteString
scEnDecryptNew f key input =
    unsafePerformIO $
    BSU.unsafeUseAsCStringLen input $ \(inputBytes, inputLen) ->
    with key $ \key' ->
    do outputC <- f key' inputBytes (fromIntegral inputLen)
       unsafePackMallocCStringLen (outputC, inputLen)

{'fraction_non_alphanumeric': 0.10032644178454843, 'fraction_numerical': 0.0017410228509249185, 'mean_word_length': 3.789473684210526, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 101, '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': '2356591', 'n_tokens_mistral': 3248, 'n_tokens_neox': 3068, 'n_words': 741}