mlkem-0.2.2.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.
--
-- A complication here is that we distinguish between mutable and immutable
-- values.  And for resiliency against asynchronous exceptions, we need to
-- schedule block scrubbing with a finalizer right at the beginning when the
-- block is still in mutable form.  Fortunately, for the perspective of the GC,
-- ByteArray# and MutableByteArray# are really the same heap object in disguise
-- and unsafeFreezeByteArray# is a true no-op.  So the finalizer set on the
-- initial MutableByteArray# value gets transferred transparently to the final
-- ByteArray# form.
--
-- See GHC note [primOpEffect of unsafe freezes and thaws]
--
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module ScrubbedBlock
    ( ScrubbedBlock, 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 Unsafe.Coerce

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

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

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

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 n = Block.newPinned n >>= scrubbed  -- always pinned

thaw :: PrimMonad m => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m))
thaw (ScrubbedBlock b) = Block.thawPinned b >>= scrubbed  -- always pinned

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


{- internal -}

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

checkPinned :: Block ty -> ScrubbedBlock ty
checkPinned b = assertPinned b (ScrubbedBlock b)

scrubbed :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (MutableBlock ty (PrimState prim))
scrubbed b = unsafePrimFromIO (scheduleBlockScrubbing b >> return b)

wakeUpAfterInception :: MutableBlock ty s -> MutableBlock ty RealWorld
wakeUpAfterInception = unsafeCoerce  -- sometimes disappointing

scheduleBlockScrubbing :: MutableBlock ty s -> IO ()
scheduleBlockScrubbing b = addBlockFinalizer b (scrub $ Block.unsafeCastMut b')
  where b' = wakeUpAfterInception b
{-# NOINLINE scheduleBlockScrubbing #-}

scrub :: MutableBlock Word8 RealWorld -> IO ()
scrub b = Block.getMutableLength b >>= \(CountOf len) -> erase len b

addBlockFinalizer :: MutableBlock ty s -> IO () -> IO ()
addBlockFinalizer (Block.MutablePrimArray mbarr) (IO finalizer) = IO $ \s ->
   case mkWeak# mbarr () 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, () #)