mlkem 0.1.1.0 → 0.2.0.0
raw patch · 18 files changed
+281/−175 lines, 18 filesdep +primitivedep +ramdep −basementdep ~cryptonPVP ok
version bump matches the API change (PVP)
Dependencies added: primitive, ram
Dependencies removed: basement
Dependency ranges changed: crypton
API changes (from Hackage documentation)
Files
- CHANGELOG.md +7/−0
- README.md +8/−0
- mlkem.cabal +21/−15
- src/Auxiliary.hs +40/−26
- src/Base.hs +53/−0
- src/Block.hs +69/−32
- src/BlockN.hs +7/−10
- src/Builder.hs +2/−3
- src/Crypto.hs +3/−2
- src/Crypto/PubKey/ML_KEM.hs +4/−3
- src/Internal.hs +1/−2
- src/K_PKE.hs +1/−3
- src/Machine.hs +1/−1
- src/Marking.hs +6/−12
- src/Matrix.hs +1/−3
- src/ScrubbedBlock.hs +12/−34
- src/SecureBlock.hs +1/−6
- src/Vector.hs +44/−23
CHANGELOG.md view
@@ -1,5 +1,12 @@ # Changelog for `mlkem` +## 0.2.0.0 - 2026-03-26++* Flag `use_crypton` is now enabled by default. It requires crypton >= 1.1.1+ and changes the dependency to `ram` instead of `memory`++* Dependency to `basement` is replaced with `primitive`+ ## 0.1.1.0 - 2025-11-30 * Fixed encoding and decoding on big-endian architectures
README.md view
@@ -29,9 +29,17 @@ the program exits. Also, depending on optimizations applied, lambdas may capture variables and move them to the heap. This could theoretically include machine words containing secret information that would not then be destroyed.+Cautious users can run the benchmarks with info-table profiling and verify that+closures containing non pointers capture only non-secret variables like loop+indices or algorithm parameters. Best performance is obtained with the LLVM code generator. On ARM, define macro `__ARM_FEATURE_UNALIGNED` if unaligned access is supported by the target.++Randomness is provided either from explicit inputs or through a user-selected+instance of the `MonadRandom` type class from `crypton`. A good implementation+would combine multiple sources of entropy, reseed periodically, and protect+its internal state in memory. ## Testing
mlkem.cabal view
@@ -1,11 +1,11 @@ cabal-version: 2.2 --- This file has been generated from package.yaml by hpack version 0.38.1.+-- This file has been generated from package.yaml by hpack version 0.39.1. -- -- see: https://github.com/sol/hpack name: mlkem-version: 0.1.1.0+version: 0.2.0.0 synopsis: Module-Lattice-based Key-Encapsulation Mechanism description: Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM) implemented in Haskell.@@ -31,13 +31,14 @@ flag use_crypton description: Use crypton instead of cryptonite manual: True- default: False+ default: True library exposed-modules: Crypto.PubKey.ML_KEM other-modules: Auxiliary+ Base Block BlockN Builder@@ -62,16 +63,17 @@ ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -Wno-unticked-promoted-constructors -O2 build-depends: base >=4.7 && <5- , basement >=0.0.8 , deepseq- , memory+ , primitive >=0.7.2 default-language: Haskell2010 if flag(use_crypton) build-depends:- crypton+ crypton >=1.1.1+ , ram else build-depends: cryptonite >=0.26+ , memory test-suite mlkem-test type: exitcode-stdio-1.0@@ -90,12 +92,11 @@ build-depends: aeson , base >=4.7 && <5- , basement >=0.0.8 , bytestring , deepseq , directory- , memory , mlkem+ , primitive >=0.7.2 , process , tasty , tasty-hunit@@ -105,16 +106,19 @@ default-language: Haskell2010 if flag(use_crypton) build-depends:- crypton+ crypton >=1.1.1+ , ram else build-depends: cryptonite >=0.26+ , memory test-suite mlkem-test-full type: exitcode-stdio-1.0 main-is: Tests.hs other-modules: Auxiliary+ Base Block BlockN Builder@@ -147,11 +151,10 @@ build-depends: aeson , base >=4.7 && <5- , basement >=0.0.8 , bytestring , deepseq , directory- , memory+ , primitive >=0.7.2 , process , tasty , tasty-hunit@@ -161,10 +164,12 @@ default-language: Haskell2010 if flag(use_crypton) build-depends:- crypton+ crypton >=1.1.1+ , ram else build-depends: cryptonite >=0.26+ , memory benchmark mlkem-bench type: exitcode-stdio-1.0@@ -178,15 +183,16 @@ ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -Wno-unticked-promoted-constructors -threaded -rtsopts -with-rtsopts=-N -with-rtsopts=-A48m build-depends: base >=4.7 && <5- , basement >=0.0.8 , criterion , deepseq- , memory , mlkem+ , primitive >=0.7.2 default-language: Haskell2010 if flag(use_crypton) build-depends:- crypton+ crypton >=1.1.1+ , ram else build-depends: cryptonite >=0.26+ , memory
src/Auxiliary.hs view
@@ -9,9 +9,12 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UnboxedTuples #-} module Auxiliary ( Zq, Rq, Tq, (..+), (..-) , ntt, nttInv, rcompress, rdecompress@@ -23,15 +26,13 @@ #endif ) where -import Basement.NormalForm-import Basement.PrimType-import Basement.Types.OffsetSize- import Crypto.Hash.Algorithms import Data.ByteArray (ByteArrayAccess, Bytes, View) import qualified Data.ByteArray as B +import Data.Primitive.Types (Prim(..))+ import Control.DeepSeq (NFData(..)) import Control.Monad import Control.Monad.ST@@ -47,6 +48,7 @@ import Unsafe.Coerce +import Base import Block (blockIndex) import BlockN (BlockN, MutableBlockN) import Builder (Builder)@@ -112,24 +114,37 @@ deriving Eq #endif -instance PrimType Zq where- type PrimSize Zq = 2- primSizeInBytes _ = primSizeInBytes (Proxy :: Proxy Word16)- {-# INLINE primSizeInBytes #-}- primShiftToBytes _ = primShiftToBytes (Proxy :: Proxy Word16)- {-# INLINE primShiftToBytes #-}- primBaUIndex ba (Offset i) = Zq (primBaUIndex ba (Offset i))- {-# INLINE primBaUIndex #-}- primMbaURead mba (Offset i) = Zq <$> primMbaURead mba (Offset i)- {-# INLINE primMbaURead #-}- primMbaUWrite mba (Offset i) (Zq a) = primMbaUWrite mba (Offset i) a- {-# INLINE primMbaUWrite #-}- primAddrIndex addr (Offset i) = Zq (primAddrIndex addr (Offset i))- {-# INLINE primAddrIndex #-}- primAddrRead addr (Offset i) = Zq <$> primAddrRead addr (Offset i)- {-# INLINE primAddrRead #-}- primAddrWrite addr (Offset i) (Zq a) = primAddrWrite addr (Offset i) a- {-# INLINE primAddrWrite #-}+instance Prim Zq where+ sizeOf# (Zq a) = sizeOf# a+ {-# INLINE sizeOf# #-}+ alignment# (Zq a) = alignment# a+ {-# INLINE alignment# #-}+#if MIN_VERSION_primitive(0,9,0)+ sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word16)+ {-# INLINE sizeOfType# #-}+ alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word16)+ {-# INLINE alignmentOfType# #-}+#endif+ indexByteArray# ba i = Zq (indexByteArray# ba i)+ {-# INLINE indexByteArray# #-}+ readByteArray# mba i s =+ case readByteArray# mba i s of+ (# s', a #) -> (# s', Zq a #)+ {-# INLINE readByteArray# #-}+ writeByteArray# mba i (Zq a) = writeByteArray# mba i a+ {-# INLINE writeByteArray# #-}+ setByteArray# mba i len (Zq a) = setByteArray# mba i len a+ {-# INLINE setByteArray# #-}+ indexOffAddr# addr i = Zq (indexOffAddr# addr i)+ {-# INLINE indexOffAddr# #-}+ readOffAddr# addr i s =+ case readOffAddr# addr i s of+ (# s', a #) -> (# s', Zq a #)+ {-# INLINE readOffAddr# #-}+ writeOffAddr# addr i (Zq a) = writeOffAddr# addr i a+ {-# INLINE writeOffAddr# #-}+ setOffAddr# addr i len (Zq a) = setOffAddr# addr i len a+ {-# INLINE setOffAddr# #-} instance Add Zq where zero = Zq 0@@ -196,7 +211,9 @@ newtype Tq marking = Tq (BlockN marking N Zq) #ifdef ML_KEM_TESTING- deriving (Eq, Show)+ deriving (Eq, Show, NFData)+#else+ deriving NFData #endif instance Classified marking => Add (Tq marking) where@@ -208,9 +225,6 @@ {-# SPECIALIZE instance Add (Tq Pub) #-} instance Leak Tq--instance NFData (Tq marking) where- rnf (Tq a) = toNormalForm a instance BiMul (Tq Pub) (Tq Sec) where (..*) = multiplyNTTs
+ src/Base.hs view
@@ -0,0 +1,53 @@+-- |+-- Module : Base+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Utilities similar to ones provided in basement but adapted for primitive+--+{-# LANGUAGE CPP #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+module Base+ ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimState+ , natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#)+#ifdef ML_KEM_TESTING+ , checkBounds+#endif+ ) where++import Control.Monad.Primitive++import Data.Primitive.Types++import Data.Bits++import GHC.TypeLits++type PrimType = Prim++newtype CountOf ty = CountOf Int+ deriving (Show, Eq, Ord)++newtype Offset ty = Offset Int+ deriving (Show, Eq, Ord, Num)++offsetShiftL :: Int -> Offset ty -> Offset ty2+offsetShiftL n (Offset o) = Offset (o `unsafeShiftL` n)++offsetShiftR :: Int -> Offset ty -> Offset ty2+offsetShiftR n (Offset o) = Offset (o `unsafeShiftR` n)++unsafePrimFromIO :: PrimMonad m => IO a -> m a+unsafePrimFromIO = unsafeIOToPrim++(.==#) :: Offset ty -> CountOf ty -> Bool+(.==#) (Offset ofs) (CountOf sz) = ofs == sz+{-# INLINE (.==#) #-}++#ifdef ML_KEM_TESTING+checkBounds :: CountOf ty -> Offset ty -> a -> a+checkBounds (CountOf sz) (Offset ofs) a+ | ofs >= 0 && ofs < sz = a+ | otherwise = error ("offset not in valid range: " ++ show ofs)+#endif
src/Block.hs view
@@ -3,67 +3,75 @@ -- License : BSD-3-Clause -- Copyright : (c) 2025 Olivier Chéron ----- An array of primitive (unlifted) elements. This module currently exposes--- the implementation from basement and fixes a lack of inlining in the--- @create@ function.+-- 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+ , create, foldZipWith, iterModify, Block.length, mutableContents , Block.new, Block.newPinned, Block.thaw, Block.unsafeCast- , Block.unsafeFreeze, Block.unsafeThaw, Block.withMutablePtr+ , Block.unsafeFreeze, Block.unsafeThaw+#ifdef ML_KEM_TESTING+ , Block.toList+#endif ) where -import Basement.Block (Block)-import Basement.Block.Mutable (MutableBlock)-import qualified Basement.Block as Block hiding (create, map)-import qualified Basement.Block.Mutable as Block-import Basement.Monad-import Basement.PrimType-import Basement.Types.OffsetSize+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 = Block.index-blockRead = Block.read-blockWrite = Block.write+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 = Block.unsafeIndex-blockRead = Block.unsafeRead-blockWrite = Block.unsafeWrite+blockIndex b (Offset i) = indexPrimArray b i+blockRead mb (Offset i) = readPrimArray mb i+blockWrite mb (Offset i) = writePrimArray mb i #endif create :: PrimType ty => CountOf ty -> (Offset ty -> ty) -> Block ty-create n initializer = runST $ do- mb <- Block.new n+create (CountOf n) initializer = runST $ do+ mb <- newPrimArray n loop mb 0- Block.unsafeFreeze mb+ unsafeFreezePrimArray mb where loop !mb i- | i .==# n = pure ()- | otherwise = Block.unsafeWrite mb i (initializer i) >> loop mb (i + 1)+ | i == n = pure ()+ | otherwise = writePrimArray mb i (initializer $ Offset i) >> loop mb (i + 1) {-# INLINE create #-} iterModify :: (PrimType ty, PrimMonad prim) => (ty -> ty) -> MutableBlock ty (PrimState prim) -> prim ()-iterModify f ma = loop 0+iterModify f ma = getSizeofMutablePrimArray ma >>= (`loop` 0) where- !sz = Block.mutableLength ma- loop i- | i .==# sz = pure ()- | otherwise = Block.unsafeRead ma i >>= \x -> Block.unsafeWrite ma i (f x) >> loop (i+1)+ loop n i+ | i == n = pure ()+ | otherwise = readPrimArray ma i >>= \x -> writePrimArray ma i (f x) >> loop n (i+1) {-# INLINE iterModify #-} foldZipWith :: (PrimType a, PrimType b)@@ -71,13 +79,42 @@ foldZipWith f c a b = assert (sa == sb) $ loop c 0 where- CountOf sa = Block.length a- CountOf sb = Block.length b+ sa = sizeofPrimArray a+ sb = sizeofPrimArray b loop !acc i | i == sa = acc | otherwise = do- let va = Block.unsafeIndex a (Offset i)- let vb = Block.unsafeIndex b (Offset i)+ let va = indexPrimArray a i+ let vb = indexPrimArray b i loop (f acc va vb) (i + 1) {-# INLINE foldZipWith #-}++length :: PrimType ty => Block ty -> CountOf ty+length = CountOf . sizeofPrimArray++mutableContents :: MutableBlock ty s -> Ptr ty+mutableContents = mutablePrimArrayContents -- pinned only++new :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))+new (CountOf n) = newPrimArray n++newPinned :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))+newPinned (CountOf n) = newPinnedPrimArray n++thaw :: (PrimMonad prim, PrimType ty) => Block ty -> prim (MutableBlock ty (PrimState prim))+thaw b = thawPrimArray b 0 (sizeofPrimArray b)++#ifdef ML_KEM_TESTING+toList :: PrimType ty => Block ty -> [ty]+toList = primArrayToList+#endif++unsafeCast :: Block a -> Block b+unsafeCast (PrimArray b) = PrimArray b++unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (Block ty)+unsafeFreeze = unsafeFreezePrimArray++unsafeThaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))+unsafeThaw = unsafeThawPrimArray
src/BlockN.hs view
@@ -19,18 +19,14 @@ #endif ) where -import Basement.Monad-import Basement.Nat-import Basement.NormalForm-import Basement.PrimType-import Basement.Types.OffsetSize-+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)@@ -40,15 +36,15 @@ newtype BlockN marking (n :: Nat) a = BlockN { unBlockN :: SecureBlock marking a } #ifdef ML_KEM_TESTING-instance (Classified marking, PrimType a) => Eq (BlockN marking n a) where+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 NormalForm (BlockN marking n a) where- toNormalForm = SecureBlock.toNormalForm . unBlockN+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 = create (const zero)@@ -113,7 +109,7 @@ f (SecureBlock.index a (Offset i)) (SecureBlock.index b (Offset i)) {-# INLINE zipWith #-} -unsafeCast :: PrimType b => BlockN marking n a -> SecureBlock marking b+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@@ -125,6 +121,7 @@ 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 #-} thaw :: (Classified marking, PrimMonad prim, PrimType a) => BlockN marking n a -> prim (MutableBlockN marking n a (PrimState prim)) thaw = fmap MutableBlockN . SecureBlock.thaw . unBlockN
src/Builder.hs view
@@ -13,8 +13,6 @@ , runToBlock, unsafeCreate ) where -import Basement.Types.OffsetSize- import Data.ByteArray (ByteArray) import Control.Monad.ST@@ -27,6 +25,7 @@ import qualified GHC.Exts as Exts +import Base import Block (Block) import Marking (Classified, Leak(..), SecurityMarking(..)) import SecureBytes (SecureBytes)@@ -77,7 +76,7 @@ runToBlock :: Builder Pub -> Block Word8 runToBlock b = runST $ do mb <- Block.newPinned (CountOf $ builderLength b)- Block.withMutablePtr mb (copyBuilderToPtr b)+ copyBuilderToPtr b (Block.mutableContents mb) Block.unsafeFreeze mb unsafeCreate :: Int -> (Ptr a -> IO ()) -> Builder marking
src/Crypto.hs view
@@ -32,6 +32,7 @@ import GHC.TypeNats +import Foreign.Marshal.Utils (fillBytes) import Foreign.Ptr (Ptr, castPtr, plusPtr) import Foreign.Storable (pokeByteOff) @@ -202,7 +203,7 @@ g c = (B.convert $ B.takeView ab 32, B.dropView ab 32) where ab = Builder.run $ hashWith SHA3_512 c --- Override cryptonite types and hashing functions.+-- Override cryptonite/crypton types and hashing functions. -- -- Standard type Digest is a newtype over an unpinned Block Word8, which -- requires a trampoline to implement most Ptr access to the underlying byte@@ -226,4 +227,4 @@ hashMutableUpdate (ctx :: MutableContext a) ba B.withByteArray ctx $ \pctx -> do hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig- ScrubbedBlock.erasePtr (B.length ctx) pctx+ fillBytes pctx 0 (B.length ctx)
src/Crypto/PubKey/ML_KEM.hs view
@@ -48,9 +48,10 @@ generate :: (ParamSet a, MonadRandom m) => proxy a -> m (EncapsulationKey a, DecapsulationKey a) generate p = do- d <- getRandomBytes 32- z <- getRandomBytes 32- return (Internal.keyGen p (d :: ScrubbedBytes) z)+ seed <- getRandomBytes 64+ let d = B.takeView seed 32+ z = B.drop 32 seed+ return (Internal.keyGen p d z) -- | Generate an ML-KEM key pair from the specified seed (d, z). Length of -- inputs must be 32 bytes.
src/Internal.hs view
@@ -18,14 +18,13 @@ , keyGen, toPublic, encaps, decaps ) where -import Basement.Nat- import Control.DeepSeq (NFData(..)) import Control.Monad import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes) import qualified Data.ByteArray as B +import Base import qualified Builder import qualified Crypto import qualified K_PKE as K
src/K_PKE.hs view
@@ -16,9 +16,6 @@ , EncryptionKey, ekEncode, ekDecode ) where -import Basement.Nat-import Basement.Types.OffsetSize- import Control.DeepSeq (NFData(..)) import Control.Monad @@ -28,6 +25,7 @@ import Unsafe.Coerce import Auxiliary (Rq, Tq, (..+), (..-))+import Base import Builder (Builder) import Iterate import Marking (SecurityMarking(..), Leak(..))
src/Machine.hs view
@@ -22,7 +22,7 @@ || defined(powerpc64le_HOST_ARCH) #define MLKEM_ALLOW_UNALIGNED_OP 1 --- Little-endian conversion in `basement` and `memory` is avoided at compile+-- Little-endian conversion in `memory` / `ram` is avoided at compile -- time only for AMD/Intel, here we will short circuit on ARM too #if (defined(arm_HOST_ARCH) || defined(aarch64_HOST_ARCH)) \ && !defined(WORDS_BIGENDIAN)
src/Marking.hs view
@@ -21,14 +21,7 @@ #endif ) where -import Basement.Monad-import Basement.NormalForm-import Basement.PrimType-import Basement.Types.OffsetSize-#ifdef ML_KEM_TESTING-import qualified Basement.Compat.IsList-#endif-+import Control.DeepSeq (NFData(..)) import Control.Monad.ST import Data.ByteArray (Bytes, ScrubbedBytes)@@ -40,6 +33,7 @@ import Unsafe.Coerce +import Base import Block (Block, MutableBlock, blockIndex) import ScrubbedBlock (ScrubbedBlock) import qualified Block@@ -67,7 +61,7 @@ unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty) #ifdef ML_KEM_TESTING- eq :: PrimType ty => SecureBlock marking ty -> SecureBlock marking ty -> Bool+ eq :: (Eq ty, PrimType ty) => SecureBlock marking ty -> SecureBlock marking ty -> Bool showsPrec :: (PrimType ty, Show ty) => Int -> SecureBlock marking ty -> ShowS lengthBlock :: PrimType ty => SecureBlock marking ty -> CountOf ty #endif@@ -134,11 +128,11 @@ #ifdef ML_KEM_TESTING toList :: PrimType ty => SecureBlock marking ty -> [ty]-toList = Basement.Compat.IsList.toList . unwrap+toList = Block.toList . unwrap #endif toNormalForm :: SecureBlock marking ty -> ()-toNormalForm = Basement.NormalForm.toNormalForm . unwrap+toNormalForm = rnf . unwrap -unsafeCast :: PrimType b => SecureBlock marking a -> SecureBlock marking b+unsafeCast :: SecureBlock marking a -> SecureBlock marking b unsafeCast = wrap . Block.unsafeCast . unwrap
src/Matrix.hs view
@@ -16,9 +16,7 @@ #endif ) where -import Basement.Nat-import Basement.Types.OffsetSize-+import Base import Math import Vector (Vector) import qualified Vector
src/ScrubbedBlock.hs view
@@ -8,39 +8,24 @@ -- package memory but for blocks. -- {-# LANGUAGE BangPatterns #-}-{-# LANGUAGE CPP #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE UnboxedTuples #-} module ScrubbedBlock ( ScrubbedBlock, create, foldZipWith, ScrubbedBlock.length- , new, thaw, unsafeFreeze, Block.withMutablePtr- , erasePtr+ , new, thaw, unsafeFreeze ) where -import Basement.Block (Block(..), MutableBlock(..), isPinned)-import Basement.Block.Mutable (mutableLengthBytes, unsafeCopyBytesRO)-import Basement.Compat.Primitive--import Basement.Monad-import Basement.PrimType-import Basement.Types.OffsetSize+import Data.Primitive.PrimArray as Block import Control.Exception (assert) import Control.Monad.ST import Data.Word -import Foreign.Ptr (Ptr)--import Block (blockWrite)+import Base+import Block (Block, MutableBlock, blockWrite) import qualified Block -#if MIN_VERSION_base(4,19,0)-import GHC.Base (setAddrRange#)-import GHC.Exts (Ptr(Ptr))-#else-import Data.Memory.PtrMethods (memSet)-#endif import GHC.Base (IO(IO), Int(I#), setByteArray#) import GHC.Exts (mkWeak#) @@ -74,10 +59,12 @@ new = Block.newPinned -- always pinned thaw :: (PrimType ty, PrimMonad m) => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m))-thaw (ScrubbedBlock !b) = do- mb <- new (Block.length b)- unsafeCopyBytesRO mb 0 b 0 (mutableLengthBytes mb)+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@@ -86,7 +73,7 @@ {- internal -} assertPinned :: Block ty -> a -> a-assertPinned mb = assert (isPinned mb == Pinned)+assertPinned mb = assert (Block.isPrimArrayPinned mb) scrubbed :: PrimMonad prim => Block ty -> prim (ScrubbedBlock ty) scrubbed b = assertPinned b $ unsafePrimFromIO $@@ -101,19 +88,10 @@ where CountOf len = Block.length b addBlockFinalizer :: Block ty -> IO () -> IO ()-addBlockFinalizer (Block barr) (IO finalizer) = IO $ \s ->+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) (MutableBlock mbarr) = IO $ \s1 ->+erase (I# len) (Block.MutablePrimArray mbarr) = IO $ \s1 -> case setByteArray# mbarr 0# len 0# s1 of s2 -> (# s2, () #)--erasePtr :: Int -> Ptr Word8 -> IO ()-#if MIN_VERSION_base(4,19,0)-erasePtr (I# n) (Ptr addr) = IO $ \s1 ->- case setAddrRange# addr n 0# s1 of- s2 -> (# s2, () #)-#else-erasePtr n ptr = memSet ptr 0 n-#endif
src/SecureBlock.hs view
@@ -15,12 +15,7 @@ #endif ) where -import Basement.Monad-import Basement.PrimType-#ifdef ML_KEM_TESTING-import Basement.Types.OffsetSize-#endif-+import Base import Block (MutableBlock) import Marking import qualified Block
src/Vector.hs view
@@ -4,13 +4,12 @@ -- Copyright : (c) 2025 Olivier Chéron -- -- A vector of lifted elements with the vector dimension at type level.--- Currently backed by type t'Array' from basement.+-- Backed by type t'SmallArray' from primitive. -- {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-}-{-# LANGUAGE GeneralizedNewtypeDeriving #-}-{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ScopedTypeVariables #-} module Vector ( Vector, Vector.concatMap@@ -21,30 +20,28 @@ #endif ) where -import Basement.BoxedArray (Array)-import qualified Basement.BoxedArray as Array-#ifdef ML_KEM_TESTING-import Basement.Compat.IsList-#endif-import Basement.Nat-import Basement.NormalForm-import Basement.Types.OffsetSize+import Data.Primitive.SmallArray import Control.DeepSeq (NFData(..)) #ifdef ML_KEM_TESTING import Control.Monad #endif+import Control.Monad.ST #if !(MIN_VERSION_base(4,20,0)) import Data.List (foldl') #endif import Data.Proxy +import Base import Iterate import Math +type Array = SmallArray+type MArray ty s = SmallMutableArray s ty+ newtype Vector (n :: Nat) a = Vector { unVector :: Array a }- deriving (Eq, Show, NormalForm)+ deriving (Eq, Show) instance Functor (Vector n) where fmap = mapVector@@ -60,12 +57,35 @@ neg = mapVector neg {-# INLINE neg #-} +arrayCreate :: forall ty. CountOf ty -> (Offset ty -> ty) -> Array ty+arrayCreate n initializer = runST (arrayNew n >>= iter initializer)+ where+ iter :: PrimMonad prim => (Offset ty -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)+ iter f ma = loop 0+ where+ loop s@(Offset i)+ | s .==# n = unsafeFreezeSmallArray ma+ | otherwise = writeSmallArray ma i (f s) >> loop (s + 1)+ {-# INLINE loop #-}+ {-# INLINE iter #-}++arrayLength :: Array ty -> CountOf ty+arrayLength = CountOf . sizeofSmallArray++arrayMap :: (a -> b) -> Array a -> Array b+arrayMap f a = arrayCreate (CountOf sz) $ \(Offset i) -> f $ arrayIndex a (Offset i)+ where CountOf sz = arrayLength a++arrayNew :: PrimMonad prim => CountOf ty -> prim (MArray ty (PrimState prim))+arrayNew (CountOf c) = newSmallArray c placeholder+ where placeholder = error "arrayNew: unexpected evaluation"+ create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a create = genericCreate {-# INLINE create #-} genericCreate :: forall n a a'. KnownNat n => (Offset a' -> a) -> Vector n a-genericCreate f = Vector $ Array.create (CountOf sz) (\(Offset !i) -> f (Offset i))+genericCreate f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i)) where !sz = fromIntegral $ natVal (Proxy :: Proxy n) {-# INLINE [1] genericCreate #-} @@ -78,18 +98,19 @@ {-# INLINE [1] genericCreateZipRight #-} mapVector :: (a -> b) -> Vector n a -> Vector n b-mapVector f = Vector <$> fmap f . unVector+mapVector f = Vector <$> arrayMap f . unVector {-# INLINE [1] mapVector #-} arrayIndex :: Array a -> Offset a -> a #ifdef ML_KEM_TESTING-arrayIndex = Array.index+arrayIndex a off@(Offset i) =+ checkBounds (arrayLength a) off $ indexSmallArray a i replicateM :: forall n m a. (KnownNat n, Applicative m) => m a -> m (Vector n a)-replicateM f = Vector . fromList <$> Control.Monad.replicateM sz f+replicateM f = Vector . smallArrayFromList <$> Control.Monad.replicateM sz f where !sz = fromIntegral $ natVal (Proxy :: Proxy n) #else-arrayIndex = Array.unsafeIndex+arrayIndex a (Offset i) = indexSmallArray a i #endif index :: Vector n a -> Offset a -> a@@ -101,14 +122,14 @@ mapToList :: (a -> b) -> Vector n a -> [b] mapToList f (Vector a) = Prelude.map (f . arrayIndex a . Offset) (offsets sa)- where CountOf sa = Array.length a+ where CountOf sa = arrayLength a zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c zipWith f (Vector a) (Vector !b) = Vector $- Array.create (CountOf sa) $ \(Offset i) ->+ arrayCreate (CountOf sa) $ \(Offset i) -> f (arrayIndex a (Offset i)) (arrayIndex b (Offset i)) where- CountOf sa = Array.length a+ CountOf sa = arrayLength a {-# INLINE [1] zipWith #-} fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c@@ -117,18 +138,18 @@ where ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i)) gg = g (arrayIndex a 0) (arrayIndex b 0)- CountOf !sa = Array.length a+ CountOf !sa = arrayLength a {-# INLINE fold1ZipWith #-} foldIndexWith :: (c -> Offset a -> a -> c) -> c -> Vector n a -> c foldIndexWith f c (Vector a) = foldl' g c (offsets sa) where g x i = f x (Offset i) (arrayIndex a (Offset i))- CountOf !sa = Array.length a+ CountOf !sa = arrayLength a {-# INLINE foldIndexWith #-} toNormalForm :: NFData a => Vector n a -> ()-toNormalForm = Array.foldl' (\acc x -> acc `seq` rnf x) () . unVector+toNormalForm = foldl' (\acc x -> acc `seq` rnf x) () . unVector {-# RULES "mapVector/mapVector" [2] forall f g a. mapVector f (mapVector g a) = mapVector (f . g) a