<commit_before>{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures   #-}
{-# LANGUAGE DataKinds        #-}
module Raaz.Primitive.Util
       ( allocBufferFor
       , processByteSource
       , computeDigest
       , transformAndDigest
       , BufferPtr
       , Buffer, getBufferPointer, bufferSize, processBuffer
       , module Raaz.Primitive.Implementation
       ) where

import Control.Monad.IO.Class          (liftIO)
import Data.ByteString          as B
import Data.ByteString.Internal as IB
import Data.Proxy
import Data.Monoid
import Foreign.Ptr                     (castPtr)
import System.IO.Unsafe                (unsafePerformIO)
import GHC.TypeLits

import Raaz.Core
import Raaz.Core.Types.Internal
import Raaz.Primitive.Implementation

-- | The pointer type associated with the buffer used by the
-- implementation.
type BufferPtr = AlignedPointer BufferAlignment

-- | A memory buffer than can handle up to @n@ blocks of data. This
-- happens when you need to do some transformation on internal data
-- using a primitive. An example is using a stream cipher for
-- pseudo-random generation. Note that and additional blocks that
-- might be required at the end for padding is taken care of and one
-- does not need to worry about it here. I.e. @Buffer n@ has enough
-- space for the @n@ blocks and if required any additional blocks that
-- meant for padding.
newtype Buffer (n :: Nat) = Buffer { unBuffer :: Pointer }

-- | Get the underlying pointer for the buffer.
getBufferPointer :: Buffer n -> BufferPtr
getBufferPointer = nextAlignedPtr . unBuffer

-- | The size of data (measured in blocks) that can be safely
-- processed inside this buffer.
bufferSize :: KnownNat n => Proxy (Buffer n) -> BLOCKS Prim
bufferSize = flip blocksOf Proxy . fromIntegral . natVal . nProxy
  where nProxy :: Proxy (Buffer n) -> Proxy n
        nProxy  _ = Proxy

-- | Internal function used by allocation.  WARNING: Not to be exposed
-- else can be confusing with `bufferSize`.
actualBufferSize :: KnownNat n => Proxy (Buffer n) -> BLOCKS Prim
actualBufferSize bproxy = bufferSize bproxy <> additionalBlocks


-- | Process the data in the buffer.
processBuffer :: KnownNat n => Buffer n -> MT Internals ()
processBuffer buf = processBlocks (getBufferPointer buf) $ bufferSize $ pure buf

instance KnownNat n => Memory (Buffer n) where
  memoryAlloc = allocThisBuffer
    where allocThisBuffer = Buffer <$> pointerAlloc sz
          bufferProxy     :: Alloc (Buffer n) -> Proxy (Buffer n)
          bufferProxy _   = Proxy
          algn            = ptrAlignment (Proxy :: Proxy BufferPtr)
          sz              = atLeastAligned (actualBufferSize $ bufferProxy allocThisBuffer) algn

  unsafeToPointer = unBuffer

-- | Allocate a buffer for a primitive.
allocBufferFor :: (KnownNat BufferAlignment, MonadIOCont m)
               => BLOCKS Prim
               -> (BufferPtr  -> m a) -> m a
allocBufferFor blks = allocaAligned totalSize
  where totalSize = blks `mappend` additionalBlocks

-- | Process a byte source.
processByteSource :: (KnownNat BufferAlignment, ByteSource src) => src -> MT Internals ()
processByteSource src = allocBufferFor blks $ \ ptr -> do
  processChunks (processBlocks ptr blks) (processLast ptr) src blks (forgetAlignment ptr)
  where blks       = atLeast l1Cache :: BLOCKS Prim
<commit_msg>add some inlines<commit_after>{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures   #-}
{-# LANGUAGE DataKinds        #-}
module Raaz.Primitive.Util
       ( allocBufferFor
       , processByteSource
       , computeDigest
       , transformAndDigest
       , BufferPtr
       , Buffer, getBufferPointer, bufferSize, processBuffer
       , module Raaz.Primitive.Implementation
       ) where

import Control.Monad.IO.Class          (liftIO)
import Data.ByteString          as B
import Data.ByteString.Internal as IB
import Data.Proxy
import Data.Monoid
import Foreign.Ptr                     (castPtr)
import System.IO.Unsafe                (unsafePerformIO)
import GHC.TypeLits

import Raaz.Core
import Raaz.Core.Types.Internal
import Raaz.Primitive.Implementation

-- | The pointer type associated with the buffer used by the
-- implementation.
type BufferPtr = AlignedPointer BufferAlignment

-- | A memory buffer than can handle up to @n@ blocks of data. This
-- happens when you need to do some transformation on internal data
-- using a primitive. An example is using a stream cipher for
-- pseudo-random generation. Note that and additional blocks that
-- might be required at the end for padding is taken care of and one
-- does not need to worry about it here. I.e. @Buffer n@ has enough
-- space for the @n@ blocks and if required any additional blocks that
-- meant for padding.
newtype Buffer (n :: Nat) = Buffer { unBuffer :: Pointer }

-- | Get the underlying pointer for the buffer.
getBufferPointer :: Buffer n -> BufferPtr
getBufferPointer = nextAlignedPtr . unBuffer

{-# INLINE bufferSize #-}
-- | The size of data (measured in blocks) that can be safely
-- processed inside this buffer.
bufferSize :: KnownNat n => Proxy (Buffer n) -> BLOCKS Prim
bufferSize = flip blocksOf Proxy . fromIntegral . natVal . nProxy
  where nProxy :: Proxy (Buffer n) -> Proxy n
        nProxy  _ = Proxy

-- | Internal function used by allocation.  WARNING: Not to be exposed
-- else can be confusing with `bufferSize`.
actualBufferSize :: KnownNat n => Proxy (Buffer n) -> BLOCKS Prim
actualBufferSize bproxy = bufferSize bproxy <> additionalBlocks


-- | Process the data in the buffer.
{-# INLINE processBuffer #-}
processBuffer :: KnownNat n => Buffer n -> MT Internals ()
processBuffer buf = processBlocks (getBufferPointer buf) $ bufferSize $ pure buf

instance KnownNat n => Memory (Buffer n) where
  memoryAlloc = allocThisBuffer
    where allocThisBuffer = Buffer <$> pointerAlloc sz
          bufferProxy     :: Alloc (Buffer n) -> Proxy (Buffer n)
          bufferProxy _   = Proxy
          algn            = ptrAlignment (Proxy :: Proxy BufferPtr)
          sz              = atLeastAligned (actualBufferSize $ bufferProxy allocThisBuffer) algn

  unsafeToPointer = unBuffer

-- | Allocate a buffer for a primitive.
allocBufferFor :: (KnownNat BufferAlignment, MonadIOCont m)
               => BLOCKS Prim
               -> (BufferPtr  -> m a) -> m a
allocBufferFor blks = allocaAligned totalSize
  where totalSize = blks `mappend` additionalBlocks

-- | Process a byte source.
processByteSource :: (KnownNat BufferAlignment, ByteSource src) => src -> MT Internals ()
processByteSource src = allocBufferFor blks $ \ ptr -> do
  processChunks (processBlocks ptr blks) (processLast ptr) src blks (forgetAlignment ptr)
  where blks       = atLeast l1Cache :: BLOCKS Prim

{'fraction_non_alphanumeric': 0.07969075230449003, 'fraction_numerical': 0.0002973535533749628, 'mean_word_length': 3.0695704779189352, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 41, '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': '3817127', 'n_tokens_mistral': 1858, 'n_tokens_neox': 1807, 'n_words': 720}