-- |
-- 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 :: forall ty. PrimType ty => Block ty -> Offset ty -> ty
blockIndex Block ty
b (Offset Int
i) = Block ty -> Int -> ty
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray Block ty
b Int
i
blockRead :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MutableBlock ty (PrimState prim) -> Offset ty -> prim ty
blockRead MutableBlock ty (PrimState prim)
mb (Offset Int
i) = MutableBlock ty (PrimState prim) -> Int -> prim ty
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutableBlock ty (PrimState prim)
mb Int
i
blockWrite :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MutableBlock ty (PrimState prim) -> Offset ty -> ty -> prim ()
blockWrite MutableBlock ty (PrimState prim)
mb (Offset Int
i) = MutableBlock ty (PrimState prim) -> Int -> ty -> prim ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutableBlock ty (PrimState prim)
mb Int
i
#endif

create :: PrimType ty
       => CountOf ty
       -> (Offset ty -> ty)
       -> Block ty
create :: forall ty.
PrimType ty =>
CountOf ty -> (Offset ty -> ty) -> Block ty
create (CountOf Int
n) Offset ty -> ty
initializer = (forall s. ST s (Block ty)) -> Block ty
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Block ty)) -> Block ty)
-> (forall s. ST s (Block ty)) -> Block ty
forall a b. (a -> b) -> a -> b
$ do
    mb <- Int -> ST s (MutablePrimArray (PrimState (ST s)) ty)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
n
    loop mb 0
    unsafeFreezePrimArray mb
  where
    loop :: MutablePrimArray (PrimState f) ty -> Int -> f ()
loop !MutablePrimArray (PrimState f) ty
mb Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = MutablePrimArray (PrimState f) ty -> Int -> ty -> f ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState f) ty
mb Int
i (Offset ty -> ty
initializer (Offset ty -> ty) -> Offset ty -> ty
forall a b. (a -> b) -> a -> b
$ Int -> Offset ty
forall ty. Int -> Offset ty
Offset Int
i) f () -> f () -> f ()
forall a b. f a -> f b -> f b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MutablePrimArray (PrimState f) ty -> Int -> f ()
loop MutablePrimArray (PrimState f) ty
mb (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE create #-}

iterModify :: (PrimType ty, PrimMonad prim)
           => (ty -> ty)
           -> MutableBlock ty (PrimState prim)
           -> prim ()
iterModify :: forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
(ty -> ty) -> MutableBlock ty (PrimState prim) -> prim ()
iterModify ty -> ty
f MutableBlock ty (PrimState prim)
ma = MutableBlock ty (PrimState prim) -> prim Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutableBlock ty (PrimState prim)
ma prim Int -> (Int -> prim ()) -> prim ()
forall a b. prim a -> (a -> prim b) -> prim b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> Int -> prim ()
forall {f :: * -> *}.
(PrimState f ~ PrimState prim, PrimMonad f) =>
Int -> Int -> f ()
`loop` Int
0)
  where
    loop :: Int -> Int -> f ()
loop Int
n Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = MutablePrimArray (PrimState f) ty -> Int -> f ty
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutableBlock ty (PrimState prim)
MutablePrimArray (PrimState f) ty
ma Int
i f ty -> (ty -> f ()) -> f ()
forall a b. f a -> (a -> f b) -> f b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ty
x -> MutablePrimArray (PrimState f) ty -> Int -> ty -> f ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutableBlock ty (PrimState prim)
MutablePrimArray (PrimState f) ty
ma Int
i (ty -> ty
f ty
x) f () -> f () -> f ()
forall a b. f a -> f b -> f b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> f ()
loop Int
n (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
{-# INLINE iterModify #-}

foldZipWith :: (PrimType a, PrimType b)
            => (c -> a -> b -> c) -> c -> Block a -> Block b -> c
foldZipWith :: forall a b c.
(PrimType a, PrimType b) =>
(c -> a -> b -> c) -> c -> Block a -> Block b -> c
foldZipWith c -> a -> b -> c
f c
c Block a
a Block b
b = Bool -> c -> c
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
sa Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sb) (c -> c) -> c -> c
forall a b. (a -> b) -> a -> b
$
    c -> Int -> c
loop c
c Int
0
  where
    sa :: Int
sa = Block a -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray Block a
a
    sb :: Int
sb = Block b -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray Block b
b

    loop :: c -> Int -> c
loop !c
acc Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sa = c
acc
        | Bool
otherwise = do
            let va :: a
va = Block a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray Block a
a Int
i
            let vb :: b
vb = Block b -> Int -> b
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray Block b
b Int
i
            c -> Int -> c
loop (c -> a -> b -> c
f c
acc a
va b
vb) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE foldZipWith #-}

length :: PrimType ty => Block ty -> CountOf ty
length :: forall ty. PrimType ty => Block ty -> CountOf ty
length = Int -> CountOf ty
forall ty. Int -> CountOf ty
CountOf (Int -> CountOf ty) -> (Block ty -> Int) -> Block ty -> CountOf ty
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Block ty -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray

mutableContents :: MutableBlock ty s -> Ptr ty
mutableContents :: forall ty s. MutableBlock ty s -> Ptr ty
mutableContents = MutablePrimArray s ty -> Ptr ty
forall s a. MutablePrimArray s a -> Ptr a
mutablePrimArrayContents -- pinned only

new :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))
new :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MutableBlock ty (PrimState prim))
new (CountOf Int
n) = Int -> prim (MutablePrimArray (PrimState prim) ty)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
n

newPinned :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))
newPinned :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MutableBlock ty (PrimState prim))
newPinned (CountOf Int
n) = Int -> prim (MutablePrimArray (PrimState prim) ty)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPinnedPrimArray Int
n

thaw :: (PrimMonad prim, PrimType ty) => Block ty -> prim (MutableBlock ty (PrimState prim))
thaw :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
Block ty -> prim (MutableBlock ty (PrimState prim))
thaw Block ty
b = Block ty
-> Int -> Int -> prim (MutablePrimArray (PrimState prim) ty)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimArray a -> Int -> Int -> m (MutablePrimArray (PrimState m) a)
thawPrimArray Block ty
b Int
0 (Block ty -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray Block ty
b)

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

unsafeCast :: Block a -> Block b
unsafeCast :: forall a b. Block a -> Block b
unsafeCast (PrimArray ByteArray#
b) = ByteArray# -> PrimArray b
forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
b

unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (Block ty)
unsafeFreeze :: forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreeze = MutablePrimArray (PrimState prim) ty -> prim (PrimArray ty)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray

unsafeThaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))
unsafeThaw :: forall (prim :: * -> *) ty.
PrimMonad prim =>
Block ty -> prim (MutableBlock ty (PrimState prim))
unsafeThaw = PrimArray ty -> prim (MutablePrimArray (PrimState prim) ty)
forall (prim :: * -> *) ty.
PrimMonad prim =>
Block ty -> prim (MutableBlock ty (PrimState prim))
unsafeThawPrimArray