mlkem 0.2.0.0 → 0.2.1.0
raw patch · 17 files changed
+511/−248 lines, 17 filesPVP ok
version bump matches the API change (PVP)
API changes (from Hackage documentation)
+ Crypto.PubKey.ML_KEM: generateOpen :: (ParamSet a, ByteArray d, ByteArray z, MonadRandom m) => proxy a -> m (EncapsulationKey a, DecapsulationKey a, d, z)
Files
- CHANGELOG.md +7/−0
- mlkem.cabal +5/−1
- src/Auxiliary.hs +106/−105
- src/Base.hs +14/−5
- src/Block.hs +24/−33
- src/BlockN.hs +128/−21
- src/Crypto/PubKey/ML_KEM.hs +17/−2
- src/Equality.hs +25/−0
- src/Fusion.hs +102/−0
- src/K_PKE.hs +9/−6
- src/Marking.hs +1/−6
- src/Math.hs +14/−0
- src/Matrix.hs +6/−8
- src/ScrubbedBlock.hs +4/−24
- src/SecureBlock.hs +3/−6
- src/Vector.hs +32/−20
- tests/Tests.hs +14/−11
CHANGELOG.md view
@@ -1,5 +1,12 @@ # Changelog for `mlkem` +## 0.2.1.0 - 2026-05-28++* Function `generateOpen` is added to return not only the key pair but also the+ seed that was used++* Optimizations to remove intermediate allocations where possible+ ## 0.2.0.0 - 2026-03-26 * Flag `use_crypton` is now enabled by default. It requires crypton >= 1.1.1
mlkem.cabal view
@@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack name: mlkem-version: 0.2.0.0+version: 0.2.1.0 synopsis: Module-Lattice-based Key-Encapsulation Mechanism description: Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM) implemented in Haskell.@@ -44,6 +44,8 @@ Builder ByteArrayST Crypto+ Equality+ Fusion Internal Iterate K_PKE@@ -125,6 +127,8 @@ ByteArrayST Crypto Crypto.PubKey.ML_KEM+ Equality+ Fusion Internal Iterate K_PKE
src/Auxiliary.hs view
@@ -146,6 +146,9 @@ setOffAddr# addr i len (Zq a) = setOffAddr# addr i len a {-# INLINE setOffAddr# #-} +instance PrimSized Zq where+ type PrimSize Zq = 2+ instance Add Zq where zero = Zq 0 Zq a .+ Zq b = Zq $ reduceSimple (a + b)@@ -182,22 +185,23 @@ instance Classified marking => Add (Rq marking) where zero = Rq zero Rq a .+ Rq b = Rq (a .+ b)+ {-# INLINE (.+) #-} Rq a .- Rq b = Rq (a .- b)+ {-# INLINE (.-) #-} neg (Rq a) = Rq (neg a)- {-# SPECIALIZE instance Add (Rq Sec) #-}- {-# SPECIALIZE instance Add (Rq Pub) #-}+ {-# INLINE neg #-} infixl 6 ..+, ..- -- Transformation called only at expected location in the LWE problem, after -- adding noise to secret information. (..+) :: Rq Sec -> Rq Sec -> Rq Pub-Rq a ..+ Rq b = Rq $ BlockN.zipWith (.+) a b-{-# NOINLINE (..+) #-}+a ..+ b = leak (a .+ b)+{-# INLINE (..+) #-} (..-) :: Rq Pub -> Rq Sec -> Rq Sec-Rq a ..- Rq b = Rq $ BlockN.zipWith (.-) a b-{-# NOINLINE (..-) #-}+Rq a ..- Rq b = Rq $ BlockN.zipWith (flip (.-)) b a+{-# INLINE (..-) #-} instance Leak Rq @@ -219,20 +223,21 @@ instance Classified marking => Add (Tq marking) where zero = Tq zero Tq a .+ Tq b = Tq (a .+ b)+ {-# INLINE (.+) #-} Tq a .- Tq b = Tq (a .- b)+ {-# INLINE (.-) #-} neg (Tq a) = Tq (neg a)- {-# SPECIALIZE instance Add (Tq Sec) #-}- {-# SPECIALIZE instance Add (Tq Pub) #-}+ {-# INLINE neg #-} instance Leak Tq instance BiMul (Tq Pub) (Tq Sec) where (..*) = multiplyNTTs- {-# NOINLINE (..*) #-}+ {-# INLINE (..*) #-} instance BiMulAdd (Tq Pub) (Tq Sec) where- biMulAdd = multiplyNTTsAdd- {-# NOINLINE biMulAdd #-}+ biMulFold = multiplyNTTsFold+ {-# INLINE biMulFold #-} #ifdef ML_KEM_TESTING instance Mul (Tq Sec) where@@ -255,61 +260,66 @@ -- Computes the NTT representation of the given polynomial ntt :: Classified marking => Rq marking -> Tq marking-ntt (Rq !a) = runST $ do- b <- BlockN.thaw a- outer b 1 128- Tq <$> BlockN.unsafeFreeze b+ntt (Rq a) = Tq $ BlockN.runThaw a mutNtt+{-# INLINE ntt #-}++mutNtt :: MutableBlockN marking N Zq s -> ST s ()+mutNtt !b = outer 1 128 where- outer !b !i len = when (len >= 2) $ inner b i len 0+ outer !i len = when (len >= 2) $ inner i len 0 - inner !b !i !len start+ inner !i !len start | start < 256 = do let zeta = BlockN.index zetaPowBitRev i -- 17 ^ bitRev7 i- loop b zeta (start + len) len start- inner b (i + 1) len (start + offsetShiftL 1 len)- | otherwise = outer b i (offsetShiftR 1 len)+ loop zeta (start + len) len start+ inner (i + 1) len (start + offsetShiftL 1 len)+ | otherwise = outer i (offsetShiftR 1 len) - loop !b !zeta end len j =+ loop !zeta end len j = when (j < end) $ do t <- (zeta .*) <$> BlockN.read b (j + len) x <- BlockN.read b j BlockN.write b (j + len) (x .- t) BlockN.write b j (x .+ t)- loop b zeta end len (j + 1)-{-# SPECIALIZE ntt :: Rq Sec -> Tq Sec #-}-{-# SPECIALIZE ntt :: Rq Pub -> Tq Pub #-}+ loop zeta end len (j + 1)+{-# NOINLINE mutNtt #-} -- Computes the polynomial that corresponds to the given NTT representation nttInv :: Tq Sec -> Rq Sec-nttInv (Tq !a) = runST $ do- b <- BlockN.thaw a- outer b 127 2+nttInv (Tq a) = Rq $ BlockN.runThaw a mutNttInv+{-# INLINE nttInv #-}++mutNttInv :: MutableBlockN Sec N Zq s -> ST s ()+mutNttInv !b = do+ outer 127 2 BlockN.iterModify (\x -> x .* Zq 3303) b- Rq <$> BlockN.unsafeFreeze b where- outer !b !i len = when (len <= 128) $ inner b i len 0+ outer !i len = when (len <= 128) $ inner i len 0 - inner !b !i !len start+ inner !i !len start | start < 256 = do let zeta = BlockN.index zetaPowBitRev i -- 17 ^ bitRev7 i- loop b zeta (start + len) len start- inner b (i - 1) len (start + offsetShiftL 1 len)- | otherwise = outer b i (offsetShiftL 1 len)+ loop zeta (start + len) len start+ inner (i - 1) len (start + offsetShiftL 1 len)+ | otherwise = outer i (offsetShiftL 1 len) - loop !b !zeta end len j =+ loop !zeta end len j = when (j < end) $ do t <- BlockN.read b j x <- BlockN.read b (j + len) BlockN.write b j (t .+ x) BlockN.write b (j + len) (zeta .* (x .- t))- loop b zeta end len (j + 1)+ loop zeta end len (j + 1)+{-# NOINLINE mutNttInv #-} -- Computes the product of two NTT representations multiplyNTTs :: Tq Pub -> Tq Sec -> Tq Sec-multiplyNTTs (Tq !f) (Tq !g) = runST $ do- b <- BlockN.new (Proxy :: Proxy Sec)- loop b 0- Tq <$> BlockN.unsafeFreeze b+multiplyNTTs f g = Tq $+ BlockN.runNew (Proxy :: Proxy Sec) $ mutMultiplyNTTs f g+{-# INLINE multiplyNTTs #-}++mutMultiplyNTTs :: Tq Pub -> Tq Sec -> MutableBlockN Sec N Zq s -> ST s ()+mutMultiplyNTTs (Tq !f) (Tq !g) bb = loop bb 0 where loop :: MutableBlockN Sec N Zq s -> Offset Zq -> ST s () loop !b i = when (i < 128) $ do@@ -332,22 +342,24 @@ !c0 = reduce (a0 `mul` b0 + a1 `mul` b1g) !c1 = reduce (a0 `mul` b1 + a1 `mul` b0) +multiplyNTTsFold :: Foldable t => Tq Sec -> t (Tq Pub, Tq Sec) -> Tq Sec+multiplyNTTsFold (Tq c) =+ Tq . BlockN.runFold c (uncurry multiplyNTTsAdd)+{-# INLINE multiplyNTTsFold #-}+ -- Multiply then add a third term-multiplyNTTsAdd :: Tq Pub -> Tq Sec -> Tq Sec -> Tq Sec-multiplyNTTsAdd (Tq !f) (Tq !g) (Tq !h) = runST $ do- b <- BlockN.new (Proxy :: Proxy Sec)- loop b 0- Tq <$> BlockN.unsafeFreeze b+multiplyNTTsAdd :: Tq Pub -> Tq Sec -> MutableBlockN Sec N Zq s -> ST s ()+multiplyNTTsAdd (Tq !f) (Tq !g) bb = loop bb 0 where loop :: MutableBlockN Sec N Zq s -> Offset Zq -> ST s () loop !b i = when (i < 128) $ do let ii = offsetShiftL 1 i- a0 = BlockN.index f ii+ c0 <- BlockN.read b ii+ c1 <- BlockN.read b (ii + 1)+ let a0 = BlockN.index f ii a1 = BlockN.index f (ii + 1) b0 = BlockN.index g ii b1 = BlockN.index g (ii + 1)- c0 = BlockN.index h ii- c1 = BlockN.index h (ii + 1) (d0, d1) = baseCaseMultiplyAdd a0 a1 b0 b1 c0 c1 (BlockN.index gamma i) BlockN.write b ii d0 BlockN.write b (ii + 1) d1@@ -364,17 +376,15 @@ -- Values of 17 ^ BitRev7(𝑖) mod 𝑞 for 𝑖 ∈ {0, … , 127} zetaPowBitRev :: BlockN Pub 128 Zq-zetaPowBitRev = runST $ do- out <- BlockN.new (Proxy :: Proxy Pub)+zetaPowBitRev = BlockN.runNew (Proxy :: Proxy Pub) $ \out -> foldM_ (loop out) one offsets- BlockN.unsafeFreeze out where offsets = Prelude.map (fromIntegral . bitRev7) [0 .. 127] loop b acc i = BlockN.write b i acc >> return (Zq 17 .* acc) -- Values of 17 ^ 2.BitRev7(𝑖)+1 mod 𝑞 for 𝑖 ∈ {0, … , 127} gamma :: BlockN Pub 128 Zq-gamma = BlockN.map (\z -> z .* z .* Zq 17) zetaPowBitRev+gamma = BlockN.mapEqPrimSize (\z -> z .* z .* Zq 17) zetaPowBitRev -- Compress a field element with 𝑑 < 12 compress :: Int -> Zq -> Word16@@ -393,22 +403,18 @@ -- Compress a polynomial with 𝑑 < 12 rcompress :: Classified marking => Int -> Rq marking -> BlockN marking N Word16-rcompress !d (Rq a) = BlockN.map (compress d) a-{-# SPECIALIZE NOINLINE rcompress :: Int -> Rq Sec -> BlockN Sec N Word16 #-}-{-# SPECIALIZE NOINLINE rcompress :: Int -> Rq Pub -> BlockN Pub N Word16 #-}+rcompress !d (Rq a) = BlockN.mapEqPrimSize (compress d) a+{-# INLINE rcompress #-} -- Decompress a polynomial with 𝑑 < 12 rdecompress :: Classified marking => Int -> BlockN marking N Word16 -> Rq marking-rdecompress !d = Rq . BlockN.map (decompress d)-{-# SPECIALIZE NOINLINE rdecompress :: Int -> BlockN Sec N Word16 -> Rq Sec #-}-{-# SPECIALIZE NOINLINE rdecompress :: Int -> BlockN Pub N Word16 -> Rq Pub #-}+rdecompress !d = Rq . BlockN.mapEqPrimSize (decompress d)+{-# INLINE rdecompress #-} -- Generates a pseudorandom element of T𝑞 from a seed and two indices sampleNTT :: SecureBytes Pub -> Word8 -> Word8 -> Tq Pub-sampleNTT seed !x !y = runST $ do- b <- BlockN.new (Proxy :: Proxy Pub)- runXof b (280 * 3) 0 0- Tq <$> BlockN.unsafeFreeze b+sampleNTT seed !x !y = Tq $+ BlockN.runNew (Proxy :: Proxy Pub) $ \b -> runXof b (280 * 3) 0 0 where runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of SomeNat proxy -> do@@ -479,10 +485,13 @@ -- Takes a seed as input and outputs a pseudorandom sample from the -- distribution D_eta samplePolyCBD :: Word -> SecureBytes Sec -> Rq Sec-samplePolyCBD !eta !input = runST $ ST.withByteArray input $ \p -> do- f <- BlockN.new (Proxy :: Proxy Sec)- loop p f 0 zeroPos- Rq <$> BlockN.unsafeFreeze f+samplePolyCBD eta input = Rq $+ BlockN.runNew (Proxy :: Proxy Sec) $ mutSamplePolyCBD eta input+{-# INLINE samplePolyCBD #-}++mutSamplePolyCBD :: Word -> SecureBytes Sec -> MutableBlockN Sec N Zq s -> ST s ()+mutSamplePolyCBD !eta !input ff =+ ST.withByteArray input $ \p -> loop p ff 0 zeroPos where loop :: Ptr WordLE -> MutableBlockN Sec N Zq s -> Offset Zq -> BitPos -> ST s () loop !p !f !i !bp = when (i < Offset n) $ do@@ -499,6 +508,7 @@ let (howMany, bp') = nextPos j bp bits = x .&. getMask howMany getBits p bp' (acc + fromIntegral (popCount bits)) (j - howMany)+{-# NOINLINE mutSamplePolyCBD #-} -- Encodes an array of 𝑑-bit integers into a byte array for 1 ≤ 𝑑 ≤ 12 byteEncode :: Int -> BlockN marking N Word16 -> Builder marking@@ -506,19 +516,18 @@ {-# INLINE byteEncode #-} runByteEncode :: Int -> BlockN marking N Word16 -> Ptr WordLE -> ST s ()-runByteEncode !d !f dst = outer dst zeroPos 0 0+runByteEncode !d !f dst = loop dst 0 zeroPos 0 (get 0) d where- outer :: Ptr WordLE -> BitPos -> WordM -> Int -> ST s ()- outer !b !bp !o pos = when (pos < n) $- inner b pos bp o (BlockN.index f (Offset pos)) d- {-# NOINLINE outer #-}+ get = BlockN.index f . Offset+ {-# INLINE get #-} - inner :: Ptr WordLE -> Int -> BitPos -> WordM -> Word16 -> Int -> ST s ()- inner !b !pos !bp !o !a j- | j == 0 = outer b bp o (pos + 1)- | bitPos bp + howMany < wordBits = inner b pos bp' o' a' j'- | otherwise = pokeWordPos b bp o' >> inner b pos bp' 0 a' j'+ loop !b !pos !bp !o !a j+ | j == 0, pos' == n = return ()+ | j == 0 = loop b pos' bp o (get pos') d+ | bitPos bp + howMany < wordBits = loop b pos bp' o' a' j'+ | otherwise = pokeWordPos b bp o' >> loop b pos bp' 0 a' j' where+ pos' = pos + 1 (howMany, bp') = nextPos j bp x = fromIntegral a .&. getMask howMany o' = o .|. (x `unsafeShiftL` bitPos bp)@@ -554,57 +563,49 @@ -- Decodes a byte array into an array of 𝑑-bit integers for 1 ≤ 𝑑 ≤ 12 byteDecode :: forall marking ba. (Classified marking, ByteArrayAccess ba) => Int -> ba -> BlockN marking N Word16-byteDecode !d !b = runST $- ST.withByteArray b $ \p -> do- f <- BlockN.new (Proxy :: Proxy marking)- outer f p zeroPos 0- BlockN.unsafeFreeze f+byteDecode d b = BlockN.runNew (Proxy :: Proxy marking) $ mutByteDecode d b+{-# INLINE byteDecode #-}++mutByteDecode :: ByteArrayAccess ba => Int -> ba -> MutableBlockN marking N Word16 s -> ST s ()+mutByteDecode !d !b !f = ST.withByteArray b $ \p -> outer p zeroPos 0 where- outer :: MutableBlockN marking N Word16 s -> Ptr WordLE -> BitPos -> Offset Word16 -> ST s ()- outer !f !p !bp i = when (i < Offset n) $ inner f p i bp 0 0+ outer !p !bp i = when (i < Offset n) $ inner p i bp 0 0 - inner :: MutableBlockN marking N Word16 s -> Ptr WordLE -> Offset Word16 -> BitPos -> Word16 -> Int -> ST s ()- inner !f !p !i !bp !v j- | j == d = BlockN.write f i v >> outer f p bp (i + 1)+ inner !p !i !bp !v j+ | j == d = BlockN.write f i v >> outer p bp (i + 1) | otherwise = do let (howMany, bp') = nextPos (d - j) bp y <- get p bp howMany let v' = v .|. (fromIntegral y `unsafeShiftL` j) j' = j + howMany- inner f p i bp' v' j'+ inner p i bp' v' j' get :: Ptr WordLE -> BitPos -> Int -> ST s WordM get p bp howMany = do x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp return (x .&. getMask howMany)-{-# SPECIALIZE byteDecode :: forall ba. ByteArrayAccess ba => Int -> View ba -> BlockN Sec N Word16 #-}-{-# SPECIALIZE byteDecode :: forall ba. ByteArrayAccess ba => Int -> View ba -> BlockN Pub N Word16 #-}-{-# SPECIALIZE byteDecode :: Int -> View Bytes -> BlockN Sec N Word16 #-}-{-# SPECIALIZE byteDecode :: Int -> View Bytes -> BlockN Pub N Word16 #-}+{-# SPECIALIZE mutByteDecode :: forall marking s. Int -> View Bytes -> MutableBlockN marking N Word16 s -> ST s () #-} -- Optimization of byteDecode when 𝑑=1 byteDecode1 :: ByteArrayAccess ba => ba -> BlockN Sec N Word16-byteDecode1 !b = runST $- ST.withByteArray b $ \p -> do- f <- BlockN.new (Proxy :: Proxy Sec)- outer f p 0- BlockN.unsafeFreeze f+byteDecode1 b = BlockN.runNew (Proxy :: Proxy Sec) $ mutByteDecode1 b+{-# INLINE byteDecode1 #-}++mutByteDecode1 :: ByteArrayAccess ba => ba -> MutableBlockN Sec N Word16 s -> ST s ()+mutByteDecode1 !b !f = ST.withByteArray b $ \p -> outer p 0 where- outer :: MutableBlockN Sec N Word16 s -> Ptr WordLE -> Int -> ST s ()- outer !f !p i = when (i < n) $ do+ outer !p i = when (i < n) $ do x <- peekWord p- inner f (p `plusPtr` wordBytes) x i 0+ inner (p `plusPtr` wordBytes) x i 0 - inner :: MutableBlockN Sec N Word16 s -> Ptr WordLE -> WordM -> Int -> Int -> ST s ()- inner !f !p !acc !i j- | j == wordBits = outer f p i+ inner !p !acc !i j+ | j == wordBits = outer p i | otherwise = do let v = fromIntegral (acc .&. 1) BlockN.write f (Offset i) v- inner f p (acc `unsafeShiftR` 1) (i + 1) (j + 1)+ inner p (acc `unsafeShiftR` 1) (i + 1) (j + 1) -- byteDecode with 𝑑=12 and conversion to the field byteDecode12 :: (Classified marking, ByteArrayAccess ba) => ba -> Tq marking-byteDecode12 = Tq . BlockN.map toZq . byteDecode 12-{-# SPECIALIZE byteDecode12 :: ByteArrayAccess ba => View ba -> Tq Sec #-}-{-# SPECIALIZE byteDecode12 :: ByteArrayAccess ba => View ba -> Tq Pub #-}+byteDecode12 = Tq . BlockN.mapEqPrimSize toZq . byteDecode 12+{-# INLINE byteDecode12 #-}
src/Base.hs view
@@ -7,10 +7,12 @@ -- {-# LANGUAGE CPP #-} {-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE TypeFamilies #-} module Base- ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimState- , natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#)+ ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimSized(..)+ , PrimState, natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#) #ifdef ML_KEM_TESTING , checkBounds #endif@@ -18,16 +20,23 @@ import Control.Monad.Primitive -import Data.Primitive.Types+import Data.Primitive.Types (Prim) import Data.Bits+import Data.Word -import GHC.TypeLits+import GHC.TypeNats type PrimType = Prim +class PrimType a => PrimSized a where+ type PrimSize a :: Nat++instance PrimSized Word16 where+ type PrimSize Word16 = 2+ newtype CountOf ty = CountOf Int- deriving (Show, Eq, Ord)+ deriving (Show, Eq, Ord, Num) newtype Offset ty = Offset Int deriving (Show, Eq, Ord, Num)
src/Block.hs view
@@ -11,9 +11,9 @@ {-# LANGUAGE CPP #-} module Block ( Block, MutableBlock, blockIndex, blockRead, blockWrite- , create, foldZipWith, iterModify, Block.length, mutableContents- , Block.new, Block.newPinned, Block.thaw, Block.unsafeCast- , Block.unsafeFreeze, Block.unsafeThaw+ , foldZipWith, Block.length, mutableContents+ , Block.new, Block.newPinned, Block.thaw, thawPinned+ , Block.unsafeCast, unsafeCastMut, Block.unsafeFreeze, Block.unsafeThaw #ifdef ML_KEM_TESTING , Block.toList #endif@@ -21,14 +21,14 @@ import Control.Monad.Primitive +import Data.Primitive.ByteArray import Data.Primitive.PrimArray import Control.Exception (assert)-import Control.Monad.ST -import Foreign.Ptr+import Foreign.Ptr (Ptr) -import Base hiding (PrimMonad, PrimState)+import Base type Block = PrimArray type MutableBlock ty s = MutablePrimArray s ty@@ -49,31 +49,6 @@ blockWrite mb (Offset i) = writePrimArray mb i #endif -create :: PrimType ty- => CountOf ty- -> (Offset ty -> ty)- -> Block ty-create (CountOf n) initializer = runST $ do- mb <- newPrimArray n- loop mb 0- unsafeFreezePrimArray mb- where- loop !mb i- | 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 = getSizeofMutablePrimArray ma >>= (`loop` 0)- where- 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) => (c -> a -> b -> c) -> c -> Block a -> Block b -> c foldZipWith f c a b = assert (sa == sb) $@@ -102,9 +77,22 @@ 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)+thaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))+thaw (PrimArray !barr) = unsafeSTToPrim $+ -- as optimization, combine both steps in a known monad and avoid+ -- round trip between byte length and element count+ thawByteArray ba 0 (sizeofByteArray ba) >>= \(MutableByteArray mbarr) ->+ return (MutablePrimArray mbarr)+ where ba = ByteArray barr +thawPinned :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))+thawPinned (PrimArray !barr) = unsafeSTToPrim $ do+ let ba = ByteArray barr+ n = sizeofByteArray ba+ mb@(MutableByteArray mbarr) <- newPinnedByteArray n+ copyByteArray mb 0 ba 0 n+ return (MutablePrimArray mbarr)+ #ifdef ML_KEM_TESTING toList :: PrimType ty => Block ty -> [ty] toList = primArrayToList@@ -112,6 +100,9 @@ unsafeCast :: Block a -> Block b unsafeCast (PrimArray b) = PrimArray b++unsafeCastMut :: MutableBlock a m -> MutableBlock b m+unsafeCastMut (MutablePrimArray mb) = MutablePrimArray mb unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (Block ty) unsafeFreeze = unsafeFreezePrimArray
src/BlockN.hs view
@@ -8,26 +8,27 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-}-{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE TypeFamilies #-} {-# 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+ ( BlockN, MutableBlockN, index, iterModify, mapEqPrimSize+ , BlockN.read, runNew, runThaw, runFold, BlockN.unsafeCast+ , BlockN.write, BlockN.zipWith #ifdef ML_KEM_TESTING- , BlockN.fromList, BlockN.replicate, BlockN.toList+ , create, 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 Block (MutableBlock, blockRead, blockWrite, unsafeCastMut)+import Equality+import Fusion import Marking (Classified, SecurityMarking) import SecureBlock (SecureBlock) import qualified SecureBlock@@ -47,7 +48,7 @@ rnf = SecureBlock.toNormalForm . unBlockN instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where- zero = create (const zero)+ zero = BlockN.replicate zero {-# INLINE zero #-} (.+) = BlockN.zipWith (.+) {-# INLINE (.+) #-}@@ -58,20 +59,32 @@ 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+ #-}+ index :: PrimType a => BlockN marking n a -> Offset a -> a index = SecureBlock.index . unBlockN -#ifdef ML_KEM_TESTING replicate :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => a -> BlockN marking n a replicate = create . const +#ifdef ML_KEM_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 $ runST $ do- mb <- new (Proxy :: Proxy marking)- go mb 0 elems- unsafeFreeze mb+ | otherwise = Just $ runNew (Proxy :: Proxy marking) $ \mb -> go mb 0 elems where !sz = fromIntegral $ natVal (Proxy :: Proxy n) @@ -86,29 +99,69 @@ create :: forall marking n ty. (Classified marking, KnownNat n, PrimType ty) => (Offset ty -> ty) -> BlockN marking n ty-create initializer = BlockN $ SecureBlock.create (CountOf sz) initializer- where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+create initializer = runNew (Proxy :: Proxy marking) $ 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) = create $ \(Offset i) -> f (SecureBlock.index a (Offset i))-{-# INLINE map #-}+{-# INLINE [2] map #-} -iterModify :: (PrimType ty, PrimMonad prim)+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 #-}++iterModify :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim) => (ty -> ty) -> MutableBlockN marking n ty (PrimState prim) -> prim ()-iterModify f = SecureBlock.iterModify f . unMutableBlockN+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) = create $ \(Offset i) -> f (SecureBlock.index a (Offset i)) (SecureBlock.index b (Offset i))-{-# INLINE zipWith #-}+{-# 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 f a b = runContext (seqContext b (iterMapIxContext g (thawContext a)))+ where g i x = f x (index b $ Offset i)+{-# INLINE zipWithEndoL #-}+ unsafeCast :: BlockN marking n a -> SecureBlock marking b unsafeCast = SecureBlock.unsafeCast . unBlockN @@ -123,8 +176,62 @@ 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))+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 #-}++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 -> 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+ #-}
src/Crypto/PubKey/ML_KEM.hs view
@@ -12,7 +12,8 @@ module Crypto.PubKey.ML_KEM ( EncapsulationKey, DecapsulationKey, Ciphertext, SharedSecret -- * Operations- , generate, generateWith, encapsulate, encapsulateWith, decapsulate+ , generate, generateOpen, generateWith, encapsulate, encapsulateWith+ , decapsulate -- * Parameter sets , ParamSet, ML_KEM_512, ML_KEM_768, ML_KEM_1024 -- * Conversions and checks@@ -22,7 +23,7 @@ import Crypto.Random -import Data.ByteArray (ByteArrayAccess, ScrubbedBytes)+import Data.ByteArray (ByteArray, ByteArrayAccess, ScrubbedBytes) import qualified Data.ByteArray as B import Internal@@ -52,6 +53,20 @@ let d = B.takeView seed 32 z = B.drop 32 seed return (Internal.keyGen p d z)++-- | Generate a random seed (d, z) and the expanded key pair, returning+-- everything. This is Algorithm 19b introduced in Section 7 of+-- <https://www.rfc-editor.org/rfc/rfc9935 RFC 9935>.+--+-- Later use 'generateWith' to re-expand a seed value (d, z) that has been+-- recovered from storage.+generateOpen :: (ParamSet a, ByteArray d, ByteArray z, MonadRandom m)+ => proxy a -> m (EncapsulationKey a, DecapsulationKey a, d, z)+generateOpen p = do+ d <- getRandomBytes 32+ z <- getRandomBytes 32+ let (ek, dk) = Internal.keyGen p d (B.convert z)+ return (ek, dk, d, z) -- | Generate an ML-KEM key pair from the specified seed (d, z). Length of -- inputs must be 32 bytes.
+ src/Equality.hs view
@@ -0,0 +1,25 @@+-- |+-- Module : Equality+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Generate a constraint @'PrimSize' a ~ 'PrimSize' b@+--+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE TypeOperators #-}+module Equality+ ( EqPrimSize, ensureEqPrimSize+ ) where++import Data.Type.Equality++import Base++type EqPrimSize a b = (PrimType a, PrimType b, PrimSize a ~ PrimSize b)++eqPrimSize :: PrimSize a ~ PrimSize b => k a b -> PrimSize a :~: PrimSize b+eqPrimSize _ = Refl++ensureEqPrimSize :: EqPrimSize a b => k a b -> c -> c+ensureEqPrimSize op = case eqPrimSize op of Refl -> id
+ src/Fusion.hs view
@@ -0,0 +1,102 @@+-- |+-- Module : Fusion+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Infrastructure to decrease intermediate allocations and prefer in-place+-- mutation when possible+--+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE TypeFamilyDependencies #-}+module Fusion+ ( Fusion(..), MapF(..)+ , Context, runContext, newContext, thawContext, mapContext, modifyContext+ , foldContext, seqContext+ ) where++import Control.Monad ( forM_, (>=>) )+import Control.Monad.ST++-- class of values that can be mutated in the ST monad+class Fusion a where+ type Mut a s = mut | mut -> a+ newF :: ST s (Mut a s)+ thawF :: a -> ST s (Mut a s)+ unsafeFreezeF :: Mut a s -> ST s a++-- a transformation step in the fusion pipeline, with two implementations+-- provided: one that operates on an existing mutation context, and one that+-- initiates a new context from the input+data MapF a b = MapF+ { mapUpdate :: forall s. Mut a s -> ST s (Mut b s)+ , mapInit :: forall s. a -> ST s (Mut b s)+ }++-- MapF is almost a category except for the 'Fusion' constraint on objects+--+-- idMapF :: Fusion a => MapF a a+-- idMapF = MapF { mapUpdate = pure, mapInit = thawF }++composeMapF :: MapF b c -> MapF a b -> MapF a c+composeMapF m2 m1 = MapF+ { mapUpdate = mapUpdate m1 >=> mapUpdate m2+ , mapInit = mapInit m1 >=> mapUpdate m2+ }++-- fusion context+newtype Context a = Context (forall s. ST s (Mut a s))++newContext :: Fusion a => Context a+newContext = Context newF++thawContext :: Fusion a => a -> Context a+thawContext a = Context $ thawF a+{-# INLINE [0] thawContext #-}++modifyContext :: (forall s. Mut a s -> ST s ()) -> Context a -> Context a+modifyContext f = bindContext $ \ma -> f ma >> return ma++mapContext :: MapF a b -> Context a -> Context b+mapContext m = bindContext (mapUpdate m)+{-# INLINE [0] mapContext #-}++initContext :: MapF a b -> a -> Context b+initContext m a = Context $ mapInit m a++bindContext :: (forall s. Mut a s -> ST s (Mut b s)) -> Context a -> Context b+bindContext f (Context ctx) = Context $ ctx >>= f++foldContext :: Foldable t => (forall s. b -> Mut a s -> ST s ()) -> Context a -> t b -> Context a+foldContext f c bs = modifyContext (\ma -> forM_ bs $ \b -> f b ma) c++runContext :: Fusion a => Context a -> a+runContext (Context ctx) = runST (ctx >>= unsafeFreezeF)+{-# INLINE [0] runContext #-}++seqContext :: a -> Context b -> Context b+seqContext = seq+{-# INLINE [0] seqContext #-}+++-- Fusion rules+--+-- "thawContext/runContext" is the canonical optimization that eliminates an+-- allocation + value copy. Instead, it sequences two transformations on the+-- same mutation context.+--+-- "mapContext/seqContext" moves strictness annotations upstream so that they+-- do not prevent other rules from firing.+--+-- "mapContext/mapContext" is not strictly needed: the function is ultimately+-- inlined to the same code. But we keep it so that simplifications fire early+-- and do not wait for the final phase.+--+-- "mapContext/thawContext" is the rule that invokes mapInit instead of copying+-- the input and calling mapUpdate.++{-# RULES+"thawContext/runContext" [~0] forall c. thawContext (runContext c) = c+"mapContext/seqContext" [~0] forall a m c. mapContext m (seqContext a c) = seqContext a (mapContext m c)+"mapContext/mapContext" [~0] forall m1 m2 c. mapContext m2 (mapContext m1 c) = mapContext (composeMapF m2 m1) c+"mapContext/thawContext" [1] forall m a. mapContext m (thawContext a) = initContext m a+ #-}
src/K_PKE.hs view
@@ -29,7 +29,7 @@ import Builder (Builder) import Iterate import Marking (SecurityMarking(..), Leak(..))-import Vector (Vector)+import Vector (Vector, dot) import qualified Auxiliary as Aux import qualified Crypto import qualified Builder@@ -49,10 +49,11 @@ class Leak t => LeakVec vec t where leakVec :: vec (t Sec) -> vec (t Pub)- leakVec = unsafeCoerce -instance LeakVec (Vector k) Tq-instance LeakVec (Vector k) Rq+instance LeakVec (Vector k) Tq where+ leakVec = unsafeCoerce+instance LeakVec (Vector k) Rq where+ leakVec = fmap unsafeCoerce -- for rule "mapVector/zipWith" newtype DecryptionKey (k :: Nat) = DecryptionKey { dkS :: Vector k (Tq Sec) } data EncryptionKey (k :: Nat) = EncryptionKey { ekT :: Vector k (Tq Pub), ekRho :: Bytes, ekA :: Vector k (Vector k (Tq Pub)) }@@ -104,9 +105,11 @@ createVector :: (KnownNat k, ByteArrayAccess s) => Word -> s -> Int -> Vector k (Rq Sec) createVector !eta !s !j = Vector.create $ \(Offset i) -> sample eta s (i + j)+{-# INLINE createVector #-} sample :: ByteArrayAccess s => Word -> s -> Int -> Rq Sec sample eta s = Aux.samplePolyCBD eta . Crypto.prf eta s . fromIntegral+{-# INLINE sample #-} -- Uses randomness to generate an encryption key and a corresponding decryption key keyGen :: (KnownNat k, ByteArrayAccess d) => Params k -> d -> (EncryptionKey k, DecryptionKey k)@@ -136,7 +139,7 @@ yy = Aux.ntt <$> y u = leakVec $ (Aux.nttInv <$> Matrix.muly aa yy) .+ e1 mu = Aux.rdecompress 1 (Aux.byteDecode1 m)- v = Aux.nttInv (tt `Matrix.mulz` yy) .+ e2 ..+ mu+ v = Aux.nttInv (tt `dot` yy) .+ e2 ..+ mu c1 = Vector.concatMap (Aux.byteEncode du . Aux.rcompress du) u c2 = Aux.byteEncode dv (Aux.rcompress dv v) @@ -148,5 +151,5 @@ c2 = B.view c (32 * du * k) (32 * dv) u' = Vector.create $ \(Offset i) -> Aux.rdecompress du . Aux.byteDecode du $ B.view c (32 * du * i) (32 * du) :: Rq Pub v' = Aux.rdecompress dv (Aux.byteDecode dv c2)- w = v' ..- Aux.nttInv ((Aux.ntt <$> u') `Matrix.mulz` dkS dk)+ w = v' ..- Aux.nttInv ((Aux.ntt <$> u') `dot` dkS dk) m = Aux.byteEncode1 (Aux.rcompress 1 w)
src/Marking.hs view
@@ -55,9 +55,8 @@ class Classified (marking :: SecurityMarking) where type SecureBlock marking = (block :: Type -> Type) | block -> marking - create :: PrimType ty => CountOf ty -> (Offset ty -> ty) -> SecureBlock marking ty new :: (PrimType ty, PrimMonad prim) => proxy marking -> CountOf ty -> prim (MutableBlock ty (PrimState prim))- thaw :: (PrimType ty, PrimMonad m) => SecureBlock marking ty -> m (MutableBlock ty (PrimState m))+ thaw :: PrimMonad m => SecureBlock marking ty -> m (MutableBlock ty (PrimState m)) unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty) #ifdef ML_KEM_TESTING@@ -74,8 +73,6 @@ instance Classified Pub where type SecureBlock Pub = Block - create = Block.create- {-# INLINE create #-} new _ = Block.new thaw = Block.thaw unsafeFreeze = Block.unsafeFreeze@@ -95,8 +92,6 @@ instance Classified Sec where type SecureBlock Sec = ScrubbedBlock - create = ScrubbedBlock.create- {-# INLINE create #-} new _ = ScrubbedBlock.new thaw = ScrubbedBlock.thaw unsafeFreeze = ScrubbedBlock.unsafeFreeze
src/Math.hs view
@@ -10,11 +10,16 @@ -- typically a public operand (left) with a secret operand (right), producing a -- secret output. --+{-# LANGUAGE CPP #-} {-# LANGUAGE MultiParamTypeClasses #-} module Math ( Add(..), Mul(..), MulAdd(..), BiMul(..), BiMulAdd(..) ) where +#if !(MIN_VERSION_base(4,20,0))+import Data.List (foldl')+#endif+ infixl 7 .* infixr 7 ..* infixl 6 .+, .-@@ -37,5 +42,14 @@ (..*) :: b -> a -> a class BiMul b a => BiMulAdd b a where+ {-# MINIMAL biMulAdd | biMulFold #-}+ -- invariant: biMulAdd a b c == a ..* b .+ c biMulAdd :: b -> a -> a -> a+ biMulAdd b a x = biMulFold x [(b, a)]+ {-# INLINE biMulAdd #-}++ -- repeated biMulAdd+ biMulFold :: Foldable t => a -> t (b, a) -> a+ biMulFold = foldl' $ \c (b, a) -> biMulAdd b a c+ {-# INLINE biMulFold #-}
src/Matrix.hs view
@@ -4,13 +4,12 @@ -- Copyright : (c) 2025 Olivier Chéron -- -- A matrix here is simply a vector of vectors. The module also implements--- 'mulz' as dot product and two utility functions 'mulw' and 'muly' that--- multiply a matrix and a vector.+-- two utility functions 'mulw' and 'muly' that multiply a matrix and a vector. -- {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} module Matrix- ( create, mulw, muly, mulz+ ( create, mulw, muly #ifdef ML_KEM_TESTING , transpose #endif@@ -30,13 +29,12 @@ mulw :: (KnownNat n, BiMulAdd b a) => Vector m (Vector n b) -> Vector m a -> Vector n a -> Vector n a mulw a !u !b = Vector.create $ \(Offset i) ->- Vector.foldIndexWith (\c (Offset j) vu -> biMulAdd (index a (Offset i) (Offset j)) vu c) (Vector.index b (Offset i)) u+ Vector.biMulFoldIndexWith (\(Offset j) vu -> (index a (Offset i) (Offset j), vu)) (Vector.index b (Offset i)) u+{-# INLINE mulw #-} muly :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a-muly a !u = fmap (`mulz` u) a--mulz :: BiMulAdd b a => Vector n b -> Vector n a -> a-mulz = Vector.fold1ZipWith (\c a b -> biMulAdd a b c) (..*)+muly a !u = fmap (`Vector.dot` u) a+{-# INLINE muly #-} #ifdef ML_KEM_TESTING transpose :: (KnownNat m, KnownNat n) => Vector m (Vector n ty) -> Vector n (Vector m ty)
src/ScrubbedBlock.hs view
@@ -7,11 +7,10 @@ -- 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+ ( ScrubbedBlock, foldZipWith, ScrubbedBlock.length , new, thaw, unsafeFreeze ) where @@ -23,7 +22,7 @@ import Data.Word import Base-import Block (Block, MutableBlock, blockWrite)+import Block (Block, MutableBlock) import qualified Block import GHC.Base (IO(IO), Int(I#), setByteArray#)@@ -32,20 +31,6 @@ 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) =@@ -58,13 +43,8 @@ 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 #-}+thaw :: PrimMonad m => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m))+thaw (ScrubbedBlock b) = Block.thawPinned b -- always pinned unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty) unsafeFreeze mb = Block.unsafeFreeze mb >>= scrubbed
src/SecureBlock.hs view
@@ -8,20 +8,17 @@ -- {-# LANGUAGE CPP #-} module SecureBlock- ( SecureBlock, create, index, iterModify, new, thaw+ ( SecureBlock, index, new, thaw , unsafeCast, unsafeFreeze, toNormalForm #ifdef ML_KEM_TESTING , eq, Marking.showsPrec, toList, SecureBlock.length #endif ) where +#ifdef ML_KEM_TESTING import Base-import Block (MutableBlock)+#endif import Marking-import qualified Block--iterModify :: (PrimType ty, PrimMonad prim) => (ty -> ty) -> MutableBlock ty (PrimState prim) -> prim ()-iterModify = Block.iterModify #ifdef ML_KEM_TESTING length :: (Classified marking, PrimType ty) => SecureBlock marking ty -> CountOf ty
src/Vector.hs view
@@ -12,11 +12,11 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ScopedTypeVariables #-} module Vector- ( Vector, Vector.concatMap- , Vector.fold1ZipWith, Vector.foldIndexWith, Vector.toNormalForm+ ( Vector, Vector.concatMap, Vector.dot+ , Vector.fold1ZipWith, Vector.biMulFoldIndexWith, Vector.toNormalForm , Vector.create, Vector.index #ifdef ML_KEM_TESTING- , Vector.replicateM, Vector.zipWith+ , Vector.replicateM #endif ) where @@ -87,19 +87,19 @@ genericCreate :: forall n a a'. KnownNat n => (Offset a' -> a) -> Vector n a genericCreate f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i)) where !sz = fromIntegral $ natVal (Proxy :: Proxy n)-{-# INLINE [1] genericCreate #-}+{-# INLINE [0] genericCreate #-} genericCreateZipLeft :: KnownNat n => (a -> b -> c) -> (Offset a' -> a) -> Vector n b -> Vector n c genericCreateZipLeft f g a = genericCreate $ \off@(Offset i) -> f (g off) (index a (Offset i))-{-# INLINE [1] genericCreateZipLeft #-}+{-# INLINE [0] genericCreateZipLeft #-} genericCreateZipRight :: KnownNat n => (a -> b -> c) -> (Offset b' -> b) -> Vector n a -> Vector n c genericCreateZipRight f g a = genericCreate $ \off@(Offset i) -> f (index a (Offset i)) (g off)-{-# INLINE [1] genericCreateZipRight #-}+{-# INLINE [0] genericCreateZipRight #-} mapVector :: (a -> b) -> Vector n a -> Vector n b mapVector f = Vector <$> arrayMap f . unVector-{-# INLINE [1] mapVector #-}+{-# INLINE [0] mapVector #-} arrayIndex :: Array a -> Offset a -> a #ifdef ML_KEM_TESTING@@ -130,7 +130,7 @@ f (arrayIndex a (Offset i)) (arrayIndex b (Offset i)) where CountOf sa = arrayLength a-{-# INLINE [1] zipWith #-}+{-# INLINE [0] zipWith #-} fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c fold1ZipWith f g (Vector a) (Vector !b) =@@ -141,23 +141,35 @@ 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)+biMulFoldIndexWith :: BiMulAdd b a => (Offset ty -> t -> (b, a)) -> a -> Vector n t -> a+biMulFoldIndexWith f c (Vector a) =+ biMulFold c (map g $ offsets sa) where- g x i = f x (Offset i) (arrayIndex a (Offset i))+ g i = f (Offset i) (arrayIndex a (Offset i)) CountOf !sa = arrayLength a-{-# INLINE foldIndexWith #-}+{-# INLINE biMulFoldIndexWith #-} +dot :: BiMulAdd b a => Vector n b -> Vector n a -> a+dot (Vector b) (Vector a) =+ biMulFold (arrayIndex b 0 ..* arrayIndex a 0) (map g $ offsetsFrom 1 sb)+ where+ g i = (arrayIndex b (Offset i), arrayIndex a (Offset i))+ CountOf !sb = arrayLength b+{-# INLINE dot #-}+ toNormalForm :: NFData a => Vector n a -> () 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-"mapVector/genericCreate" [2] forall f g. mapVector f (genericCreate g) = genericCreate (f . g)-"zipWith/genericCreate left" [2] forall f g a. Vector.zipWith f (genericCreate g) a = genericCreateZipLeft f g a-"zipWith/genericCreate right" [2] forall f g a. Vector.zipWith f a (genericCreate g) = genericCreateZipRight f g a-"genericCreateZipLeft/genericCreate" [2] forall f g h. genericCreateZipLeft f g (genericCreate h) = genericCreate $ \(Offset i) -> f (g (Offset i)) (h (Offset i))-"genericCreateZipRight/genericCreate" [2] forall f g h. genericCreateZipRight f g (genericCreate h) = genericCreate $ \(Offset i) -> f (h (Offset i)) (g (Offset i))-"zipWith/mapVector left" [2] forall f g a. Vector.zipWith f (mapVector g a) = Vector.zipWith (f . g) a-"zipWith/mapVector right" [2] forall f g a b. Vector.zipWith f a (mapVector g b) = Vector.zipWith (\aa bb -> f aa (g bb)) a b+"mapVector/mapVector" [~0] forall f g a. mapVector f (mapVector g a) = mapVector (f . g) a+"mapVector/genericCreate" [~0] forall f g. mapVector f (genericCreate g) = genericCreate (f . g)+"zipWith/genericCreate left" [~0] forall f g b. Vector.zipWith f (genericCreate g) b = genericCreateZipLeft f g b+"zipWith/genericCreate right" [~0] forall f g a. Vector.zipWith f a (genericCreate g) = genericCreateZipRight f g a+"genericCreateZipLeft/mapVector" [~0] forall f g h b. genericCreateZipLeft f g (mapVector h b) = genericCreateZipLeft (\aa bb -> f aa (h bb)) g b+"genericCreateZipRight/mapVector" [~0] forall f g h a. genericCreateZipRight f g (mapVector h a) = genericCreateZipRight (f . h) g a+"genericCreateZipLeft/genericCreate" [~0] forall f g h. genericCreateZipLeft f g (genericCreate h) = genericCreate $ \(Offset i) -> f (g (Offset i)) (h (Offset i))+"genericCreateZipRight/genericCreate" [~0] forall f g h. genericCreateZipRight f g (genericCreate h) = genericCreate $ \(Offset i) -> f (h (Offset i)) (g (Offset i))+"zipWith/mapVector left" [~0] forall f g a. Vector.zipWith f (mapVector g a) = Vector.zipWith (f . g) a+"zipWith/mapVector right" [~0] forall f g a b. Vector.zipWith f a (mapVector g b) = Vector.zipWith (\aa bb -> f aa (g bb)) a b+"mapVector/zipWith" [~0] forall f g a b. mapVector f (Vector.zipWith g a b) = Vector.zipWith (\aa bb -> f (g aa bb)) a b #-}
tests/Tests.hs view
@@ -93,7 +93,7 @@ newtype D = D Int deriving Show instance Arbitrary D where- arbitrary = D <$> choose (0, 12)+ arbitrary = sized $ \n -> D <$> choose (0, min n 12) data Dim = forall (n :: Nat). KnownNat n => Dim (Proxy n) @@ -101,7 +101,7 @@ show (Dim n) = show n instance Arbitrary Dim where- arbitrary = toDim <$> choose (1, 9)+ arbitrary = sized $ \n -> toDim <$> choose (1, min (1 + n) 9) toDim :: Int -> Dim toDim n = case someNatVal (fromIntegral n) of SomeNat p -> Dim p@@ -283,6 +283,9 @@ [ Just c === Lib.decode p (B.convert c :: Bytes) , Just kk === Lib.decode p (B.convert kk :: Bytes) ]+ , testProperty "generation from seed" $ \(P p) -> ioProperty $ do+ (ek, dk, d, z) <- Lib.generateOpen p+ return (Just (ek, dk) === Lib.generateWith p (d :: Bytes) (z :: Bytes)) , testProperty "toPublic" $ \(P p) -> ioProperty $ do (ek, dk) <- Lib.generate p return (ek === toPublic dk)@@ -469,6 +472,15 @@ , testProperty "double negation" $ \(Dim n) -> do a <- arbitraryVector n return (a === neg (neg a))+ , testProperty "dot product commutative" $ \(Dim n) -> do+ (u, v) <- (,) <$> arbitraryVector n <*> arbitraryVector n+ return (u `dot` v === v `dot` u)+ , testProperty "dot product distributive" $ \(Dim n) -> do+ (u, v, w) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n+ return $ conjoin+ [ u `dot` (v .+ w) === (u `dot` v) .+ (u `dot` w)+ , (u .+ v) `dot` w === (u `dot` w) .+ (v `dot` w)+ ] ] , testGroup "Matrix" [ testProperty "mulw distributive left" $ \(Dim n) (Dim m) -> do@@ -486,15 +498,6 @@ , testProperty "muly distributive right" $ \(Dim n) (Dim m) -> do (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector m return (a `muly` (u .+ v) === (a `muly` u) .+ (a `muly` v))- , testProperty "mulz commutative" $ \(Dim n) -> do- (u, v) <- (,) <$> arbitraryVector n <*> arbitraryVector n- return (u `mulz` v === v `mulz` u)- , testProperty "mulz distributive" $ \(Dim n) -> do- (u, v, w) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n- return $ conjoin- [ u `mulz` (v .+ w) === (u `mulz` v) .+ (u `mulz` w)- , (u .+ v) `mulz` w === (u `mulz` w) .+ (v `mulz` w)- ] ] #endif ]