{-# 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'
  , fnv1a32
  , fnv1a64
  , toByteString
  , pinnedToByteString
  ) where

import Prelude hiding (length)

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

import qualified Data.ByteString.Internal as ByteString
import qualified Data.Primitive 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
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
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
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
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 = (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

-- | Hash byte sequence with 64-bit variant of FNV-1a.
fnv1a64 :: Bytes -> Word64
fnv1a64 :: Bytes -> Word64
fnv1a64 = (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

-- | 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)

-- | 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
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
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)

-- | /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