mldsa-0.1.0.0: src/BlockN.hs

-- |
-- Module      : BlockN
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- A secure block with length at type level
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
module BlockN
    ( BlockN, MutableBlockN, BlockN.erase, BlockN.foldl', index, iterModify
    , iterSet, mapEqPrimSize, BlockN.read, BlockN.replicate, runNew, runThaw
    , runFold, BlockN.seq, BlockN.unsafeCast, BlockN.write, BlockN.zipWith
    , zipWithEqPrimSizeR
#ifdef ML_DSA_TESTING
    , create, BlockN.fromList, BlockN.toList
#endif
    ) where

import Control.DeepSeq (NFData(..))
import Control.Monad.ST

import Data.Proxy

import Base
import Block (MutableBlock, blockRead, blockWrite, erase, unsafeCastMut)
import Equality
import Fusion
import Marking (Classified, SecurityMarking)
import SecureBlock (SecureBlock)
import qualified SecureBlock
import Math

newtype BlockN marking (n :: Nat) a = BlockN { unBlockN :: SecureBlock marking a }

#ifdef ML_DSA_TESTING
instance (Classified marking, Eq a, PrimType a) => Eq (BlockN marking n a) where
    BlockN a == BlockN b = SecureBlock.eq a b

instance (Classified marking, PrimType a, Show a) => Show (BlockN marking n a) where
    showsPrec d = SecureBlock.showsPrec d . unBlockN
#endif

instance NFData (BlockN marking n a) where
    rnf = SecureBlock.toNormalForm . unBlockN

instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where
    zero = BlockN.replicate zero
    {-# INLINE zero #-}
    (.+) = BlockN.zipWithEndoL (.+)
    {-# INLINE (.+) #-}
    (.-) = BlockN.zipWithEndoL (.-)
    {-# INLINE (.-) #-}
    neg = BlockN.mapEndo neg
    {-# INLINE neg #-}

newtype MutableBlockN (marking :: SecurityMarking) (n :: Nat) a m = MutableBlockN { unMutableBlockN :: MutableBlock a m }

instance (Classified marking, KnownNat n, PrimType a) => Fusion (BlockN marking n a) where
    type Mut (BlockN marking n a) s = MutableBlockN marking n a s
    newF = new Proxy
    thawF = thaw
    unsafeFreezeF = unsafeFreeze

-- Endomorphism specialization: a different implementation is substituted
-- wherever possible with rewrite rules.  Identical input and output types give
-- a chance for a transformation to be fused on an existing mutable block.

{-# RULES
"mapEndo" [~2] forall f. BlockN.map f = mapEndo f
"zipWithEndoL" [~2] forall f. BlockN.zipWith f = zipWithEndoL f
  #-}

foldl' :: forall marking n a b. (KnownNat n, PrimType a) => (b -> a -> b) -> b -> BlockN marking n a -> b
foldl' f b (BlockN !a) = loop b 0
  where
    !sz = fromIntegral $ natVal (Proxy :: Proxy n)
    loop !acc i
        | i .==# sz = acc
        | otherwise = loop (acc `f` SecureBlock.index a i) (i + 1)
{-# INLINE foldl' #-}

index :: PrimType a => BlockN marking n a -> Offset a -> a
index = SecureBlock.index . unBlockN

replicate :: (Classified marking, KnownNat n, PrimType a) => a -> BlockN marking n a
replicate = create . const

#ifdef ML_DSA_TESTING
fromList :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => [a] -> Maybe (BlockN marking n a)
fromList elems
    | Prelude.length elems /= sz = Nothing
    | otherwise = Just $ runNew (Proxy :: Proxy marking) $ \mb -> go mb 0 elems
  where
    !sz = fromIntegral $ natVal (Proxy :: Proxy n)

    go !mb !i list = case list of
        []     -> return ()
        (x:xs) -> write mb i x >> go mb (i + 1) xs

toList :: PrimType a => BlockN marking n a -> [a]
toList = SecureBlock.toList . unBlockN
#endif

create :: (Classified marking, KnownNat n, PrimType ty)
       => (Offset ty -> ty)
       -> BlockN marking n ty
create initializer = runNew Proxy $ iterSet initializer
{-# INLINE create #-}

map :: (Classified marking, KnownNat n, PrimType a, PrimType b)
    => (a -> b) -> BlockN marking n a -> BlockN marking n b
map f (BlockN a) = BlockN.seq a $
    create $ \(Offset i) -> f (SecureBlock.index a (Offset i))
{-# INLINE [2] map #-}

mapEndo :: (Classified marking, KnownNat n, PrimType a)
        => (a -> a) -> BlockN marking n a -> BlockN marking n a
mapEndo = mapEqPrimSize
{-# INLINE mapEndo #-}

mapEqPrimSize :: (Classified marking, KnownNat n, EqPrimSize a b) => (a -> b) -> BlockN marking n a -> BlockN marking n b
mapEqPrimSize f = runContext . iterMapContext f . thawContext
{-# INLINE mapEqPrimSize #-}

erase :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)
      => MutableBlockN marking n ty (PrimState prim) -> prim ()
erase (MutableBlockN ma) = Block.erase sz ma
  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)

iterModify :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)
           => (ty -> ty)
           -> MutableBlockN marking n ty (PrimState prim)
           -> prim ()
iterModify f = iterModifyIx (\_ x -> f x)
{-# INLINE iterModify #-}

iterModifyIx :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)
             => (Offset ty -> ty -> ty)
             -> MutableBlockN marking n ty (PrimState prim)
             -> prim ()
iterModifyIx f (MutableBlockN !ma) = loop 0
  where
    !sz = fromIntegral $ natVal (Proxy :: Proxy n)

    loop i
        | i .==# sz = pure ()
        | otherwise = blockRead ma i >>= \x -> blockWrite ma i (f i x) >> loop (i + 1)
{-# INLINE iterModifyIx #-}

iterSet :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)
        => (Offset ty -> ty)
        -> MutableBlockN marking n ty (PrimState prim)
        -> prim ()
iterSet f (MutableBlockN !ma) = loop 0
  where
    !sz = fromIntegral $ natVal (Proxy :: Proxy n)

    loop i
        | i .==# sz = pure ()
        | otherwise = blockWrite ma i (f i) >> loop (i + 1)
{-# INLINE iterSet #-}

zipWith :: (Classified mc, KnownNat n, PrimType a, PrimType b, PrimType c)
        => (a -> b -> c) -> BlockN ma n a -> BlockN mb n b -> BlockN mc n c
zipWith f (BlockN a) (BlockN b) =
    BlockN.seq a $ BlockN.seq b $ create $ \(Offset i) ->
        f (SecureBlock.index a (Offset i)) (SecureBlock.index b (Offset i))
{-# INLINE [2] zipWith #-}

zipWithEndoL :: (Classified ma, KnownNat n, PrimType a, PrimType b)
             => (a -> b -> a) -> BlockN ma n a -> BlockN mb n b -> BlockN ma n a
zipWithEndoL = flip . zipWithEqPrimSizeR . flip
{-# INLINE zipWithEndoL #-}

zipWithEqPrimSizeR :: (Classified marking, KnownNat n, PrimType a, EqPrimSize b c)
                   => (a -> b -> c) -> BlockN ma n a -> BlockN marking n b -> BlockN marking n c
zipWithEqPrimSizeR f a b = runContext (seqContext a (iterMapIxContext g (thawContext b)))
  where g = f . index a . Offset
{-# INLINE zipWithEqPrimSizeR #-}

unsafeCast :: BlockN marking n a -> SecureBlock marking b
unsafeCast = SecureBlock.unsafeCast . unBlockN

read :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> prim a
read = blockRead . unMutableBlockN

write :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> a -> prim ()
write = blockWrite . unMutableBlockN

new :: forall proxy marking n a prim. (Classified marking, KnownNat n, PrimMonad prim, PrimType a) => proxy marking -> prim (MutableBlockN marking n a (PrimState prim))
new prx = MutableBlockN <$> SecureBlock.new prx (CountOf sz)
  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)
{-# INLINE new #-}

runThaw :: (Classified marking, KnownNat n, PrimType a) => BlockN marking n a -> (forall s. MutableBlockN marking n a s -> ST s ()) -> BlockN marking n a
runThaw a f = runContext (modifyContext f (thawContext a))
{-# INLINE runThaw #-}

runNew :: (Classified marking, KnownNat n, PrimType a) => proxy marking -> (forall s. MutableBlockN marking n a s -> ST s ()) -> BlockN marking n a
runNew _ f = runContext (modifyContext f newContext)
{-# INLINE runNew #-}

runFold :: (Classified marking, KnownNat n, PrimType a, Foldable t) => BlockN marking n a -> (forall s. b -> MutableBlockN marking n a s -> ST s ()) -> t b -> BlockN marking n a
runFold a f = runContext . foldContext f (thawContext a)
{-# INLINE runFold #-}

seq :: (Classified marking, KnownNat n, PrimType a) => b -> BlockN marking n a -> BlockN marking n a
seq b a = runContext (seqContext b (thawContext a))
{-# INLINE seq #-}

thaw :: (Classified marking, PrimMonad prim) => BlockN marking n a -> prim (MutableBlockN marking n a (PrimState prim))
thaw = fmap MutableBlockN . SecureBlock.thaw . unBlockN

unsafeFreeze :: (Classified marking, PrimMonad prim) => MutableBlockN marking n a (PrimState prim) -> prim (BlockN marking n a)
unsafeFreeze = fmap BlockN . SecureBlock.unsafeFreeze . unMutableBlockN

unsafeMapIx :: forall marking n a b prim. (KnownNat n, EqPrimSize a b, PrimMonad prim) => (Int -> a -> b) -> MutableBlockN marking n a (PrimState prim) -> prim (MutableBlockN marking n b (PrimState prim))
unsafeMapIx f (MutableBlockN !ma) = MutableBlockN . ensureEqPrimSize witness <$> loop 0
  where
    witness = undefined :: a -> b
    !sz = fromIntegral $ natVal (Proxy :: Proxy n)
    loop i
        | i == sz = return (unsafeCastMut ma)
        | otherwise = do
            a <- blockRead ma (Offset i)
            blockWrite (unsafeCastMut ma) (Offset i) (f i a)
            loop (i + 1)
{-# INLINE unsafeMapIx #-}

--

iterMapContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b)
iterMapContext f = iterMapIxContext (\_ x -> f x)
{-# INLINE iterMapContext #-}

iterMapIxContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (Int -> a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b)
iterMapIxContext f = mapContext m
  where m = MapF { mapUpdate = unsafeMapIx f
                 , mapInit = \x -> newF >>= \mb -> Prelude.seq x (iterSet (g x) mb) >> return mb
                 }
        g x (Offset i) = f i (index x (Offset i))
{-# INLINE [1] iterMapIxContext #-}


-- Fusion rules
--
-- "iterMapIxContext/iterMapIxContext" merges element-wise transformations as
-- single operations.  For example @a .+ b .+ c@ becomes a single loop that
-- processes all input blocks in parallel and writes to the destination block.
--
-- "iterMapIxContext/seqContext" moves strictness annotations upstream so that
-- they do not prevent other rules from firing.

{-# RULES
"iterMapIxContext/seqContext" [~1] forall a f c. iterMapIxContext f (seqContext a c) = seqContext a (iterMapIxContext f c)
"iterMapIxContext/iterMapIxContext" [~1] forall f g c. iterMapIxContext f (iterMapIxContext g c) = iterMapIxContext (\i a -> f i (g i a)) c
  #-}