mldsa-0.1.0.0: src/ByteArrayST.hs

-- |
-- 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
    , fillBytes
    ) where

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

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

import Data.Word

import Foreign.Ptr (Ptr)
import qualified Foreign.Marshal.Utils as S
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

fillBytes :: Ptr a -> Word8 -> Int -> ST s ()
fillBytes a v = unsafeIOToST . S.fillBytes a v