-- |
-- Module      : ByteArrayST
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- Byte array primitives in the @ST@ monad instead of @IO@
--
{-# LANGUAGE RankNTypes #-}
module ByteArrayST
    ( unsafeCreate, withByteArray
    , peek, peekElemOff, pokeElemOff, pokeByteOff
    ) where

import Data.ByteArray (ByteArray, ByteArrayAccess)
import qualified Data.ByteArray as B

import Control.Monad.ST
import Control.Monad.ST.Unsafe

import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable)
import qualified Foreign.Storable as S

unsafeCreate :: ByteArray ba => Int -> (forall s. Ptr p -> ST s ()) -> ba
unsafeCreate sz f = B.unsafeCreate sz (stToIO . f)
{-# INLINE unsafeCreate #-}

withByteArray :: ByteArrayAccess ba => ba -> (Ptr p -> ST s a) -> ST s a
withByteArray b f = unsafeIOToST $ B.withByteArray b (unsafeSTToIO . f)
{-# INLINE withByteArray #-}

peek :: Storable a => Ptr a -> ST s a
peek = unsafeIOToST . S.peek

peekElemOff :: Storable a => Ptr a -> Int -> ST s a
peekElemOff a = unsafeIOToST . S.peekElemOff a

pokeElemOff :: Storable a => Ptr a -> Int -> a -> ST s ()
pokeElemOff a off = unsafeIOToST . S.pokeElemOff a off

pokeByteOff :: Storable a => Ptr a -> Int -> a -> ST s ()
pokeByteOff a off = unsafeIOToST . S.pokeByteOff a off
