mlkem-0.2.0.0: src/ScrubbedBlock.hs

-- |
-- Module      : ScrubbedBlock
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- A block that is always pinned in memory and automatically erased by a
-- finalizer when not referenced anymore.  Same pattern as ScrubbedBytes from
-- package memory but for blocks.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module ScrubbedBlock
    ( ScrubbedBlock, create, foldZipWith, ScrubbedBlock.length
    , new, thaw, unsafeFreeze
    ) where

import Data.Primitive.PrimArray as Block

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

import Data.Word

import Base
import Block (Block, MutableBlock, blockWrite)
import qualified Block

import GHC.Base (IO(IO), Int(I#), setByteArray#)
import GHC.Exts (mkWeak#)

newtype ScrubbedBlock ty = ScrubbedBlock (Block ty)
    deriving (Eq, Show)

create :: PrimType ty
       => CountOf ty
       -> (Offset ty -> ty)
       -> ScrubbedBlock ty
create n initializer = runST $ do
    mb <- new n
    loop mb 0
    unsafeFreeze mb
  where
    loop !mb i
        | i .==# n = pure ()
        | otherwise = blockWrite mb i (initializer i) >> loop mb (i + 1)
{-# INLINE create #-}

foldZipWith :: (PrimType a, PrimType b)
            => (c -> a -> b -> c) -> c -> ScrubbedBlock a -> ScrubbedBlock b -> c
foldZipWith f c (ScrubbedBlock a) (ScrubbedBlock b) =
    Block.foldZipWith f c a b
{-# INLINE foldZipWith #-}

length :: PrimType ty => ScrubbedBlock ty -> CountOf ty
length (ScrubbedBlock b) = Block.length b

new :: (PrimType ty, PrimMonad prim) => CountOf ty -> prim (MutableBlock ty (PrimState prim))
new = Block.newPinned  -- always pinned

thaw :: (PrimType ty, PrimMonad m) => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m))
thaw (ScrubbedBlock b) = do
    let !n@(CountOf sz) = Block.length b
    mb <- new n
    Block.copyPrimArray mb 0 b 0 sz
    return mb
{-# INLINE thaw #-}

unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty)
unsafeFreeze mb = Block.unsafeFreeze mb >>= scrubbed


{- internal -}

assertPinned :: Block ty -> a -> a
assertPinned mb = assert (Block.isPrimArrayPinned mb)

scrubbed :: PrimMonad prim => Block ty -> prim (ScrubbedBlock ty)
scrubbed b = assertPinned b $ unsafePrimFromIO $
    scheduleBlockScrubbing b >> return (ScrubbedBlock b)

scheduleBlockScrubbing :: Block ty -> IO ()
scheduleBlockScrubbing b = addBlockFinalizer b (scrub $ Block.unsafeCast b)
{-# NOINLINE scheduleBlockScrubbing #-}

scrub :: Block Word8 -> IO ()
scrub b = Block.unsafeThaw b >>= erase len
  where CountOf len = Block.length b

addBlockFinalizer :: Block ty -> IO () -> IO ()
addBlockFinalizer (Block.PrimArray barr) (IO finalizer) = IO $ \s ->
   case mkWeak# barr () finalizer s of { (# s1, _ #) -> (# s1, () #) }

erase :: Int -> MutableBlock Word8 RealWorld -> IO ()
erase (I# len) (Block.MutablePrimArray mbarr) = IO $ \s1 ->
    case setByteArray# mbarr 0# len 0# s1 of
        s2 -> (# s2, () #)