398 lines
13 KiB
Haskell
398 lines
13 KiB
Haskell
-- |
|
|
-- Module : Data.ByteArray.Sized
|
|
-- License : BSD-style
|
|
-- Maintainer : Nicolas Di Prima <nicolas@primetype.co.uk>
|
|
-- Stability : stable
|
|
-- Portability : Good
|
|
--
|
|
|
|
{-# LANGUAGE CPP #-}
|
|
{-# LANGUAGE DataKinds #-}
|
|
{-# LANGUAGE TypeApplications #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
|
|
{-# LANGUAGE NoImplicitPrelude #-}
|
|
{-# LANGUAGE BangPatterns #-}
|
|
{-# LANGUAGE TypeFamilies #-}
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE TypeOperators #-}
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
{-# LANGUAGE FlexibleInstances #-}
|
|
{-# LANGUAGE MultiParamTypeClasses #-}
|
|
{-# LANGUAGE FunctionalDependencies #-}
|
|
{-# LANGUAGE UndecidableInstances #-}
|
|
#if __GLASGOW_HASKELL__ >= 806
|
|
{-# LANGUAGE NoStarIsType #-}
|
|
#endif
|
|
|
|
module Data.ByteArray.Sized
|
|
( ByteArrayN(..)
|
|
, SizedByteArray
|
|
, unSizedByteArray
|
|
, sizedByteArray
|
|
, unsafeSizedByteArray
|
|
|
|
, -- * ByteArrayN operators
|
|
alloc
|
|
, create
|
|
, allocAndFreeze
|
|
, unsafeCreate
|
|
, inlineUnsafeCreate
|
|
, empty
|
|
, pack
|
|
, unpack
|
|
, cons
|
|
, snoc
|
|
, xor
|
|
, index
|
|
, splitAt
|
|
, take
|
|
, drop
|
|
, append
|
|
, copy
|
|
, copyRet
|
|
, copyAndFreeze
|
|
, replicate
|
|
, zero
|
|
, convert
|
|
, fromByteArrayAccess
|
|
, unsafeFromByteArrayAccess
|
|
) where
|
|
|
|
import Basement.Imports
|
|
import Basement.NormalForm
|
|
import Basement.Nat
|
|
import Basement.Numerical.Additive ((+))
|
|
import Basement.Numerical.Subtractive ((-))
|
|
|
|
import Basement.Sized.List (ListN, unListN, toListN)
|
|
|
|
import Foreign.Storable
|
|
import Foreign.Ptr
|
|
import Data.Maybe (fromMaybe)
|
|
|
|
import Data.Memory.Internal.Compat
|
|
import Data.Memory.PtrMethods
|
|
|
|
import Data.Proxy (Proxy(..))
|
|
|
|
import Data.ByteArray.Types (ByteArrayAccess(..), ByteArray)
|
|
import qualified Data.ByteArray.Types as ByteArray (allocRet)
|
|
|
|
#if MIN_VERSION_basement(0,0,7)
|
|
import Basement.BlockN (BlockN)
|
|
import qualified Basement.BlockN as BlockN
|
|
import qualified Basement.PrimType as Base
|
|
import Basement.Types.OffsetSize (Countable)
|
|
#endif
|
|
|
|
-- | Type class to emulate exactly the behaviour of 'ByteArray' but with
|
|
-- a known length at compile time
|
|
--
|
|
class (ByteArrayAccess c, KnownNat n) => ByteArrayN (n :: Nat) c | c -> n where
|
|
-- | just like 'allocRet' but with the size at the type level
|
|
allocRet :: forall p a
|
|
. Proxy n
|
|
-> (Ptr p -> IO a)
|
|
-> IO (a, c)
|
|
|
|
-- | Wrapper around any collection type with the size as type parameter
|
|
--
|
|
newtype SizedByteArray (n :: Nat) ba = SizedByteArray { unSizedByteArray :: ba }
|
|
deriving (Eq, Show, Typeable, Ord, NormalForm)
|
|
|
|
-- | create a 'SizedByteArray' from the given 'ByteArrayAccess' if the
|
|
-- size is the same as the target size.
|
|
--
|
|
sizedByteArray :: forall n ba . (KnownNat n, ByteArrayAccess ba)
|
|
=> ba
|
|
-> Maybe (SizedByteArray n ba)
|
|
sizedByteArray ba
|
|
| length ba == n = Just $ SizedByteArray ba
|
|
| otherwise = Nothing
|
|
where
|
|
n = fromInteger $ natVal (Proxy @n)
|
|
|
|
-- | just like the 'sizedByteArray' function but throw an exception if
|
|
-- the size is invalid.
|
|
unsafeSizedByteArray :: forall n ba . (ByteArrayAccess ba, KnownNat n) => ba -> SizedByteArray n ba
|
|
unsafeSizedByteArray = fromMaybe (error "The size is invalid") . sizedByteArray
|
|
|
|
instance (ByteArrayAccess ba, KnownNat n) => ByteArrayAccess (SizedByteArray n ba) where
|
|
length _ = fromInteger $ natVal (Proxy @n)
|
|
withByteArray (SizedByteArray ba) = withByteArray ba
|
|
|
|
instance (KnownNat n, ByteArray ba) => ByteArrayN n (SizedByteArray n ba) where
|
|
allocRet p f = do
|
|
(a, ba) <- ByteArray.allocRet n f
|
|
pure (a, SizedByteArray ba)
|
|
where
|
|
n = fromInteger $ natVal p
|
|
|
|
#if MIN_VERSION_basement(0,0,7)
|
|
instance ( ByteArrayAccess (BlockN n ty)
|
|
, PrimType ty
|
|
, KnownNat n
|
|
, Countable ty n
|
|
, KnownNat nbytes
|
|
, nbytes ~ (Base.PrimSize ty * n)
|
|
) => ByteArrayN nbytes (BlockN n ty) where
|
|
allocRet _ f = do
|
|
mba <- BlockN.new @n
|
|
a <- BlockN.withMutablePtrHint True False mba (f . castPtr)
|
|
ba <- BlockN.freeze mba
|
|
return (a, ba)
|
|
#endif
|
|
|
|
|
|
-- | Allocate a new bytearray of specific size, and run the initializer on this memory
|
|
alloc :: forall n ba p . (ByteArrayN n ba, KnownNat n)
|
|
=> (Ptr p -> IO ())
|
|
-> IO ba
|
|
alloc f = snd <$> allocRet (Proxy @n) f
|
|
|
|
-- | Allocate a new bytearray of specific size, and run the initializer on this memory
|
|
create :: forall n ba p . (ByteArrayN n ba, KnownNat n)
|
|
=> (Ptr p -> IO ())
|
|
-> IO ba
|
|
create = alloc @n
|
|
{-# NOINLINE create #-}
|
|
|
|
-- | similar to 'allocN' but hide the allocation and initializer in a pure context
|
|
allocAndFreeze :: forall n ba p . (ByteArrayN n ba, KnownNat n)
|
|
=> (Ptr p -> IO ()) -> ba
|
|
allocAndFreeze f = unsafeDoIO (alloc @n f)
|
|
{-# NOINLINE allocAndFreeze #-}
|
|
|
|
-- | similar to 'createN' but hide the allocation and initializer in a pure context
|
|
unsafeCreate :: forall n ba p . (ByteArrayN n ba, KnownNat n)
|
|
=> (Ptr p -> IO ()) -> ba
|
|
unsafeCreate f = unsafeDoIO (alloc @n f)
|
|
{-# NOINLINE unsafeCreate #-}
|
|
|
|
inlineUnsafeCreate :: forall n ba p . (ByteArrayN n ba, KnownNat n)
|
|
=> (Ptr p -> IO ()) -> ba
|
|
inlineUnsafeCreate f = unsafeDoIO (alloc @n f)
|
|
{-# INLINE inlineUnsafeCreate #-}
|
|
|
|
-- | Create an empty byte array
|
|
empty :: forall ba . ByteArrayN 0 ba => ba
|
|
empty = unsafeDoIO (alloc @0 $ \_ -> return ())
|
|
|
|
-- | Pack a list of bytes into a bytearray
|
|
pack :: forall n ba . (ByteArrayN n ba, KnownNat n) => ListN n Word8 -> ba
|
|
pack l = inlineUnsafeCreate @n (fill $ unListN l)
|
|
where fill [] _ = return ()
|
|
fill (x:xs) !p = poke p x >> fill xs (p `plusPtr` 1)
|
|
{-# INLINE fill #-}
|
|
{-# NOINLINE pack #-}
|
|
|
|
-- | Un-pack a bytearray into a list of bytes
|
|
unpack :: forall n ba
|
|
. (ByteArrayN n ba, KnownNat n, NatWithinBound Int n, ByteArrayAccess ba)
|
|
=> ba -> ListN n Word8
|
|
unpack bs = fromMaybe (error "the impossible appened") $ toListN @n $ loop 0
|
|
where !len = length bs
|
|
loop i
|
|
| i == len = []
|
|
| otherwise =
|
|
let !v = unsafeDoIO $ withByteArray bs (`peekByteOff` i)
|
|
in v : loop (i+1)
|
|
|
|
-- | prepend a single byte to a byte array
|
|
cons :: forall ni no bi bo
|
|
. ( ByteArrayN ni bi, ByteArrayN no bo, ByteArrayAccess bi
|
|
, KnownNat ni, KnownNat no
|
|
, (ni + 1) ~ no
|
|
)
|
|
=> Word8 -> bi -> bo
|
|
cons b ba = unsafeCreate @no $ \d -> withByteArray ba $ \s -> do
|
|
pokeByteOff d 0 b
|
|
memCopy (d `plusPtr` 1) s len
|
|
where
|
|
!len = fromInteger $ natVal (Proxy @ni)
|
|
|
|
-- | append a single byte to a byte array
|
|
snoc :: forall bi bo ni no
|
|
. ( ByteArrayN ni bi, ByteArrayN no bo, ByteArrayAccess bi
|
|
, KnownNat ni, KnownNat no
|
|
, (ni + 1) ~ no
|
|
)
|
|
=> bi -> Word8 -> bo
|
|
snoc ba b = unsafeCreate @no $ \d -> withByteArray ba $ \s -> do
|
|
memCopy d s len
|
|
pokeByteOff d len b
|
|
where
|
|
!len = fromInteger $ natVal (Proxy @ni)
|
|
|
|
-- | Create a xor of bytes between a and b.
|
|
--
|
|
-- the returns byte array is the size of the smallest input.
|
|
xor :: forall n a b c
|
|
. ( ByteArrayN n a, ByteArrayN n b, ByteArrayN n c
|
|
, ByteArrayAccess a, ByteArrayAccess b
|
|
, KnownNat n
|
|
)
|
|
=> a -> b -> c
|
|
xor a b =
|
|
unsafeCreate @n $ \pc ->
|
|
withByteArray a $ \pa ->
|
|
withByteArray b $ \pb ->
|
|
memXor pc pa pb n
|
|
where
|
|
n = fromInteger (natVal (Proxy @n))
|
|
|
|
-- | return a specific byte indexed by a number from 0 in a bytearray
|
|
--
|
|
-- unsafe, no bound checking are done
|
|
index :: forall n na ba
|
|
. ( ByteArrayN na ba, ByteArrayAccess ba
|
|
, KnownNat na, KnownNat n
|
|
, n <= na
|
|
)
|
|
=> ba -> Proxy n -> Word8
|
|
index b pi = unsafeDoIO $ withByteArray b $ \p -> peek (p `plusPtr` i)
|
|
where
|
|
i = fromInteger $ natVal pi
|
|
|
|
-- | Split a bytearray at a specific length in two bytearray
|
|
splitAt :: forall nblhs nbi nbrhs bi blhs brhs
|
|
. ( ByteArrayN nbi bi, ByteArrayN nblhs blhs, ByteArrayN nbrhs brhs
|
|
, ByteArrayAccess bi
|
|
, KnownNat nbi, KnownNat nblhs, KnownNat nbrhs
|
|
, nblhs <= nbi, (nbrhs + nblhs) ~ nbi
|
|
)
|
|
=> bi -> (blhs, brhs)
|
|
splitAt bs = unsafeDoIO $
|
|
withByteArray bs $ \p -> do
|
|
b1 <- alloc @nblhs $ \r -> memCopy r p n
|
|
b2 <- alloc @nbrhs $ \r -> memCopy r (p `plusPtr` n) (len - n)
|
|
return (b1, b2)
|
|
where
|
|
n = fromInteger $ natVal (Proxy @nblhs)
|
|
len = length bs
|
|
|
|
-- | Take the first @n@ byte of a bytearray
|
|
take :: forall nbo nbi bi bo
|
|
. ( ByteArrayN nbi bi, ByteArrayN nbo bo
|
|
, ByteArrayAccess bi
|
|
, KnownNat nbi, KnownNat nbo
|
|
, nbo <= nbi
|
|
)
|
|
=> bi -> bo
|
|
take bs = unsafeCreate @nbo $ \d -> withByteArray bs $ \s -> memCopy d s m
|
|
where
|
|
!m = min len n
|
|
!len = length bs
|
|
!n = fromInteger $ natVal (Proxy @nbo)
|
|
|
|
-- | drop the first @n@ byte of a bytearray
|
|
drop :: forall n nbi nbo bi bo
|
|
. ( ByteArrayN nbi bi, ByteArrayN nbo bo
|
|
, ByteArrayAccess bi
|
|
, KnownNat n, KnownNat nbi, KnownNat nbo
|
|
, (nbo + n) ~ nbi
|
|
)
|
|
=> Proxy n -> bi -> bo
|
|
drop pn bs = unsafeCreate @nbo $ \d ->
|
|
withByteArray bs $ \s ->
|
|
memCopy d (s `plusPtr` ofs) nb
|
|
where
|
|
ofs = min len n
|
|
nb = len - ofs
|
|
len = length bs
|
|
n = fromInteger $ natVal pn
|
|
|
|
-- | append one bytearray to the other
|
|
append :: forall nblhs nbrhs nbout blhs brhs bout
|
|
. ( ByteArrayN nblhs blhs, ByteArrayN nbrhs brhs, ByteArrayN nbout bout
|
|
, ByteArrayAccess blhs, ByteArrayAccess brhs
|
|
, KnownNat nblhs, KnownNat nbrhs, KnownNat nbout
|
|
, (nbrhs + nblhs) ~ nbout
|
|
)
|
|
=> blhs -> brhs -> bout
|
|
append blhs brhs = unsafeCreate @nbout $ \p ->
|
|
withByteArray blhs $ \plhs ->
|
|
withByteArray brhs $ \prhs -> do
|
|
memCopy p plhs (length blhs)
|
|
memCopy (p `plusPtr` length blhs) prhs (length brhs)
|
|
|
|
-- | Duplicate a bytearray into another bytearray, and run an initializer on it
|
|
copy :: forall n bs1 bs2 p
|
|
. ( ByteArrayN n bs1, ByteArrayN n bs2
|
|
, ByteArrayAccess bs1
|
|
, KnownNat n
|
|
)
|
|
=> bs1 -> (Ptr p -> IO ()) -> IO bs2
|
|
copy bs f = alloc @n $ \d -> do
|
|
withByteArray bs $ \s -> memCopy d s (length bs)
|
|
f (castPtr d)
|
|
|
|
-- | Similar to 'copy' but also provide a way to return a value from the initializer
|
|
copyRet :: forall n bs1 bs2 p a
|
|
. ( ByteArrayN n bs1, ByteArrayN n bs2
|
|
, ByteArrayAccess bs1
|
|
, KnownNat n
|
|
)
|
|
=> bs1 -> (Ptr p -> IO a) -> IO (a, bs2)
|
|
copyRet bs f =
|
|
allocRet (Proxy @n) $ \d -> do
|
|
withByteArray bs $ \s -> memCopy d s (length bs)
|
|
f (castPtr d)
|
|
|
|
-- | Similiar to 'copy' but expect the resulting bytearray in a pure context
|
|
copyAndFreeze :: forall n bs1 bs2 p
|
|
. ( ByteArrayN n bs1, ByteArrayN n bs2
|
|
, ByteArrayAccess bs1
|
|
, KnownNat n
|
|
)
|
|
=> bs1 -> (Ptr p -> IO ()) -> bs2
|
|
copyAndFreeze bs f =
|
|
inlineUnsafeCreate @n $ \d -> do
|
|
copyByteArrayToPtr bs d
|
|
f (castPtr d)
|
|
{-# NOINLINE copyAndFreeze #-}
|
|
|
|
-- | Create a bytearray of a specific size containing a repeated byte value
|
|
replicate :: forall n ba . (ByteArrayN n ba, KnownNat n)
|
|
=> Word8 -> ba
|
|
replicate b = inlineUnsafeCreate @n $ \ptr -> memSet ptr b (fromInteger $ natVal $ Proxy @n)
|
|
{-# NOINLINE replicate #-}
|
|
|
|
-- | Create a bytearray of a specific size initialized to 0
|
|
zero :: forall n ba . (ByteArrayN n ba, KnownNat n) => ba
|
|
zero = unsafeCreate @n $ \ptr -> memSet ptr 0 (fromInteger $ natVal $ Proxy @n)
|
|
{-# NOINLINE zero #-}
|
|
|
|
-- | Convert a bytearray to another type of bytearray
|
|
convert :: forall n bin bout
|
|
. ( ByteArrayN n bin, ByteArrayN n bout
|
|
, KnownNat n
|
|
)
|
|
=> bin -> bout
|
|
convert bs = inlineUnsafeCreate @n (copyByteArrayToPtr bs)
|
|
|
|
-- | Convert a ByteArrayAccess to another type of bytearray
|
|
--
|
|
-- This function returns nothing if the size is not compatible
|
|
fromByteArrayAccess :: forall n bin bout
|
|
. ( ByteArrayAccess bin, ByteArrayN n bout
|
|
, KnownNat n
|
|
)
|
|
=> bin -> Maybe bout
|
|
fromByteArrayAccess bs
|
|
| l == n = Just $ inlineUnsafeCreate @n (copyByteArrayToPtr bs)
|
|
| otherwise = Nothing
|
|
where
|
|
l = length bs
|
|
n = fromInteger $ natVal (Proxy @n)
|
|
|
|
-- | Convert a ByteArrayAccess to another type of bytearray
|
|
unsafeFromByteArrayAccess :: forall n bin bout
|
|
. ( ByteArrayAccess bin, ByteArrayN n bout
|
|
, KnownNat n
|
|
)
|
|
=> bin -> bout
|
|
unsafeFromByteArrayAccess bs = case fromByteArrayAccess @n @bin @bout bs of
|
|
Nothing -> error "Invalid Size"
|
|
Just v -> v
|