mlkem-0.2.2.0: src/Block.hs

-- |
-- Module      : Block
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- An array of primitive (unlifted) elements.  This module exposes the
-- t'PrimArray' implementation from primitive but through an API similar to
-- basement @Block@.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
module Block
    ( Block, MutableBlock, blockIndex, blockRead, blockWrite
    , foldZipWith, Block.length, mutableContents, getMutableLength
    , Block.new, Block.newPinned, Block.thaw, thawPinned
    , Block.unsafeCast, unsafeCastMut, Block.unsafeFreeze, Block.unsafeThaw
#ifdef ML_KEM_TESTING
    , Block.toList
#endif
    ) where

import Control.Monad.Primitive

import Data.Primitive.ByteArray
import Data.Primitive.PrimArray

import Control.Exception (assert)

import Foreign.Ptr (Ptr)

import Base

type Block = PrimArray
type MutableBlock ty s = MutablePrimArray s ty

blockIndex :: PrimType ty => Block ty -> Offset ty -> ty
blockRead :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> Offset ty -> prim ty
blockWrite :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> Offset ty -> ty -> prim ()
#ifdef ML_KEM_TESTING
blockIndex b off@(Offset i) =
    checkBounds (Block.length b) off $ indexPrimArray b i
blockRead mb off@(Offset i) = getSizeofMutablePrimArray mb >>= \sz ->
    checkBounds (CountOf sz) off $ readPrimArray mb i
blockWrite mb off@(Offset i) a = getSizeofMutablePrimArray mb >>= \sz ->
    checkBounds (CountOf sz) off $ writePrimArray mb i a
#else
blockIndex b (Offset i) = indexPrimArray b i
blockRead mb (Offset i) = readPrimArray mb i
blockWrite mb (Offset i) = writePrimArray mb i
#endif

foldZipWith :: (PrimType a, PrimType b)
            => (c -> a -> b -> c) -> c -> Block a -> Block b -> c
foldZipWith f c a b = assert (sa == sb) $
    loop c 0
  where
    sa = sizeofPrimArray a
    sb = sizeofPrimArray b

    loop !acc i
        | i == sa = acc
        | otherwise = do
            let va = indexPrimArray a i
            let vb = indexPrimArray b i
            loop (f acc va vb) (i + 1)
{-# INLINE foldZipWith #-}

length :: PrimType ty => Block ty -> CountOf ty
length = CountOf . sizeofPrimArray

getMutableLength :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> prim (CountOf ty)
getMutableLength = fmap CountOf . getSizeofMutablePrimArray

mutableContents :: MutableBlock ty s -> Ptr ty
mutableContents = mutablePrimArrayContents -- pinned only

new :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))
new (CountOf n) = newPrimArray n

newPinned :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))
newPinned (CountOf n) = newPinnedPrimArray n

thaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))
thaw (PrimArray !barr) = unsafeSTToPrim $
    -- as optimization, combine both steps in a known monad and avoid
    -- round trip between byte length and element count
    thawByteArray ba 0 (sizeofByteArray ba) >>= \(MutableByteArray mbarr) ->
        return (MutablePrimArray mbarr)
  where ba = ByteArray barr

thawPinned :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))
thawPinned (PrimArray !barr) = unsafeSTToPrim $ do
    let ba = ByteArray barr
        n = sizeofByteArray ba
    mb@(MutableByteArray mbarr) <- newPinnedByteArray n
    copyByteArray mb 0 ba 0 n
    return (MutablePrimArray mbarr)

#ifdef ML_KEM_TESTING
toList :: PrimType ty => Block ty -> [ty]
toList = primArrayToList
#endif

unsafeCast :: Block a -> Block b
unsafeCast (PrimArray b) = PrimArray b

unsafeCastMut :: MutableBlock a m -> MutableBlock b m
unsafeCastMut (MutablePrimArray mb) = MutablePrimArray mb

unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (Block ty)
unsafeFreeze = unsafeFreezePrimArray

unsafeThaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))
unsafeThaw = unsafeThawPrimArray