{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_HADDOCK hide #-}

-- |
-- Module      : Crypto.Internal.ByteArray
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : stable
-- Portability : Good
--
-- Simple and efficient byte array types
module Crypto.Internal.ByteArray (
    module Data.ByteArray,
    module Data.ByteArray.Mapping,
    module Data.ByteArray.Encoding,
    constAllZero,
    allocAndFreezePrimIO,
    allocAndFreezePrim,
) where

import Data.ByteArray
import Data.ByteArray.Encoding
import Data.ByteArray.Mapping

import Data.Bits ((.|.))
import qualified Data.Primitive.ByteArray as Prim
import Data.Word (Word8)
import Foreign.Ptr (Ptr, castPtr)
import Foreign.Storable (peekByteOff)

import Crypto.Internal.Compat (unsafeDoIO)

-- | Allocate a pinned 'Prim.ByteArray' of the given size, populate it via a
-- 'Ptr', then freeze and return it.  The pointer must not be retained after
-- the action returns.
allocAndFreezePrimIO :: Int -> (Ptr p -> IO ()) -> IO Prim.ByteArray
allocAndFreezePrimIO :: forall p. Int -> (Ptr p -> IO ()) -> IO ByteArray
allocAndFreezePrimIO Int
n Ptr p -> IO ()
f = do
    mba <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
Prim.newPinnedByteArray Int
n
    f (castPtr (Prim.mutableByteArrayContents mba))
    Prim.unsafeFreezeByteArray mba

-- | The allocation is strictly local,
-- the computation is deterministic, and no IO effects escape.
allocAndFreezePrim :: Int -> (Ptr p -> IO ()) -> Prim.ByteArray
allocAndFreezePrim :: forall p. Int -> (Ptr p -> IO ()) -> ByteArray
allocAndFreezePrim Int
n = IO ByteArray -> ByteArray
forall a. IO a -> a
unsafeDoIO (IO ByteArray -> ByteArray)
-> ((Ptr p -> IO ()) -> IO ByteArray)
-> (Ptr p -> IO ())
-> ByteArray
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Ptr p -> IO ()) -> IO ByteArray
forall p. Int -> (Ptr p -> IO ()) -> IO ByteArray
allocAndFreezePrimIO Int
n

constAllZero :: ByteArrayAccess ba => ba -> Bool
constAllZero :: forall ba. ByteArrayAccess ba => ba -> Bool
constAllZero ba
b = IO Bool -> Bool
forall a. IO a -> a
unsafeDoIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ba -> (Ptr (ZonkAny 0) -> IO Bool) -> IO Bool
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. ba -> (Ptr p -> IO a) -> IO a
withByteArray ba
b ((Ptr (ZonkAny 0) -> IO Bool) -> IO Bool)
-> (Ptr (ZonkAny 0) -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr (ZonkAny 0)
p -> Ptr (ZonkAny 0) -> Int -> Word8 -> IO Bool
forall b. Ptr b -> Int -> Word8 -> IO Bool
loop Ptr (ZonkAny 0)
p Int
0 Word8
0
  where
    loop :: Ptr b -> Int -> Word8 -> IO Bool
    loop :: forall b. Ptr b -> Int -> Word8 -> IO Bool
loop Ptr b
p Int
i !Word8
acc
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$! Word8
acc Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0
        | Bool
otherwise = do
            e <- Ptr b -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr b
p Int
i
            loop p (i + 1) (acc .|. e)
    len :: Int
len = ba -> Int
forall ba. ByteArrayAccess ba => ba -> Int
Data.ByteArray.length ba
b