{-# language BangPatterns #-}
{-# language BlockArguments #-}
{-# language MagicHash #-}
{-# language NamedFieldPuns #-}
{-# language TypeApplications #-}
{-# language UnboxedTuples #-}

module Data.Bytes.Pure
  ( empty
  , emptyPinned
  , emptyPinnedU
  , pin
  , contents
  , unsafeCopy
  , toByteArray
  , toByteArrayClone
  , toPinnedByteArray
  , toPinnedByteArrayClone
  , fromByteArray
  , length
  , foldl
  , foldl'
  , foldr
  , ifoldl'
  , foldr'
  , fnv1a32
  , fnv1a64
  , toByteString
  , pinnedToByteString
  , fromByteString
  ) where

import Prelude hiding (length,foldl,foldr)

import Control.Monad.Primitive (PrimState,PrimMonad)
import Control.Monad.ST.Run (runByteArrayST)
import Data.Bits (xor)
import Data.Bytes.Types (Bytes(Bytes))
import Data.ByteString (ByteString)
import Data.Primitive (ByteArray,MutableByteArray)
import Data.Word (Word64,Word32,Word8)
import Foreign.Ptr (Ptr,plusPtr)
import GHC.IO (unsafeIOToST)

import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Internal as ByteString
import qualified Data.ByteString.Unsafe as ByteString
import qualified Data.Primitive as PM
import qualified Data.Primitive.Ptr as PM
import qualified GHC.Exts as Exts
import qualified GHC.ForeignPtr as ForeignPtr

-- | The empty byte sequence.
empty :: Bytes
empty :: Bytes
empty = ByteArray -> Int -> Int -> Bytes
Bytes ByteArray
forall a. Monoid a => a
mempty Int
0 Int
0

-- | The empty pinned byte sequence.
emptyPinned :: Bytes
emptyPinned :: Bytes
emptyPinned = ByteArray -> Int -> Int -> Bytes
Bytes
  ( (forall s. ST s ByteArray) -> ByteArray
runByteArrayST
    (Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newPinnedByteArray Int
0 ST s (MutableByteArray s)
-> (MutableByteArray s -> ST s ByteArray) -> ST s ByteArray
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MutableByteArray s -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray)
  ) Int
0 Int
0

-- | The empty pinned byte sequence.
emptyPinnedU :: ByteArray
emptyPinnedU :: ByteArray
emptyPinnedU = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST
  (Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newPinnedByteArray Int
0 ST s (MutableByteArray s)
-> (MutableByteArray s -> ST s ByteArray) -> ST s ByteArray
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MutableByteArray s -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray)

-- | Yields a pinned byte sequence whose contents are identical to those
-- of the original byte sequence. If the @ByteArray@ backing the argument
-- was already pinned, this simply aliases the argument and does not perform
-- any copying.
pin :: Bytes -> Bytes
pin :: Bytes -> Bytes
pin b :: Bytes
b@(Bytes ByteArray
arr Int
_ Int
len) = case ByteArray -> Bool
PM.isByteArrayPinned ByteArray
arr of
  Bool
True -> Bytes
b
  Bool
False -> ByteArray -> Int -> Int -> Bytes
Bytes
    ( (forall s. ST s ByteArray) -> ByteArray
runByteArrayST do
        MutableByteArray s
dst <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newPinnedByteArray Int
len
        MutableByteArray (PrimState (ST s)) -> Int -> Bytes -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Bytes -> m ()
unsafeCopy MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
0 Bytes
b
        MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst
    ) Int
0 Int
len

-- | Convert the sliced 'Bytes' to an unsliced 'ByteArray'. This
-- reuses the array backing the sliced 'Bytes' if the slicing metadata
-- implies that all of the bytes are used. Otherwise, it makes a copy.
toByteArray :: Bytes -> ByteArray
{-# inline toByteArray #-}
toByteArray :: Bytes -> ByteArray
toByteArray b :: Bytes
b@(Bytes ByteArray
arr Int
off Int
len)
  | Int
off Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0, ByteArray -> Int
PM.sizeofByteArray ByteArray
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = ByteArray
arr
  | Bool
otherwise = Bytes -> ByteArray
toByteArrayClone Bytes
b

-- | Variant of 'toByteArray' that unconditionally makes a copy of
-- the array backing the sliced 'Bytes' even if the original array
-- could be reused. Prefer 'toByteArray'.
toByteArrayClone :: Bytes -> ByteArray
{-# inline toByteArrayClone #-}
toByteArrayClone :: Bytes -> ByteArray
toByteArrayClone (Bytes ByteArray
arr Int
off Int
len) = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
m <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newByteArray Int
len
  MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
PM.copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
m Int
0 ByteArray
arr Int
off Int
len
  MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
m

-- | Copy the byte sequence into a mutable buffer. The buffer must have
-- enough space to accomodate the byte sequence, but this this is not
-- checked.
unsafeCopy :: PrimMonad m
  => MutableByteArray (PrimState m) -- ^ Destination
  -> Int -- ^ Destination Offset
  -> Bytes -- ^ Source
  -> m ()
{-# inline unsafeCopy #-}
unsafeCopy :: MutableByteArray (PrimState m) -> Int -> Bytes -> m ()
unsafeCopy MutableByteArray (PrimState m)
dst Int
dstIx (Bytes ByteArray
src Int
srcIx Int
len) =
  MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
PM.copyByteArray MutableByteArray (PrimState m)
dst Int
dstIx ByteArray
src Int
srcIx Int
len

-- | Create a slice of 'Bytes' that spans the entire argument array.
fromByteArray :: ByteArray -> Bytes
{-# inline fromByteArray #-}
fromByteArray :: ByteArray -> Bytes
fromByteArray ByteArray
b = ByteArray -> Int -> Int -> Bytes
Bytes ByteArray
b Int
0 (ByteArray -> Int
PM.sizeofByteArray ByteArray
b)

-- | The length of a slice of bytes.
length :: Bytes -> Int
{-# inline length #-}
length :: Bytes -> Int
length (Bytes ByteArray
_ Int
_ Int
len) = Int
len

-- | Hash byte sequence with 32-bit variant of FNV-1a.
fnv1a32 :: Bytes -> Word32
fnv1a32 :: Bytes -> Word32
fnv1a32 !Bytes
b = (Word32 -> Word8 -> Word32) -> Word32 -> Bytes -> Word32
forall a. (a -> Word8 -> a) -> a -> Bytes -> a
foldl'
  (\Word32
acc Word8
w -> (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word32 Word8
w Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` Word32
acc) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
0x01000193
  ) Word32
0x811c9dc5 Bytes
b

-- | Hash byte sequence with 64-bit variant of FNV-1a.
fnv1a64 :: Bytes -> Word64
fnv1a64 :: Bytes -> Word64
fnv1a64 !Bytes
b = (Word64 -> Word8 -> Word64) -> Word64 -> Bytes -> Word64
forall a. (a -> Word8 -> a) -> a -> Bytes -> a
foldl'
  (\Word64
acc Word8
w -> (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word64 Word8
w Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
acc) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0x00000100000001B3
  ) Word64
0xcbf29ce484222325 Bytes
b

-- | Left fold over bytes, non-strict in the accumulator.
foldl :: (a -> Word8 -> a) -> a -> Bytes -> a
{-# inline foldl #-}
foldl :: (a -> Word8 -> a) -> a -> Bytes -> a
foldl a -> Word8 -> a
f a
a0 (Bytes ByteArray
arr Int
off0 Int
len0) =
  Int -> Int -> a
forall t. (Eq t, Num t) => Int -> t -> a
go (Int
off0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
len0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) 
  where
  go :: Int -> t -> a
go !Int
off !t
ix = case t
ix of
    (-1) -> a
a0
    t
_ -> a -> Word8 -> a
f (Int -> t -> a
go (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (t
ix t -> t -> t
forall a. Num a => a -> a -> a
- t
1)) (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
PM.indexByteArray ByteArray
arr Int
off)

-- | Left fold over bytes, strict in the accumulator.
foldl' :: (a -> Word8 -> a) -> a -> Bytes -> a
{-# inline foldl' #-}
foldl' :: (a -> Word8 -> a) -> a -> Bytes -> a
foldl' a -> Word8 -> a
f a
a0 (Bytes ByteArray
arr Int
off0 Int
len0) = a -> Int -> Int -> a
forall t. (Eq t, Num t) => a -> Int -> t -> a
go a
a0 Int
off0 Int
len0 where
  go :: a -> Int -> t -> a
go !a
a !Int
off !t
len = case t
len of
    t
0 -> a
a
    t
_ -> a -> Int -> t -> a
go (a -> Word8 -> a
f a
a (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
PM.indexByteArray ByteArray
arr Int
off)) (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (t
len t -> t -> t
forall a. Num a => a -> a -> a
- t
1)

-- | Right fold over bytes, non-strict in the accumulator.
foldr :: (Word8 -> a -> a) -> a -> Bytes -> a
{-# inline foldr #-}
foldr :: (Word8 -> a -> a) -> a -> Bytes -> a
foldr Word8 -> a -> a
f a
a0 (Bytes ByteArray
arr Int
off0 Int
len0) = Int -> Int -> a
forall t. (Eq t, Num t) => Int -> t -> a
go Int
off0 Int
len0 where
  go :: Int -> t -> a
go !Int
off !t
len = case t
len of
    t
0 -> a
a0
    t
_ -> Word8 -> a -> a
f (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
PM.indexByteArray ByteArray
arr Int
off) (Int -> t -> a
go (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (t
len t -> t -> t
forall a. Num a => a -> a -> a
- t
1))

-- | Left fold over bytes, strict in the accumulator. The reduction function
-- is applied to each element along with its index.
ifoldl' :: (a -> Int -> Word8 -> a) -> a -> Bytes -> a
{-# inline ifoldl' #-}
ifoldl' :: (a -> Int -> Word8 -> a) -> a -> Bytes -> a
ifoldl' a -> Int -> Word8 -> a
f a
a0 (Bytes ByteArray
arr Int
off0 Int
len0) = a -> Int -> Int -> Int -> a
forall t. (Eq t, Num t) => a -> Int -> Int -> t -> a
go a
a0 Int
0 Int
off0 Int
len0 where
  go :: a -> Int -> Int -> t -> a
go !a
a !Int
ix !Int
off !t
len = case t
len of
    t
0 -> a
a
    t
_ -> a -> Int -> Int -> t -> a
go (a -> Int -> Word8 -> a
f a
a Int
ix (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
PM.indexByteArray ByteArray
arr Int
off)) (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (t
len t -> t -> t
forall a. Num a => a -> a -> a
- t
1)

-- | Right fold over bytes, strict in the accumulator.
foldr' :: (Word8 -> a -> a) -> a -> Bytes -> a
{-# inline foldr' #-}
foldr' :: (Word8 -> a -> a) -> a -> Bytes -> a
foldr' Word8 -> a -> a
f a
a0 (Bytes ByteArray
arr Int
off0 Int
len0) =
  a -> Int -> Int -> a
forall t. (Eq t, Num t) => a -> Int -> t -> a
go a
a0 (Int
off0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
len0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) 
  where
  go :: a -> Int -> t -> a
go !a
a !Int
off !t
ix = case t
ix of
    (-1) -> a
a
    t
_ -> a -> Int -> t -> a
go (Word8 -> a -> a
f (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
PM.indexByteArray ByteArray
arr Int
off) a
a) (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (t
ix t -> t -> t
forall a. Num a => a -> a -> a
- t
1)


-- | Yields a pointer to the beginning of the byte sequence. It is only safe
-- to call this on a 'Bytes' backed by a pinned @ByteArray@.
contents :: Bytes -> Ptr Word8
{-# inline contents #-}
contents :: Bytes -> Ptr Word8
contents (Bytes ByteArray
arr Int
off Int
_) = Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr (ByteArray -> Ptr Word8
PM.byteArrayContents ByteArray
arr) Int
off

-- | Convert the sliced 'Bytes' to an unsliced 'ByteArray'. This
-- reuses the array backing the sliced 'Bytes' if the slicing metadata
-- implies that all of the bytes are used and they are already pinned.
-- Otherwise, it makes a copy.
toPinnedByteArray :: Bytes -> ByteArray
{-# inline toPinnedByteArray #-}
toPinnedByteArray :: Bytes -> ByteArray
toPinnedByteArray b :: Bytes
b@(Bytes ByteArray
arr Int
off Int
len)
  | Int
off Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0, ByteArray -> Int
PM.sizeofByteArray ByteArray
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len, ByteArray -> Bool
PM.isByteArrayPinned ByteArray
arr = ByteArray
arr
  | Bool
otherwise = Bytes -> ByteArray
toPinnedByteArrayClone Bytes
b

-- | Variant of 'toPinnedByteArray' that unconditionally makes a copy of
-- the array backing the sliced 'Bytes' even if the original array
-- could be reused. Prefer 'toPinnedByteArray'.
toPinnedByteArrayClone :: Bytes -> ByteArray
toPinnedByteArrayClone :: Bytes -> ByteArray
toPinnedByteArrayClone (Bytes ByteArray
arr Int
off Int
len) = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
m <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newPinnedByteArray Int
len
  MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
PM.copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
m Int
0 ByteArray
arr Int
off Int
len
  MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
m

-- | /O(n)/ when unpinned, /O(1)/ when pinned. Create a 'ByteString' from
-- a byte sequence. This only copies the byte sequence if it is not pinned.
toByteString :: Bytes -> ByteString
toByteString :: Bytes -> ByteString
toByteString !Bytes
b = Bytes -> ByteString
pinnedToByteString (Bytes -> Bytes
pin Bytes
b)

-- | Convert a pinned 'Bytes' to a 'ByteString'
-- /O(1)/ Precondition: bytes are pinned. Behavior is undefined otherwise.
pinnedToByteString :: Bytes -> ByteString
pinnedToByteString :: Bytes -> ByteString
pinnedToByteString (Bytes y :: ByteArray
y@(PM.ByteArray ByteArray#
x) Int
off Int
len) =
  ForeignPtr Word8 -> Int -> Int -> ByteString
ByteString.PS
    (Addr# -> ForeignPtrContents -> ForeignPtr Word8
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr.ForeignPtr
      (case Ptr Word8 -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
plusPtr (ByteArray -> Ptr Word8
PM.byteArrayContents ByteArray
y) Int
off of {Exts.Ptr Addr#
p -> Addr#
p})
      (MutableByteArray# RealWorld -> ForeignPtrContents
ForeignPtr.PlainPtr (ByteArray# -> MutableByteArray# RealWorld
Exts.unsafeCoerce# ByteArray#
x))
    )
    Int
0 Int
len

-- | /O(n)/ Copy a 'ByteString' to a byte sequence.
fromByteString :: ByteString -> Bytes
fromByteString :: ByteString -> Bytes
fromByteString !ByteString
b = ByteArray -> Int -> Int -> Bytes
Bytes
  ( (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ IO ByteArray -> ST s ByteArray
forall a s. IO a -> ST s a
unsafeIOToST (IO ByteArray -> ST s ByteArray) -> IO ByteArray -> ST s ByteArray
forall a b. (a -> b) -> a -> b
$ do 
      dst :: MutableByteArray RealWorld
dst@(PM.MutableByteArray MutableByteArray# RealWorld
dst# ) <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newByteArray Int
len
      ByteString -> (CString -> IO ()) -> IO ()
forall a. ByteString -> (CString -> IO a) -> IO a
ByteString.unsafeUseAsCString ByteString
b ((CString -> IO ()) -> IO ()) -> (CString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CString
src -> do
        MutablePrimArray (PrimState IO) CChar
-> Int -> CString -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> Ptr a -> Int -> m ()
PM.copyPtrToMutablePrimArray (MutableByteArray# RealWorld -> MutablePrimArray RealWorld CChar
forall s a. MutableByteArray# s -> MutablePrimArray s a
PM.MutablePrimArray MutableByteArray# RealWorld
dst# ) Int
0 CString
src Int
len
      MutableByteArray (PrimState IO) -> IO ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst
  ) Int
0 Int
len
  where
  !len :: Int
len = ByteString -> Int
ByteString.length ByteString
b