-- |
-- Module      : Data.ByteArray.Mapping
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : stable
-- Portability : Good
--
module Data.ByteArray.Mapping
    ( toW64BE
    , toW64LE
    , fromW64BE
    , mapAsWord64
    , mapAsWord128
    ) where

import           Data.ByteArray.Types
import           Data.ByteArray.Methods
import           Data.Memory.Internal.Compat
import           Data.Memory.Internal.Imports hiding (empty)
import           Data.Memory.Endian
import           Data.Memory.ExtendedWords
import           Foreign.Storable
import           Foreign.Ptr

import           Prelude hiding (length, take, drop, span, concat, replicate, splitAt, null, pred)

-- | Transform a bytearray at a specific offset into
-- a Word64 tagged as BE (Big Endian)
--
-- no bounds checking. unsafe
toW64BE :: ByteArrayAccess bs => bs -> Int -> BE Word64
toW64BE :: forall bs. ByteArrayAccess bs => bs -> Int -> BE Word64
toW64BE bs
bs Int
ofs = IO (BE Word64) -> BE Word64
forall a. IO a -> a
unsafeDoIO (IO (BE Word64) -> BE Word64) -> IO (BE Word64) -> BE Word64
forall a b. (a -> b) -> a -> b
$ bs -> (Ptr Any -> IO (BE Word64)) -> IO (BE Word64)
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. bs -> (Ptr p -> IO a) -> IO a
withByteArray bs
bs ((Ptr Any -> IO (BE Word64)) -> IO (BE Word64))
-> (Ptr Any -> IO (BE Word64)) -> IO (BE Word64)
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> Ptr (BE Word64) -> IO (BE Word64)
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any
p Ptr Any -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
ofs)

-- | Transform a bytearray at a specific offset into
-- a Word64 tagged as LE (Little Endian)
--
-- no bounds checking. unsafe
toW64LE :: ByteArrayAccess bs => bs -> Int -> LE Word64
toW64LE :: forall bs. ByteArrayAccess bs => bs -> Int -> LE Word64
toW64LE bs
bs Int
ofs = IO (LE Word64) -> LE Word64
forall a. IO a -> a
unsafeDoIO (IO (LE Word64) -> LE Word64) -> IO (LE Word64) -> LE Word64
forall a b. (a -> b) -> a -> b
$ bs -> (Ptr Any -> IO (LE Word64)) -> IO (LE Word64)
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. bs -> (Ptr p -> IO a) -> IO a
withByteArray bs
bs ((Ptr Any -> IO (LE Word64)) -> IO (LE Word64))
-> (Ptr Any -> IO (LE Word64)) -> IO (LE Word64)
forall a b. (a -> b) -> a -> b
$ \Ptr Any
p -> Ptr (LE Word64) -> IO (LE Word64)
forall a. Storable a => Ptr a -> IO a
peek (Ptr Any
p Ptr Any -> Int -> Ptr (LE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
ofs)

-- | Serialize a @Word64@ to a @ByteArray@ in big endian format
fromW64BE :: (ByteArray ba) => Word64 -> ba
fromW64BE :: forall ba. ByteArray ba => Word64 -> ba
fromW64BE Word64
n = Int -> (Ptr (BE Word64) -> IO ()) -> ba
forall a p. ByteArray a => Int -> (Ptr p -> IO ()) -> a
allocAndFreeze Int
8 ((Ptr (BE Word64) -> IO ()) -> ba)
-> (Ptr (BE Word64) -> IO ()) -> ba
forall a b. (a -> b) -> a -> b
$ \Ptr (BE Word64)
p -> Ptr (BE Word64) -> BE Word64 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr (BE Word64)
p (Word64 -> BE Word64
forall a. ByteSwap a => a -> BE a
toBE Word64
n)

-- | map blocks of 128 bits of a bytearray, creating a new bytestring
-- of equivalent size where each blocks has been mapped through @f@
--
-- no length checking is done. unsafe
mapAsWord128 :: ByteArray bs => (Word128 -> Word128) -> bs -> bs
mapAsWord128 :: forall bs. ByteArray bs => (Word128 -> Word128) -> bs -> bs
mapAsWord128 Word128 -> Word128
f bs
bs =
    Int -> (Ptr (BE Word64) -> IO ()) -> bs
forall a p. ByteArray a => Int -> (Ptr p -> IO ()) -> a
unsafeCreate Int
len ((Ptr (BE Word64) -> IO ()) -> bs)
-> (Ptr (BE Word64) -> IO ()) -> bs
forall a b. (a -> b) -> a -> b
$ \Ptr (BE Word64)
dst ->
    bs -> (Ptr (BE Word64) -> IO ()) -> IO ()
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. bs -> (Ptr p -> IO a) -> IO a
withByteArray bs
bs ((Ptr (BE Word64) -> IO ()) -> IO ())
-> (Ptr (BE Word64) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr (BE Word64)
src ->
        Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
loop (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
16) Ptr (BE Word64)
dst Ptr (BE Word64)
src
  where
        len :: Int
len        = bs -> Int
forall ba. ByteArrayAccess ba => ba -> Int
length bs
bs
        loop :: Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
        loop :: Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
loop Int
0 Ptr (BE Word64)
_ Ptr (BE Word64)
_ = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        loop Int
i Ptr (BE Word64)
d Ptr (BE Word64)
s = do
            BE Word64
w1 <- Ptr (BE Word64) -> IO (BE Word64)
forall a. Storable a => Ptr a -> IO a
peek Ptr (BE Word64)
s
            BE Word64
w2 <- Ptr (BE Word64) -> IO (BE Word64)
forall a. Storable a => Ptr a -> IO a
peek (Ptr (BE Word64)
s Ptr (BE Word64) -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
8)
            let (Word128 Word64
r1 Word64
r2) = Word128 -> Word128
f (Word64 -> Word64 -> Word128
Word128 (BE Word64 -> Word64
forall a. ByteSwap a => BE a -> a
fromBE BE Word64
w1) (BE Word64 -> Word64
forall a. ByteSwap a => BE a -> a
fromBE BE Word64
w2))
            Ptr (BE Word64) -> BE Word64 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr (BE Word64)
d               (Word64 -> BE Word64
forall a. ByteSwap a => a -> BE a
toBE Word64
r1)
            Ptr (BE Word64) -> BE Word64 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr (BE Word64)
d Ptr (BE Word64) -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
8) (Word64 -> BE Word64
forall a. ByteSwap a => a -> BE a
toBE Word64
r2)
            Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Ptr (BE Word64)
d Ptr (BE Word64) -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
16) (Ptr (BE Word64)
s Ptr (BE Word64) -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
16)

-- | map blocks of 64 bits of a bytearray, creating a new bytestring
-- of equivalent size where each blocks has been mapped through @f@
--
-- no length checking is done. unsafe
mapAsWord64 :: ByteArray bs => (Word64 -> Word64) -> bs -> bs
mapAsWord64 :: forall bs. ByteArray bs => (Word64 -> Word64) -> bs -> bs
mapAsWord64 Word64 -> Word64
f bs
bs =
    Int -> (Ptr (BE Word64) -> IO ()) -> bs
forall a p. ByteArray a => Int -> (Ptr p -> IO ()) -> a
unsafeCreate Int
len ((Ptr (BE Word64) -> IO ()) -> bs)
-> (Ptr (BE Word64) -> IO ()) -> bs
forall a b. (a -> b) -> a -> b
$ \Ptr (BE Word64)
dst ->
    bs -> (Ptr (BE Word64) -> IO ()) -> IO ()
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. bs -> (Ptr p -> IO a) -> IO a
withByteArray bs
bs ((Ptr (BE Word64) -> IO ()) -> IO ())
-> (Ptr (BE Word64) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr (BE Word64)
src ->
        Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
loop (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8) Ptr (BE Word64)
dst Ptr (BE Word64)
src
  where
        len :: Int
len        = bs -> Int
forall ba. ByteArrayAccess ba => ba -> Int
length bs
bs

        loop :: Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
        loop :: Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
loop Int
0 Ptr (BE Word64)
_ Ptr (BE Word64)
_ = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        loop Int
i Ptr (BE Word64)
d Ptr (BE Word64)
s = do
            BE Word64
w <- Ptr (BE Word64) -> IO (BE Word64)
forall a. Storable a => Ptr a -> IO a
peek Ptr (BE Word64)
s
            let r :: Word64
r = Word64 -> Word64
f (BE Word64 -> Word64
forall a. ByteSwap a => BE a -> a
fromBE BE Word64
w)
            Ptr (BE Word64) -> BE Word64 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr (BE Word64)
d (Word64 -> BE Word64
forall a. ByteSwap a => a -> BE a
toBE Word64
r)
            Int -> Ptr (BE Word64) -> Ptr (BE Word64) -> IO ()
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Ptr (BE Word64)
d Ptr (BE Word64) -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
8) (Ptr (BE Word64)
s Ptr (BE Word64) -> Int -> Ptr (BE Word64)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
8)