-- |
-- 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 KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
module BlockN
    ( BlockN, MutableBlockN, create, index, iterModify, BlockN.map
    , BlockN.new, BlockN.read, BlockN.thaw, BlockN.unsafeCast
    , BlockN.unsafeFreeze, BlockN.write, BlockN.zipWith
#ifdef ML_KEM_TESTING
    , BlockN.fromList, BlockN.replicate, BlockN.toList
#endif
    ) where

import Control.DeepSeq (NFData(..))
#ifdef ML_KEM_TESTING
import Control.Monad.ST
#endif

import Data.Proxy

import Base
import Block (MutableBlock, blockRead, blockWrite)
import Marking (Classified, SecurityMarking)
import SecureBlock (SecureBlock)
import qualified SecureBlock
import Math

newtype BlockN marking (n :: Nat) a = BlockN { forall (marking :: SecurityMarking) (n :: Nat) a.
BlockN marking n a -> SecureBlock marking a
unBlockN :: SecureBlock marking a }

#ifdef ML_KEM_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 :: BlockN marking n a -> ()
rnf = SecureBlock marking a -> ()
forall (marking :: SecurityMarking) ty.
SecureBlock marking ty -> ()
SecureBlock.toNormalForm (SecureBlock marking a -> ())
-> (BlockN marking n a -> SecureBlock marking a)
-> BlockN marking n a
-> ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BlockN marking n a -> SecureBlock marking a
forall (marking :: SecurityMarking) (n :: Nat) a.
BlockN marking n a -> SecureBlock marking a
unBlockN

instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where
    zero :: BlockN marking n a
zero = (Offset a -> a) -> BlockN marking n a
forall (marking :: SecurityMarking) (n :: Nat) ty.
(Classified marking, KnownNat n, PrimType ty) =>
(Offset ty -> ty) -> BlockN marking n ty
create (a -> Offset a -> a
forall a b. a -> b -> a
const a
forall a. Add a => a
zero)
    {-# INLINE zero #-}
    .+ :: BlockN marking n a -> BlockN marking n a -> BlockN marking n a
(.+) = (a -> a -> a)
-> BlockN marking n a -> BlockN marking n a -> BlockN marking n a
forall (mc :: SecurityMarking) (n :: Nat) a b c
       (ma :: SecurityMarking) (mb :: SecurityMarking).
(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
BlockN.zipWith a -> a -> a
forall a. Add a => a -> a -> a
(.+)
    {-# INLINE (.+) #-}
    .- :: BlockN marking n a -> BlockN marking n a -> BlockN marking n a
(.-) = (a -> a -> a)
-> BlockN marking n a -> BlockN marking n a -> BlockN marking n a
forall (mc :: SecurityMarking) (n :: Nat) a b c
       (ma :: SecurityMarking) (mb :: SecurityMarking).
(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
BlockN.zipWith a -> a -> a
forall a. Add a => a -> a -> a
(.-)
    {-# INLINE (.-) #-}
    neg :: BlockN marking n a -> BlockN marking n a
neg = (a -> a) -> BlockN marking n a -> BlockN marking n a
forall (marking :: SecurityMarking) (n :: Nat) a b.
(Classified marking, KnownNat n, PrimType a, PrimType b) =>
(a -> b) -> BlockN marking n a -> BlockN marking n b
BlockN.map a -> a
forall a. Add a => a -> a
neg
    {-# INLINE neg #-}

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

index :: PrimType a => BlockN marking n a -> Offset a -> a
index :: forall a (marking :: SecurityMarking) (n :: Nat).
PrimType a =>
BlockN marking n a -> Offset a -> a
index = SecureBlock marking a -> Offset a -> a
forall ty (marking :: SecurityMarking).
PrimType ty =>
SecureBlock marking ty -> Offset ty -> ty
SecureBlock.index (SecureBlock marking a -> Offset a -> a)
-> (BlockN marking n a -> SecureBlock marking a)
-> BlockN marking n a
-> Offset a
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BlockN marking n a -> SecureBlock marking a
forall (marking :: SecurityMarking) (n :: Nat) a.
BlockN marking n a -> SecureBlock marking a
unBlockN

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

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 $ runST $ do
        mb <- new (Proxy :: Proxy marking)
        go mb 0 elems
        unsafeFreeze mb
  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 :: forall marking n ty. (Classified marking, KnownNat n, PrimType ty)
       => (Offset ty -> ty)
       -> BlockN marking n ty
create :: forall (marking :: SecurityMarking) (n :: Nat) ty.
(Classified marking, KnownNat n, PrimType ty) =>
(Offset ty -> ty) -> BlockN marking n ty
create Offset ty -> ty
initializer = SecureBlock marking ty -> BlockN marking n ty
forall (marking :: SecurityMarking) (n :: Nat) a.
SecureBlock marking a -> BlockN marking n a
BlockN (SecureBlock marking ty -> BlockN marking n ty)
-> SecureBlock marking ty -> BlockN marking n ty
forall a b. (a -> b) -> a -> b
$ CountOf ty -> (Offset ty -> ty) -> SecureBlock marking ty
forall ty.
PrimType ty =>
CountOf ty -> (Offset ty -> ty) -> SecureBlock marking ty
forall (marking :: SecurityMarking) ty.
(Classified marking, PrimType ty) =>
CountOf ty -> (Offset ty -> ty) -> SecureBlock marking ty
SecureBlock.create (Int -> CountOf ty
forall ty. Int -> CountOf ty
CountOf Int
sz) Offset ty -> ty
initializer
  where !sz :: Int
sz = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n)
{-# INLINE create #-}

map :: (Classified marking, KnownNat n, PrimType a, PrimType b)
    => (a -> b) -> BlockN marking n a -> BlockN marking n b
map :: forall (marking :: SecurityMarking) (n :: Nat) a b.
(Classified marking, KnownNat n, PrimType a, PrimType b) =>
(a -> b) -> BlockN marking n a -> BlockN marking n b
map a -> b
f (BlockN !SecureBlock marking a
a) = (Offset b -> b) -> BlockN marking n b
forall (marking :: SecurityMarking) (n :: Nat) ty.
(Classified marking, KnownNat n, PrimType ty) =>
(Offset ty -> ty) -> BlockN marking n ty
create ((Offset b -> b) -> BlockN marking n b)
-> (Offset b -> b) -> BlockN marking n b
forall a b. (a -> b) -> a -> b
$ \(Offset Int
i) -> a -> b
f (SecureBlock marking a -> Offset a -> a
forall ty (marking :: SecurityMarking).
PrimType ty =>
SecureBlock marking ty -> Offset ty -> ty
SecureBlock.index SecureBlock marking a
a (Int -> Offset a
forall ty. Int -> Offset ty
Offset Int
i))
{-# INLINE map #-}

iterModify :: (PrimType ty, PrimMonad prim)
           => (ty -> ty)
           -> MutableBlockN marking n ty (PrimState prim)
           -> prim ()
iterModify :: forall ty (prim :: * -> *) (marking :: SecurityMarking) (n :: Nat).
(PrimType ty, PrimMonad prim) =>
(ty -> ty)
-> MutableBlockN marking n ty (PrimState prim) -> prim ()
iterModify ty -> ty
f = (ty -> ty) -> MutableBlock ty (PrimState prim) -> prim ()
forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
(ty -> ty) -> MutableBlock ty (PrimState prim) -> prim ()
SecureBlock.iterModify ty -> ty
f (MutableBlock ty (PrimState prim) -> prim ())
-> (MutableBlockN marking n ty (PrimState prim)
    -> MutableBlock ty (PrimState prim))
-> MutableBlockN marking n ty (PrimState prim)
-> prim ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutableBlockN marking n ty (PrimState prim)
-> MutableBlock ty (PrimState prim)
forall (marking :: SecurityMarking) (n :: Nat) a m.
MutableBlockN marking n a m -> MutableBlock a m
unMutableBlockN
{-# INLINE iterModify #-}

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 :: forall (mc :: SecurityMarking) (n :: Nat) a b c
       (ma :: SecurityMarking) (mb :: SecurityMarking).
(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 a -> b -> c
f (BlockN !SecureBlock ma a
a) (BlockN !SecureBlock mb b
b) =
    (Offset c -> c) -> BlockN mc n c
forall (marking :: SecurityMarking) (n :: Nat) ty.
(Classified marking, KnownNat n, PrimType ty) =>
(Offset ty -> ty) -> BlockN marking n ty
create ((Offset c -> c) -> BlockN mc n c)
-> (Offset c -> c) -> BlockN mc n c
forall a b. (a -> b) -> a -> b
$ \(Offset Int
i) ->
        a -> b -> c
f (SecureBlock ma a -> Offset a -> a
forall ty (marking :: SecurityMarking).
PrimType ty =>
SecureBlock marking ty -> Offset ty -> ty
SecureBlock.index SecureBlock ma a
a (Int -> Offset a
forall ty. Int -> Offset ty
Offset Int
i)) (SecureBlock mb b -> Offset b -> b
forall ty (marking :: SecurityMarking).
PrimType ty =>
SecureBlock marking ty -> Offset ty -> ty
SecureBlock.index SecureBlock mb b
b (Int -> Offset b
forall ty. Int -> Offset ty
Offset Int
i))
{-# INLINE zipWith #-}

unsafeCast :: BlockN marking n a -> SecureBlock marking b
unsafeCast :: forall (marking :: SecurityMarking) (n :: Nat) a b.
BlockN marking n a -> SecureBlock marking b
unsafeCast = SecureBlock marking a -> SecureBlock marking b
forall (marking :: SecurityMarking) a b.
SecureBlock marking a -> SecureBlock marking b
SecureBlock.unsafeCast (SecureBlock marking a -> SecureBlock marking b)
-> (BlockN marking n a -> SecureBlock marking a)
-> BlockN marking n a
-> SecureBlock marking b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BlockN marking n a -> SecureBlock marking a
forall (marking :: SecurityMarking) (n :: Nat) a.
BlockN marking n a -> SecureBlock marking a
unBlockN

read :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> prim a
read :: forall (prim :: * -> *) a (marking :: SecurityMarking) (n :: Nat).
(PrimMonad prim, PrimType a) =>
MutableBlockN marking n a (PrimState prim) -> Offset a -> prim a
read = MutableBlock a (PrimState prim) -> Offset a -> prim a
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MutableBlock ty (PrimState prim) -> Offset ty -> prim ty
blockRead (MutableBlock a (PrimState prim) -> Offset a -> prim a)
-> (MutableBlockN marking n a (PrimState prim)
    -> MutableBlock a (PrimState prim))
-> MutableBlockN marking n a (PrimState prim)
-> Offset a
-> prim a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutableBlockN marking n a (PrimState prim)
-> MutableBlock a (PrimState prim)
forall (marking :: SecurityMarking) (n :: Nat) a m.
MutableBlockN marking n a m -> MutableBlock a m
unMutableBlockN

write :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> a -> prim ()
write :: forall (prim :: * -> *) a (marking :: SecurityMarking) (n :: Nat).
(PrimMonad prim, PrimType a) =>
MutableBlockN marking n a (PrimState prim)
-> Offset a -> a -> prim ()
write = MutableBlock a (PrimState prim) -> Offset a -> a -> prim ()
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MutableBlock ty (PrimState prim) -> Offset ty -> ty -> prim ()
blockWrite (MutableBlock a (PrimState prim) -> Offset a -> a -> prim ())
-> (MutableBlockN marking n a (PrimState prim)
    -> MutableBlock a (PrimState prim))
-> MutableBlockN marking n a (PrimState prim)
-> Offset a
-> a
-> prim ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutableBlockN marking n a (PrimState prim)
-> MutableBlock a (PrimState prim)
forall (marking :: SecurityMarking) (n :: Nat) a m.
MutableBlockN marking n a m -> MutableBlock a m
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 :: forall (proxy :: SecurityMarking -> *) (marking :: SecurityMarking)
       (n :: Nat) a (prim :: * -> *).
(Classified marking, KnownNat n, PrimMonad prim, PrimType a) =>
proxy marking -> prim (MutableBlockN marking n a (PrimState prim))
new proxy marking
prx = MutableBlock a (PrimState prim)
-> MutableBlockN marking n a (PrimState prim)
forall (marking :: SecurityMarking) (n :: Nat) a m.
MutableBlock a m -> MutableBlockN marking n a m
MutableBlockN (MutableBlock a (PrimState prim)
 -> MutableBlockN marking n a (PrimState prim))
-> prim (MutableBlock a (PrimState prim))
-> prim (MutableBlockN marking n a (PrimState prim))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> proxy marking
-> CountOf a -> prim (MutableBlock a (PrimState prim))
forall ty (prim :: * -> *) (proxy :: SecurityMarking -> *).
(PrimType ty, PrimMonad prim) =>
proxy marking
-> CountOf ty -> prim (MutableBlock ty (PrimState prim))
forall (marking :: SecurityMarking) ty (prim :: * -> *)
       (proxy :: SecurityMarking -> *).
(Classified marking, PrimType ty, PrimMonad prim) =>
proxy marking
-> CountOf ty -> prim (MutableBlock ty (PrimState prim))
SecureBlock.new proxy marking
prx (Int -> CountOf a
forall ty. Int -> CountOf ty
CountOf Int
sz)
  where !sz :: Int
sz = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n)
{-# INLINE new #-}

thaw :: (Classified marking, PrimMonad prim, PrimType a) => BlockN marking n a -> prim (MutableBlockN marking n a (PrimState prim))
thaw :: forall (marking :: SecurityMarking) (prim :: * -> *) a (n :: Nat).
(Classified marking, PrimMonad prim, PrimType a) =>
BlockN marking n a
-> prim (MutableBlockN marking n a (PrimState prim))
thaw = (MutableBlock a (PrimState prim)
 -> MutableBlockN marking n a (PrimState prim))
-> prim (MutableBlock a (PrimState prim))
-> prim (MutableBlockN marking n a (PrimState prim))
forall a b. (a -> b) -> prim a -> prim b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap MutableBlock a (PrimState prim)
-> MutableBlockN marking n a (PrimState prim)
forall (marking :: SecurityMarking) (n :: Nat) a m.
MutableBlock a m -> MutableBlockN marking n a m
MutableBlockN (prim (MutableBlock a (PrimState prim))
 -> prim (MutableBlockN marking n a (PrimState prim)))
-> (BlockN marking n a -> prim (MutableBlock a (PrimState prim)))
-> BlockN marking n a
-> prim (MutableBlockN marking n a (PrimState prim))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SecureBlock marking a -> prim (MutableBlock a (PrimState prim))
forall ty (m :: * -> *).
(PrimType ty, PrimMonad m) =>
SecureBlock marking ty -> m (MutableBlock ty (PrimState m))
forall (marking :: SecurityMarking) ty (m :: * -> *).
(Classified marking, PrimType ty, PrimMonad m) =>
SecureBlock marking ty -> m (MutableBlock ty (PrimState m))
SecureBlock.thaw (SecureBlock marking a -> prim (MutableBlock a (PrimState prim)))
-> (BlockN marking n a -> SecureBlock marking a)
-> BlockN marking n a
-> prim (MutableBlock a (PrimState prim))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BlockN marking n a -> SecureBlock marking a
forall (marking :: SecurityMarking) (n :: Nat) a.
BlockN marking n a -> SecureBlock marking a
unBlockN

unsafeFreeze :: (Classified marking, PrimMonad prim) => MutableBlockN marking n a (PrimState prim) -> prim (BlockN marking n a)
unsafeFreeze :: forall (marking :: SecurityMarking) (prim :: * -> *) (n :: Nat) a.
(Classified marking, PrimMonad prim) =>
MutableBlockN marking n a (PrimState prim)
-> prim (BlockN marking n a)
unsafeFreeze = (SecureBlock marking a -> BlockN marking n a)
-> prim (SecureBlock marking a) -> prim (BlockN marking n a)
forall a b. (a -> b) -> prim a -> prim b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SecureBlock marking a -> BlockN marking n a
forall (marking :: SecurityMarking) (n :: Nat) a.
SecureBlock marking a -> BlockN marking n a
BlockN (prim (SecureBlock marking a) -> prim (BlockN marking n a))
-> (MutableBlockN marking n a (PrimState prim)
    -> prim (SecureBlock marking a))
-> MutableBlockN marking n a (PrimState prim)
-> prim (BlockN marking n a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutableBlock a (PrimState prim) -> prim (SecureBlock marking a)
forall (marking :: SecurityMarking) (prim :: * -> *) ty.
(Classified marking, PrimMonad prim) =>
MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty)
forall (prim :: * -> *) ty.
PrimMonad prim =>
MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty)
SecureBlock.unsafeFreeze (MutableBlock a (PrimState prim) -> prim (SecureBlock marking a))
-> (MutableBlockN marking n a (PrimState prim)
    -> MutableBlock a (PrimState prim))
-> MutableBlockN marking n a (PrimState prim)
-> prim (SecureBlock marking a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MutableBlockN marking n a (PrimState prim)
-> MutableBlock a (PrimState prim)
forall (marking :: SecurityMarking) (n :: Nat) a m.
MutableBlockN marking n a m -> MutableBlock a m
unMutableBlockN