--------------------------------------------------------------------------------

{-# language BangPatterns #-}

{-# language ForeignFunctionInterface #-}

{-# language MagicHash #-}

{-# language ScopedTypeVariables #-}

{-# language TypeApplications #-}

{-# language UnboxedTuples #-}

--------------------------------------------------------------------------------

-- | A collection of data types, classes, and functions for interfacing

-- with another programming language using the `Prim` interface instead

-- of the `Storable` interface.

module Data.Primitive.Foreign

(

-- * Prim-Storable methods

sizeOf, alignment

, peek, peekElemOff, peekByteOff

, poke, pokeElemOff, pokeByteOff

-- * Memory allocation

-- ** Local allocation

, alloca, F.allocaBytes, F.allocaBytesAligned

-- ** Dynamic allocation

, malloc, F.mallocBytes

, calloc, F.callocBytes

, realloc, F.reallocBytes

, F.free, F.finalizerFree

-- * Marshalling arrays

-- ** Allocation

, mallocArray, mallocArray0

, allocaArray, allocaArray0

, reallocArray, reallocArray0

, callocArray, callocArray0

-- ** Marshalling

, peekArray, peekArray0

, pokeArray, pokeArray0

-- ** Combined allocation and marshalling

, newArray, newArray0

, withArray, withArray0

, withArrayLen, withArrayLen0

-- ** Copying

, copyArray

, moveArray

-- ** Finding the length

, lengthArray0

-- ** Indexing

, advancePtr

-- * General marshalling utilities

-- ** Combined allocation and marshalling

, with

, new

-- ** Marshalling of Maybe values

, F.maybeNew

, F.maybeWith

, F.maybePeek

-- ** Haskellish interface to memcpy and memmove

, F.copyBytes

, F.moveBytes

-- ** Filling up memory areas with required values

, F.fillBytes

) where

--------------------------------------------------------------------------------

import Control.Monad.Primitive (primitive)

import Data.Coerce (coerce)

import Data.Primitive.PrimArray

import Data.Primitive.Types

import Data.Void (Void)

import GHC.Exts

import GHC.Ptr

import qualified Foreign as F

--------------------------------------------------------------------------------

-- [Section: Foreign.Storable]

-- | Read a value from a memory area regarded as an array of values

-- of the same kind. The first argument specifies the start address

-- of the array and the second the index into the array (the first

-- element of the array has index @0@). The following equality holds,

--

-- > peekElemOff addr idx = fixIO $ \result ->

-- > peek (addr `plusPtr` (idx * sizeOf result))

--

-- Note that this is only a specification, not

-- necessarily the concrete implementation of the

-- function.

peekElemOff :: forall a. Prim a => Ptr a -> Int -> IO a

peekElemOff !ptr = coerce (F.peekElemOff @(PrimStorable a) (coerce ptr))

-- | Write a value to a memory area regarded as an array of

-- values of the same kind. The following equality holds:

--

-- > pokeElemOff addr idx x =

-- > poke (addr `plusPtr` (idx * sizeOf x)) x

pokeElemOff :: forall a. Prim a => Ptr a -> Int -> a -> IO ()

pokeElemOff !ptr !idx a = F.pokeElemOff (coerce ptr) idx (PrimStorable a)

-- | Read a value from a memory location given by a base

-- address and offset. The following equality holds:

--

-- > peekByteOff addr off = peek (addr `plusPtr` off)

peekByteOff :: forall a. Prim a => Ptr Void -> Int -> IO a

peekByteOff !ptr = coerce (F.peekByteOff @(PrimStorable a) ptr)

-- | Write a value to a memory location given by a base

-- address and offset. The following equality holds:

--

-- > pokeByteOff addr off x = poke (addr `plusPtr` off) x

pokeByteOff :: forall a. Prim a => Ptr Void -> Int -> a -> IO ()

pokeByteOff !ptr !idx a = F.pokeByteOff ptr idx (PrimStorable a)

-- | Read a value from the given memory location.

--

-- Note that the peek and poke functions might require properly

-- aligned addresses to function correctly. This is architecture

-- dependent; thus, portable code should ensure that when peeking

-- or poking values of some type @a@, the alignment constraint for

-- @a@, as given by the function 'alignment' is fulfilled.

peek :: forall a. Prim a => Ptr a -> IO a

peek = coerce (F.peek . coerce @(Ptr a) @(Ptr (PrimStorable a)))

-- | Write the given value to the given memory location. Alignment

-- restrictions might apply; see 'peek'.

poke :: forall a. Prim a => Ptr a -> a -> IO ()

poke !ptr a = F.poke (coerce ptr) (PrimStorable a)

--------------------------------------------------------------------------------

-- [Section: Foreign.Marshal.Alloc]

-- |@'alloca' f@ executes the computation @f@, passing as argument

-- a pointer to a temporarily allocated block of memory sufficient to

-- hold values of type @a@.

--

-- The memory is freed when @f@ terminates (either normally or via an

-- exception), so the pointer passed to @f@ must /not/ be used after this.

alloca :: forall a b. Prim a => (Ptr a -> IO b) -> IO b

alloca f = F.alloca (coerce f :: Ptr (PrimStorable a) -> IO b)

-- | Allocate a block of memory that is sufficient to hold values of type

-- @a@. The size of the area allocated is determined by the 'sizeOf'

-- method from the instance of 'Storable' for the appropriate type.

--

-- The memory may be deallocated using 'free' or 'finalizerFree' when

-- no longer required.

--

malloc :: forall a. Prim a => IO (Ptr a)

malloc = F.mallocBytes (sizeOf @a undefined)

-- | Like 'malloc' but memory is filled with bytes of value zero.

--

calloc :: forall a. Prim a => IO (Ptr a)

calloc = F.callocBytes (sizeOf @a undefined)

-- | Resize a memory area that was allocated with 'malloc' or 'mallocBytes'

-- to the size needed to store values of type @b@. The returned pointer

-- may refer to an entirely different memory area, but will be suitably

-- aligned to hold values of type @b@. The contents of the referenced

-- memory area will be the same as of the original pointer up to the

-- minimum of the original size and the size of values of type @b@.

--

-- If the argument to 'realloc' is 'nullPtr', 'realloc' behaves like

-- 'malloc'.

--

realloc :: forall a b. Prim b => Ptr a -> IO (Ptr b)

realloc ptr = coerce (F.realloc ptr :: IO (Ptr (PrimStorable b)))

--------------------------------------------------------------------------------

-- [Section: Foreign.Marshal.Array]

-- | Allocate storage for the given number of elements of a storable type

-- (like 'malloc', but for multiple elements).

mallocArray :: forall a. Prim a => Int -> IO (Ptr a)

mallocArray !idx = coerce (F.mallocArray @(PrimStorable a) idx)

-- | Like 'mallocArray', but add an extra position to hold a special

-- termination element.

mallocArray0 :: forall a. Prim a => Int -> IO (Ptr a)

mallocArray0 !idx = coerce (F.mallocArray0 @(PrimStorable a) idx)

-- | Temporarily allocate space for the given number of elements

-- (like 'alloca', but for multiple elements).

allocaArray :: forall a b. Prim a => Int -> (Ptr a -> IO b) -> IO b

allocaArray !idx f = F.allocaArray idx (coerce f :: Ptr (PrimStorable a) -> IO b)

-- | Like 'allocaArray', but add an extra position to hold a special

-- termination element.

allocaArray0 :: forall a b. Prim a => Int -> (Ptr a -> IO b) -> IO b

allocaArray0 !idx f = F.allocaArray0 idx (coerce f :: Ptr (PrimStorable a) -> IO b)

-- | Adjust the size of an array.

reallocArray :: forall a. Prim a => Ptr a -> Int -> IO (Ptr a)

reallocArray !ptr !idx = coerce (F.reallocArray @(PrimStorable a) (coerce ptr) idx)

-- | Adjust the size of an array, including an extra position for the

-- terminating element.

reallocArray0 :: forall a. Prim a => Ptr a -> Int -> IO (Ptr a)

reallocArray0 !ptr !idx = coerce (F.reallocArray0 @(PrimStorable a) (coerce ptr) idx)

-- | Like 'mallocArray', but allocated memory is filled with bytes of value zero.

callocArray :: forall a. Prim a => Int -> IO (Ptr a)

callocArray !idx = coerce (F.callocArray @(PrimStorable a) idx)

-- | Like 'mallocArray0', but allocated memory is filled with bytes of value zero.

callocArray0 :: forall a. Prim a => Int -> IO (Ptr a)

callocArray0 !idx = coerce (F.callocArray0 @(PrimStorable a) idx)

-- | Convert an array of given length into a Haskell 'PrimArray'.

peekArray :: forall a. Prim a => Int -> Ptr a -> IO (PrimArray a)

peekArray !sz@(I# n#) !(Ptr addr#) = do

marr@(MutablePrimArray ba#) <- newPrimArray sz

primitive $ \s0# ->

case (copyAddrToByteArray# addr# ba# 0# (n# *# (sizeOf# @a undefined)) s0#) of

s1# -> (# s1#, () #)

unsafeFreezePrimArray marr

-- | Convert an array terminated by the given terminator into a Haskell

-- 'PrimArray'.

peekArray0 :: forall a. (Prim a, Eq a) => a -> Ptr a -> IO (PrimArray a)

peekArray0 term !ptr = lengthArray0 term ptr >>= \size ->

peekArray size ptr

-- | Write the 'PrimArray' into memory at the given location.

pokeArray :: forall a. Prim a => Ptr a -> PrimArray a -> IO ()

pokeArray !ptr !arr = flip itraversePrimArray_ arr $ \ix atIx ->

pokeElemOff ptr ix atIx

-- | Write the 'PrimArray' into memory and terminate the elements

-- with a given terminating element.

pokeArray0 :: forall a. Prim a => a -> Ptr a -> PrimArray a -> IO ()

pokeArray0 term !ptr !arr =

let !sz = sizeofPrimArray arr

in flip itraversePrimArray_ arr $ \ix atIx ->

if ix == sz

then pokeElemOff ptr ix term

else pokeElemOff ptr ix atIx

-- | Write a 'PrimArray' into a newly allocated, consecutive

-- sequence of primitive values.

newArray :: forall a. Prim a => PrimArray a -> IO (Ptr a)

newArray !arr = do

ptr <- mallocArray (sizeofPrimArray arr)

pokeArray ptr arr

pure ptr

-- | Write a 'PrimArray' into a newly allocated, consecutive

-- sequence of primitive values, where the end is fixed by

-- the given terminating element.

newArray0 :: forall a. Prim a => a -> PrimArray a -> IO (Ptr a)

newArray0 term !arr = do

ptr <- mallocArray0 (sizeofPrimArray arr)

pokeArray0 term ptr arr

pure ptr

-- | Temporarily store a 'PrimArray' in memory.

withArray :: forall a b. Prim a => PrimArray a -> (Ptr a -> IO b) -> IO b

withArray !arr = withArrayLen arr . const

-- | Like 'withArray', but the action is also passed the size

-- of the 'PrimArray'.

withArrayLen :: forall a b. Prim a => PrimArray a -> (Int -> Ptr a -> IO b) -> IO b

withArrayLen !arr f = allocaArray len $ \ptr -> do

pokeArray ptr arr

f len ptr

where

!len = sizeofPrimArray arr

-- | Like 'withArray', but a terminator indicates where the array ends.

withArray0 :: forall a b. Prim a => a -> PrimArray a -> (Ptr a -> IO b) -> IO b

withArray0 term !arr = withArrayLen0 term arr . const

-- | Like 'withArrayLen', but a terminator indicates where the array ends.

withArrayLen0 :: forall a b. Prim a => a -> PrimArray a -> (Int -> Ptr a -> IO b) -> IO b

withArrayLen0 term !arr f = allocaArray0 len $ \ptr -> do

pokeArray0 term ptr arr

f len ptr

where

!len = sizeofPrimArray arr

-- | Return the number of elements in an array, excluding

-- the terminator

lengthArray0 :: forall a. (Prim a, Eq a)

=> a -- ^ terminating element

-> Ptr a

-> IO Int

lengthArray0 term !ptr = go 0

where

go !ix = peekElemOff ptr ix >>= \val ->

if val == term then pure ix else go (ix + 1)

-- | Advance a pointer into an array by the given number of elements.

advancePtr :: forall a. Prim a => Ptr a -> Int -> Ptr a

advancePtr !ptr !ix = ptr `plusPtr` (ix * sizeOf @a undefined)

-- | Copy the given number of elements from the source array

-- into the destination array; the memory regions /may not/ overlap.

copyArray :: forall a. Prim a

=> Ptr a -- ^ destination array

-> Ptr a -- ^ source array

-> Int -- ^ number of elements to copy

-> IO ()

copyArray !dest !src !size = F.copyBytes dest src (size * sizeOf @a undefined)

-- | Copy the given number of elements from the source array

-- into the destination array; the memory regions /may/ overlap.

moveArray :: forall a. Prim a

=> Ptr a -- ^ destination array

-> Ptr a -- ^ source array

-> Int -- ^ number of elements to copy

-> IO ()

moveArray !dest !src !size = F.moveBytes dest src (size * sizeOf @a undefined)

--------------------------------------------------------------------------------

-- [Section: Foreign.Marshal.Utils]

-- | @'with' val f@ executes the computation @f@, passing as argument

-- a pointer to a temporarily allocated block of memory into which

-- @val@ has been marshalled (the combination of 'alloca' and 'poke').

--

-- The memory is freed when @f@ terminates (either normally or via an

-- exception), so the pointer passed to @f@ must /not/ be used after this.

--

with :: forall a b. Prim a => a -> (Ptr a -> IO b) -> IO b

with val f = F.with (PrimStorable val) (coerce f :: Ptr (PrimStorable a) -> IO b)

-- | Allocate a block of memory and marshal a value into it

-- (the combination of 'malloc' and 'poke').

-- The size of the area allocated is determined by the 'Data.Primitive.Types.sizeOf'

-- method from the instance of 'Storable' for the appropriate type.

--

-- The memory may be deallocated using 'free' or

-- 'finalizerFree' when no longer required.

--

new :: forall a. Prim a => a -> IO (Ptr a)

new val = coerce (F.new (PrimStorable val))

--------------------------------------------------------------------------------