{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Codec.CBOR.ByteArray.Internal
-- Copyright   : (c) Ben Gamari 2017-2018
-- License     : BSD3-style (see LICENSE.txt)
--
-- Maintainer  : duncan@community.haskell.org
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Various bytearray utilities
--
module Codec.CBOR.ByteArray.Internal
  ( foldrByteArray
  , copyToAddr
  , isTrue#
  , sameByteArray
  , mkByteArray
  , isByteArrayPinned
  , touch
  ) where

import Control.Monad.ST
import Control.Monad
import GHC.IO (IO(..))
import GHC.Exts
import GHC.Word

import qualified Data.Primitive.ByteArray as Prim

foldrByteArray :: (Word8 -> a -> a) -> a
               -> Int             -- ^ offset
               -> Int             -- ^ length
               -> Prim.ByteArray  -- ^ array
               -> a
foldrByteArray :: (Word8 -> a -> a) -> a -> Int -> Int -> ByteArray -> a
foldrByteArray Word8 -> a -> a
f a
z Int
off0 Int
len ByteArray
ba = Int -> a
go Int
off0
  where
    go :: Int -> a
go !Int
off
      | Int
off Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = a
z
      | Bool
otherwise  =
        let x :: Word8
x = ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
Prim.indexByteArray ByteArray
ba Int
off
        in Word8 -> a -> a
f Word8
x (Int -> a
go (Int
offInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))

copyToAddr :: Prim.ByteArray -> Int -> Ptr a -> Int -> IO ()
copyToAddr :: ByteArray -> Int -> Ptr a -> Int -> IO ()
copyToAddr (Prim.ByteArray ByteArray#
ba) (I# Int#
off) (Ptr Addr#
addr) (I# Int#
len) =
    (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> case ByteArray#
-> Int# -> Addr# -> Int# -> State# RealWorld -> State# RealWorld
forall d.
ByteArray# -> Int# -> Addr# -> Int# -> State# d -> State# d
copyByteArrayToAddr# ByteArray#
ba Int#
off Addr#
addr Int#
len State# RealWorld
s of
                State# RealWorld
s' -> (# State# RealWorld
s', () #))

sameByteArray :: Prim.ByteArray -> Prim.ByteArray -> Bool
sameByteArray :: ByteArray -> ByteArray -> Bool
sameByteArray (Prim.ByteArray ByteArray#
ba1#) (Prim.ByteArray ByteArray#
ba2#) =
    case () -> () -> Int#
forall a. a -> a -> Int#
reallyUnsafePtrEquality# (ByteArray# -> ()
unsafeCoerce# ByteArray#
ba1# :: ()) (ByteArray# -> ()
unsafeCoerce# ByteArray#
ba2# :: ()) of
      Int#
r -> Int# -> Bool
isTrue# Int#
r

-- | @mkByteArray n xs@ forms a 'Prim.ByteArray' with contents @xs@. Note that
-- @n@ must be the precise length of @xs@.
mkByteArray :: Int -> [Word8] -> Prim.ByteArray
mkByteArray :: Int -> [Word8] -> ByteArray
mkByteArray Int
n [Word8]
xs = (forall s. ST s ByteArray) -> ByteArray
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
    MutableByteArray s
arr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
Prim.newByteArray Int
n
    (Int -> Word8 -> ST s ()) -> [Int] -> [Word8] -> ST s ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
Prim.writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr) [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] (Int -> [Word8] -> [Word8]
forall a. Int -> [a] -> [a]
take Int
n ([Word8] -> [Word8]) -> [Word8] -> [Word8]
forall a b. (a -> b) -> a -> b
$ [Word8]
xs [Word8] -> [Word8] -> [Word8]
forall a. [a] -> [a] -> [a]
++ Word8 -> [Word8]
forall a. a -> [a]
repeat Word8
0)
    MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
Prim.unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr

-- | A conservative estimate of pinned-ness.
isByteArrayPinned :: Prim.ByteArray -> Bool
isByteArrayPinned :: ByteArray -> Bool
isByteArrayPinned (Prim.ByteArray ByteArray#
_ba) =
#if __GLASGOW_HASKELL__ > 800
    case ByteArray# -> Int#
isByteArrayPinned# ByteArray#
_ba of
      Int#
0# -> Bool
False
      Int#
_  -> Bool
True
#else
    False
#endif

touch :: a -> IO ()
touch :: a -> IO ()
touch a
x = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case a -> State# RealWorld -> State# RealWorld
forall a. a -> State# RealWorld -> State# RealWorld
touch# a
x State# RealWorld
s of State# RealWorld
s' -> (# State# RealWorld
s', () #)