mlkem-0.2.0.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
    , create, foldZipWith, iterModify, Block.length, mutableContents
    , Block.new, Block.newPinned, Block.thaw, Block.unsafeCast
    , Block.unsafeFreeze, Block.unsafeThaw
#ifdef ML_KEM_TESTING
    , Block.toList
#endif
    ) where

import Control.Monad.Primitive

import Data.Primitive.PrimArray

import Control.Exception (assert)
import Control.Monad.ST

import Foreign.Ptr

import Base hiding (PrimMonad, PrimState)

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

create :: PrimType ty
       => CountOf ty
       -> (Offset ty -> ty)
       -> Block ty
create (CountOf n) initializer = runST $ do
    mb <- newPrimArray n
    loop mb 0
    unsafeFreezePrimArray mb
  where
    loop !mb i
        | i == n = pure ()
        | otherwise = writePrimArray mb i (initializer $ Offset i) >> loop mb (i + 1)
{-# INLINE create #-}

iterModify :: (PrimType ty, PrimMonad prim)
           => (ty -> ty)
           -> MutableBlock ty (PrimState prim)
           -> prim ()
iterModify f ma = getSizeofMutablePrimArray ma >>= (`loop` 0)
  where
    loop n i
        | i == n = pure ()
        | otherwise = readPrimArray ma i >>= \x -> writePrimArray ma i (f x) >> loop n (i+1)
{-# INLINE iterModify #-}

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

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, PrimType ty) => Block ty -> prim (MutableBlock ty (PrimState prim))
thaw b = thawPrimArray b 0 (sizeofPrimArray b)

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

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

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

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