packages feed

mlkem 0.1.0.0 → 0.1.1.0

raw patch · 19 files changed

+441/−197 lines, 19 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

CHANGELOG.md view
@@ -1,5 +1,12 @@ # Changelog for `mlkem` +## 0.1.1.0 - 2025-11-30++* Fixed encoding and decoding on big-endian architectures++* Fall back to bytes when the architecture does not support unaligned memory+  access+ ## 0.1.0.0 - 2025-11-02  * First version. Released on an unsuspecting world.
README.md view
@@ -1,5 +1,9 @@ # ML-KEM +[![BSD](https://img.shields.io/badge/License-BSD-blue)](https://en.wikipedia.org/wiki/BSD_licenses)+[![Haskell](https://img.shields.io/badge/Language-Haskell-lightgrey)](https://haskell.org/)+[![Hackage](https://img.shields.io/hackage/v/mlkem?label=Hackage&color=green)](https://hackage.haskell.org/package/mlkem)+ _Module-Lattice-based Key-Encapsulation Mechanism_ implemented in Haskell.  See [FIPS 203](https://csrc.nist.gov/pubs/fips/203/final).@@ -26,7 +30,8 @@ capture variables and move them to the heap.  This could theoretically include machine words containing secret information that would not then be destroyed. -Best performance is obtained with the LLVM code generator.+Best performance is obtained with the LLVM code generator.  On ARM, define+macro `__ARM_FEATURE_UNALIGNED` if unaligned access is supported by the target.  ## Testing
benchs/Bench.hs view
@@ -27,6 +27,8 @@     , bench "encapsulate (batch)" $ perBatchEnv (const setupEncap) (return . runEncap)     , bench "decapsulate" $ perRunEnv setupDecap (return . runDecap)     , bench "decapsulate (batch)" $ perBatchEnv (const setupDecap) (return . runDecap)+    , env (generate p) $ bench "encoding (public)" . nf (encodeDecode p) . fst+    , env (generate p) $ bench "encoding (private)" . nf (encodeDecode p) . snd     ]   where     gen32 = getRandomBytes 32 :: IO Bytes@@ -45,6 +47,13 @@         (ek, dk) <- generate p         (_, c) <- encapsulate ek         return (dk, c)++    encodeDecode :: (Encode obj, Decode obj, Eq (obj a), ParamSet a)+                 => proxy a -> obj a -> Bool+    encodeDecode prx k = decode prx (aligned $ encode k) == Just k++    aligned :: Bytes -> Bytes+    aligned = id  main :: IO () main = defaultMain
mlkem.cabal view
@@ -1,11 +1,11 @@ cabal-version: 2.2 --- This file has been generated from package.yaml by hpack version 0.38.0.+-- This file has been generated from package.yaml by hpack version 0.38.1. -- -- see: https://github.com/sol/hpack  name:           mlkem-version:        0.1.0.0+version:        0.1.1.0 synopsis:       Module-Lattice-based Key-Encapsulation Mechanism description:    Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM) implemented in                 Haskell.@@ -44,7 +44,9 @@       ByteArrayST       Crypto       Internal+      Iterate       K_PKE+      Machine       Marking       Math       Matrix@@ -57,7 +59,7 @@       Paths_mlkem   hs-source-dirs:       src-  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -O2+  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@@ -84,7 +86,7 @@       Paths_mlkem   hs-source-dirs:       tests-  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N+  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   build-depends:       aeson     , base >=4.7 && <5@@ -120,7 +122,9 @@       Crypto       Crypto.PubKey.ML_KEM       Internal+      Iterate       K_PKE+      Machine       Marking       Math       Matrix@@ -138,7 +142,7 @@   hs-source-dirs:       src       tests-  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -fno-ignore-asserts -threaded -rtsopts -with-rtsopts=-N+  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 -fno-ignore-asserts -threaded -rtsopts -with-rtsopts=-N   cpp-options: -DML_KEM_TESTING   build-depends:       aeson@@ -171,7 +175,7 @@       Paths_mlkem   hs-source-dirs:       benchs-  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N+  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
src/Auxiliary.hs view
@@ -31,7 +31,6 @@  import Data.ByteArray (ByteArrayAccess, Bytes, View) import qualified Data.ByteArray as B-import qualified Data.Memory.Endian as B  import Control.DeepSeq (NFData(..)) import Control.Monad@@ -46,10 +45,13 @@ import Foreign.Ptr (Ptr, plusPtr) import Foreign.Storable (pokeByteOff) +import Unsafe.Coerce+ import Block (blockIndex) import BlockN (BlockN, MutableBlockN) import Builder (Builder) import Crypto (BlockDigest)+import Machine import Marking (Classified, SecurityMarking(..), Leak(..)) import SecureBlock (SecureBlock) import SecureBytes (SecureBytes)@@ -176,9 +178,11 @@ -- adding noise to secret information. (..+) :: Rq Sec -> Rq Sec -> Rq Pub Rq a ..+ Rq b = Rq $ BlockN.zipWith (.+) a b+{-# NOINLINE (..+) #-}  (..-) :: Rq Pub -> Rq Sec -> Rq Sec Rq a ..- Rq b = Rq $ BlockN.zipWith (.-) a b+{-# NOINLINE (..-) #-}  instance Leak Rq @@ -205,14 +209,16 @@  instance Leak Tq -instance Classified marking => NFData (Tq marking) where+instance NFData (Tq marking) where     rnf (Tq a) = toNormalForm a  instance BiMul (Tq Pub) (Tq Sec) where     (..*) = multiplyNTTs+    {-# NOINLINE (..*) #-}  instance BiMulAdd (Tq Pub) (Tq Sec) where     biMulAdd = multiplyNTTsAdd+    {-# NOINLINE biMulAdd #-}  #ifdef ML_KEM_TESTING instance Mul (Tq Sec) where@@ -374,14 +380,14 @@ -- 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 rcompress :: Int -> Rq Sec -> BlockN Sec N Word16 #-}-{-# SPECIALIZE rcompress :: Int -> Rq Pub -> BlockN Pub N Word16 #-}+{-# SPECIALIZE NOINLINE rcompress :: Int -> Rq Sec -> BlockN Sec N Word16 #-}+{-# SPECIALIZE NOINLINE rcompress :: Int -> Rq Pub -> BlockN Pub N Word16 #-}  -- Decompress a polynomial with 𝑑 < 12 rdecompress :: Classified marking => Int -> BlockN marking N Word16 -> Rq marking rdecompress !d = Rq . BlockN.map (decompress d)-{-# SPECIALIZE rdecompress :: Int -> BlockN Sec N Word16 -> Rq Sec #-}-{-# SPECIALIZE rdecompress :: Int -> BlockN Pub N Word16 -> Rq Pub #-}+{-# SPECIALIZE NOINLINE rdecompress :: Int -> BlockN Sec N Word16 -> Rq Sec #-}+{-# SPECIALIZE NOINLINE rdecompress :: Int -> BlockN Pub N Word16 -> Rq Pub #-}  -- Generates a pseudorandom element of T𝑞 from a seed and two indices sampleNTT :: SecureBytes Pub -> Word8 -> Word8 -> Tq Pub@@ -390,7 +396,7 @@     runXof b (280 * 3) 0 0     Tq <$> BlockN.unsafeFreeze b   where-    runXof !b !xofLen !pos !j = case someNatVal (8 * fromIntegral xofLen) of+    runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of         SomeNat proxy -> do             let bytes = Crypto.unBlockDigest (doHash proxy)             loop b xofLen bytes pos j@@ -421,22 +427,14 @@         pokeByteOff d (len + 1) y     len = B.length seed -wordBits :: Int-wordBits = finiteBitSize (0 :: Word)--wordBytes :: Int-wordBytes = div wordBits 8--type WordLE = B.LE Word--peekWord :: Ptr WordLE -> ST s Word-peekWord p = B.unLE <$> ST.peek p+peekWord :: Ptr WordLE -> ST s WordM+peekWord p = fromLE <$> ST.peek p -peekWordPos :: Ptr WordLE -> BitPos -> ST s Word-peekWordPos a bp = B.unLE <$> ST.peekElemOff a (wordOff bp)+peekWordPos :: Ptr WordLE -> BitPos -> ST s WordM+peekWordPos a bp = fromLE <$> ST.peekElemOff a (wordOff bp) -pokeWordPos :: Ptr WordLE -> BitPos -> Word -> ST s ()-pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . B.LE+pokeWordPos :: Ptr WordLE -> BitPos -> WordM -> ST s ()+pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . toLE  newtype BitPos = BitPos Int @@ -457,10 +455,12 @@ nextPos requested (BitPos p) = (howMany, BitPos $ p + howMany)   where howMany = availPos requested (BitPos p) -getMask :: Int -> Word+getMask :: Int -> WordM getMask howMany     | howMany >= wordBits = maxBound     | otherwise = (1 `unsafeShiftL` howMany) - 1+    -- branch useful only when processing one byte at a time due to+    -- architecture not supporting unaligned memory access  -- Takes a seed as input and outputs a pseudorandom sample from the -- distribution D_eta@@ -487,15 +487,19 @@             getBits p bp' (acc + fromIntegral (popCount bits)) (j - howMany)  -- Encodes an array of 𝑑-bit integers into a byte array for 1 ≤ 𝑑 ≤ 12-byteEncode :: Classified marking => Int -> BlockN marking N Word16 -> Builder marking-byteEncode !d !f = Builder.create (32 * d) $ \b ->-    outer b zeroPos 0 0+byteEncode :: Int -> BlockN marking N Word16 -> Builder marking+byteEncode d f = Builder.create (32 * d) (runByteEncode d f)+{-# INLINE byteEncode #-}++runByteEncode :: Int -> BlockN marking N Word16 -> Ptr WordLE -> ST s ()+runByteEncode !d !f dst = outer dst zeroPos 0 0   where-    outer :: Ptr WordLE -> BitPos -> Word -> Int -> ST s ()+    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 #-} -    inner :: Ptr WordLE -> Int -> BitPos -> Word -> Word16 -> Int -> ST s ()+    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'@@ -506,15 +510,16 @@         o' = o .|. (x `unsafeShiftL` bitPos bp)         a' = a `unsafeShiftR` howMany         j' = j - howMany-{-# SPECIALIZE byteEncode :: Int -> BlockN Sec N Word16 -> Builder Sec #-}-{-# SPECIALIZE byteEncode :: Int -> BlockN Pub N Word16 -> Builder Pub #-}  -- Optimization of byteEncode when 𝑑=1 byteEncode1 :: BlockN Sec N Word16 -> Builder Sec-byteEncode1 !f = Builder.create 32 $ \b ->-    loop b 0 0+byteEncode1 !f = Builder.create 32 (runByteEncode1 f)+{-# INLINE byteEncode1 #-}++runByteEncode1 :: BlockN marking N Word16 -> Ptr WordLE -> ST s ()+runByteEncode1 !f dst = loop dst 0 0   where-    loop :: Ptr WordLE -> Word -> Int -> ST s ()+    loop :: Ptr WordLE -> WordM -> Int -> ST s ()     loop !b !o pos         | pos == n = return ()         | bitPos bp + 1 < wordBits = loop b o' (pos + 1)@@ -526,11 +531,15 @@         a = BlockN.index f (Offset pos)  -- byteEncode with 𝑑=12 after conversion from the field-byteEncode12 :: Classified marking => Tq marking -> Builder marking-byteEncode12 (Tq f) = byteEncode 12 $ BlockN.map (\(Zq x) -> x) f+byteEncode12 :: Tq marking -> Builder marking+byteEncode12 = byteEncode 12 . fromField+  where+    fromField :: Tq marking -> BlockN marking N Word16+    fromField (Tq f) = unsafeCoerce f+{-# INLINE byteEncode12 #-}  -- 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 :: 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)@@ -550,12 +559,12 @@                 j' = j + howMany             inner f p i bp' v' j' -    get :: Ptr WordLE -> BitPos -> Int -> ST s Word+    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 -> ba -> BlockN Sec N Word16 #-}-{-# SPECIALIZE byteDecode :: forall ba. ByteArrayAccess ba => Int -> ba -> BlockN Pub N Word16 #-}+{-# 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 #-} @@ -572,7 +581,7 @@         x <- peekWord p         inner f (p `plusPtr` wordBytes) x i 0 -    inner :: MutableBlockN Sec N Word16 s -> Ptr WordLE -> Word -> Int -> Int -> ST s ()+    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         | otherwise = do@@ -581,5 +590,7 @@             inner f 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 :: (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 #-}
src/Block.hs view
@@ -11,12 +11,9 @@ {-# LANGUAGE CPP #-} module Block     ( Block, MutableBlock, blockIndex, blockRead, blockWrite-    , create, foldZipWith, iterModify, Block.length, Block.map+    , create, foldZipWith, iterModify, Block.length     , Block.new, Block.newPinned, Block.thaw, Block.unsafeCast     , Block.unsafeFreeze, Block.unsafeThaw, Block.withMutablePtr-#ifdef ML_KEM_TESTING-    , Block.replicate-#endif     ) where  import Basement.Block (Block)@@ -30,8 +27,6 @@ import Control.Exception (assert) import Control.Monad.ST -import Data.Proxy- 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 ()@@ -58,11 +53,6 @@         | i .==# n = pure ()         | otherwise = Block.unsafeWrite mb i (initializer i) >> loop mb (i + 1) {-# INLINE create #-}--map :: (PrimType a, PrimType b) => (a -> b) -> Block a -> Block b-map f a = Block.create lenB (f . Block.unsafeIndex a . offsetCast Proxy)-  where !lenB = sizeCast (Proxy :: Proxy (a -> b)) (Block.length a)-{-# INLINE map #-}  iterModify :: (PrimType ty, PrimMonad prim)            => (ty -> ty)
src/BlockN.hs view
@@ -25,6 +25,10 @@ import Basement.PrimType import Basement.Types.OffsetSize +#ifdef ML_KEM_TESTING+import Control.Monad.ST+#endif+ import Data.Proxy  import Block (MutableBlock, blockRead, blockWrite)@@ -43,7 +47,7 @@     showsPrec d = SecureBlock.showsPrec d . unBlockN #endif -instance Classified marking => NormalForm (BlockN marking n a) where+instance NormalForm (BlockN marking n a) where     toNormalForm = SecureBlock.toNormalForm . unBlockN  instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where@@ -58,23 +62,28 @@  newtype MutableBlockN (marking :: SecurityMarking) (n :: Nat) a m = MutableBlockN { unMutableBlockN :: MutableBlock a m } -index :: (Classified marking, PrimType a) => BlockN marking n a -> Offset a -> a+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 = BlockN . SecureBlock.replicate sz-  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+replicate = create . const  fromList :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => [a] -> Maybe (BlockN marking n a) fromList elems-    | SecureBlock.length a == CountOf sz = Just (BlockN a)-    | otherwise = Nothing+    | Prelude.length elems /= sz = Nothing+    | otherwise = Just $ runST $ do+        mb <- new (Proxy :: Proxy marking)+        go mb 0 elems+        unsafeFreeze mb   where-    a = SecureBlock.fromList elems     !sz = fromIntegral $ natVal (Proxy :: Proxy n) -toList :: (Classified marking, PrimType a) => BlockN marking n a -> [a]+    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 @@ -85,8 +94,9 @@   where !sz = fromIntegral $ natVal (Proxy :: Proxy n) {-# INLINE create #-} -map :: (Classified marking, PrimType a, PrimType b) => (a -> b) -> BlockN marking n a -> BlockN marking n b-map f = BlockN . SecureBlock.map f . unBlockN+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 #-}  iterModify :: (PrimType ty, PrimMonad prim)@@ -96,14 +106,14 @@ iterModify f = SecureBlock.iterModify f . unMutableBlockN {-# INLINE iterModify #-} -zipWith :: (Classified ma, Classified mb, Classified mc, KnownNat n, PrimType a, PrimType b, PrimType c)+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 #-} -unsafeCast :: (Classified marking, PrimType b) => BlockN marking n a -> SecureBlock marking b+unsafeCast :: PrimType b => 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
src/Builder.hs view
@@ -5,7 +5,6 @@ -- -- Eliminate intermediate allocations when concatenating several buffers ---{-# LANGUAGE BangPatterns #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE KindSignatures #-} {-# LANGUAGE RankNTypes #-}@@ -26,6 +25,8 @@  import Foreign.Ptr (Ptr, castPtr, plusPtr) +import qualified GHC.Exts as Exts+ import Block (Block) import Marking (Classified, Leak(..), SecurityMarking(..)) import SecureBytes (SecureBytes)@@ -35,7 +36,7 @@  data Builder (marking :: SecurityMarking) = Builder     { builderLength :: {-# UNPACK #-} !Int-    , copyBuilderToPtr :: forall a s. Ptr a -> ST s ()+    , copyBuilderToPtr :: forall s. Ptr Word8 -> ST s ()     }  instance Semigroup (Builder marking) where@@ -47,12 +48,13 @@  instance Monoid (Builder marking) where     mempty = empty-    mconcat builders = create n (`loop` builders)+    mconcat builders = create (getSize (Exts.inline builders)) $+        go (Exts.inline builders)       where-        n = getSum $ Prelude.mconcat $ map (Sum . builderLength) builders-        loop !_ [] = return ()-        loop !p (b : bs) =-            copyBuilderToPtr b p >> loop (p `plusPtr` builderLength b) bs+        getSize = getSum . Prelude.mconcat . map (Sum . builderLength)+        go = foldr c (const $ return ())+        c b k = Exts.oneShot $ \p -> copyBuilderToPtr b p >> k (p `plusPtr` builderLength b)+    {-# INLINE mconcat #-}  instance Leak Builder
src/ByteArrayST.hs view
@@ -25,8 +25,8 @@ unsafeCreate sz f = B.unsafeCreate sz (stToIO . f) {-# INLINE unsafeCreate #-} -withByteArray :: ByteArrayAccess ba => ba -> (forall s. Ptr p -> ST s a) -> ST t a-withByteArray b f = unsafeIOToST $ B.withByteArray b (stToIO . f)+withByteArray :: ByteArrayAccess ba => ba -> (Ptr p -> ST s a) -> ST s a+withByteArray b f = unsafeIOToST $ B.withByteArray b (unsafeSTToIO . f) {-# INLINE withByteArray #-}  peek :: Storable a => Ptr a -> ST s a
src/Crypto.hs view
@@ -12,7 +12,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE ScopedTypeVariables #-} module Crypto-    ( ConstEqW(..), BoolW, andW, toBool, constSelectBytes, snoc, append+    ( ConstEqW(..), BoolW, andW, toBool, constSelectBytes, snoc, append, eq     , prf, h, j, g, BlockDigest, unBlockDigest, hashToBlock     ) where @@ -37,6 +37,7 @@  import Block (Block) import Builder (Builder)+import Machine import ScrubbedBlock (ScrubbedBlock) import Vector (Vector) import qualified Block@@ -104,10 +105,7 @@ bytesConstEqW :: (ByteArrayAccess bs1, ByteArrayAccess bs2) => bs1 -> bs2 -> BoolW bytesConstEqW a b     | B.length a /= B.length b = falseW-    | otherwise =-        assertMultW (B.length a) $-        assertMultW (B.length b) $-        foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b+    | otherwise = foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b  foldZipWith :: (ByteArrayAccess bs1, ByteArrayAccess bs2)             => (c -> Word -> Word -> c) -> c -> bs1 -> bs2 -> c@@ -170,8 +168,25 @@     nb = B.length b {-# INLINE append #-} +eq :: (ByteArrayAccess a, ByteArrayAccess b) => a -> b -> Bool+eq a b = assert (sa == sb) $ assertMultM sa $ assertMultM sb $+    runST $ ST.withByteArray a $ \pa -> ST.withByteArray b $ \pb ->+        loop (pa :: Ptr WordM) (pb :: Ptr WordM) 0+  where+    !sa = B.length a+    !sb = B.length b++    loop !pa !pb i+        | i == sa = return True+        | otherwise = do+            va <- ST.peek pa+            vb <- ST.peek pb+            if va == vb+                then loop (pa `plusPtr` wordBytes) (pb `plusPtr` wordBytes) (i + wordBytes)+                else return False+ prf :: ByteArrayAccess s => Word -> s -> Word8 -> ScrubbedBytes-prf !eta s !b = case someNatVal (8 * 64 * fromIntegral eta) of+prf !eta s !b = case someNatVal (fromIntegral (8 * 64 * eta)) of     SomeNat proxy -> unDigest (doHash proxy)   where     doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen)
+ src/Iterate.hs view
@@ -0,0 +1,47 @@+-- |+-- Module      : Iterate+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- @'offsetsFrom' x y@ is similar to @[x .. y - 1]@ but does not verify that+-- @x <= y@.  This removes an unnecessary branch.+--+{-# LANGUAGE MagicHash #-}+module Iterate+  ( offsets, offsetsFrom+  )+where++import Control.Exception (assert)++import GHC.Exts++{-# INLINE offsets #-}+offsets :: Int -> [Int]+offsets = offsetsFrom 0++{-# INLINE offsetsFrom #-}+offsetsFrom :: Int -> Int -> [Int]+offsetsFrom xx@(I# x) yy@(I# y) = assert (xx <= yy) (indices x y)++-- Implementation is derived from enumFromTo, this acts as good producer for+-- list fusion.++{-# RULES+"indices"        [~1] forall x y. indices x y = build (\c n -> indicesFB c n x y)+"indicesList"    [1] indicesFB (:) [] = indices+ #-}++indices :: Int# -> Int# -> [Int]+indices x0 y = go x0+  where+    go x = if isTrue# (x <# y) then I# x : go (x +# 1#) else []+{-# NOINLINE [1] indices #-}++indicesFB :: (Int -> r -> r) -> r -> Int# -> Int# -> r+indicesFB c n x0 y = go x0+  where+    go x = if isTrue# (x <# y) then I# x `c` go (x +# 1#) else n+    -- Be very careful not to have more than one "c" so that when indicesFB is+    -- inlined we can inline whatever is bound to "c"+{-# INLINE [0] indicesFB #-}
src/K_PKE.hs view
@@ -29,6 +29,7 @@  import Auxiliary (Rq, Tq, (..+), (..-)) import Builder (Builder)+import Iterate import Marking (SecurityMarking(..), Leak(..)) import Vector (Vector) import qualified Auxiliary as Aux@@ -82,8 +83,8 @@         !rho = B.convert $ B.view input (384 * k) 32         elem384 off = Builder.run (Aux.byteEncode12 (Vector.index tt off))     -- modulus check:-    forM_ [0 .. k - 1] $ \i -> guard (elem384 (Offset i) `B.eq` view384 (Offset i))-    let aa  = Matrix.create $ \(Offset i) (Offset j) -> Aux.sampleNTT rho (fromIntegral j) (fromIntegral i)+    forM_ (offsets k) $ \i -> guard (elem384 (Offset i) `Crypto.eq` view384 (Offset i))+    let aa = createMatrix rho     Just EncryptionKey { ekT = tt, ekRho = rho, ekA = aa }   where     k = dimension params@@ -99,15 +100,25 @@   where     view384 (Offset i) = B.view input (384 * i) 384 +createMatrix :: KnownNat k => Bytes -> Vector k (Vector k (Tq Pub))+createMatrix !rho = Matrix.create $ \(Offset i) (Offset j) ->+    Aux.sampleNTT rho (fromIntegral j) (fromIntegral i)++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)++sample :: ByteArrayAccess s => Word -> s -> Int -> Rq Sec+sample eta s = Aux.samplePolyCBD eta . Crypto.prf eta s . fromIntegral+ -- Uses randomness to generate an encryption key and a corresponding decryption key keyGen :: (KnownNat k, ByteArrayAccess d) => Params k -> d -> (EncryptionKey k, DecryptionKey k) keyGen params@Params{..} d = (ek, dk)   where     k   = dimension params     (rho, sigma) = Crypto.g (Crypto.snoc d (fromIntegral k))-    aa  = Matrix.create $ \(Offset i) (Offset j) -> Aux.sampleNTT rho (fromIntegral j) (fromIntegral i)-    s   = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta1 (Crypto.prf eta1 sigma $ fromIntegral i)-    e   = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta1 (Crypto.prf eta1 sigma $ fromIntegral (k + i))+    aa  = createMatrix rho+    s   = createVector eta1 sigma 0+    e   = createVector eta1 sigma k     !ss = Aux.ntt <$> s     ee  = Aux.ntt <$> e     !tt = leakVec $ Matrix.mulw aa ss ee@@ -121,9 +132,9 @@     k   = dimension params     tt  = ekT ek     aa  = ekA ek-    y   = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta1 (Crypto.prf eta1 r $ fromIntegral i)-    e1  = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta2 (Crypto.prf eta2 r $ fromIntegral (k + i))-    e2  = Aux.samplePolyCBD eta2 (Crypto.prf eta2 r $ fromIntegral (2 * k))+    y   = createVector eta1 r 0+    e1  = createVector eta2 r k+    e2  = sample eta2 r (2 * k)     yy  = Aux.ntt <$> y     u   = leakVec $ (Aux.nttInv <$> Matrix.muly aa yy) .+ e1     mu  = Aux.rdecompress 1 (Aux.byteDecode1 m)
+ src/Machine.hs view
@@ -0,0 +1,92 @@+-- |+-- Module      : Machine+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Architecture-dependent utilities, to read/write unaligned machine words+-- in little-endian order+--+{-# LANGUAGE CPP #-}+module Machine+    ( WordM, WordLE, assertMultM, fromLE, toLE, wordBits, wordBytes+    ) where++#include "MachDeps.h"++-- Taken from `bytestring`, a list of architectures known to accept+-- unaligned loads and stores+#if defined(i386_HOST_ARCH) || defined(x86_64_HOST_ARCH)          \+    || ((defined(arm_HOST_ARCH) || defined(aarch64_HOST_ARCH))    \+         && defined(__ARM_FEATURE_UNALIGNED))                     \+    || defined(powerpc_HOST_ARCH) || defined(powerpc64_HOST_ARCH) \+    || defined(powerpc64le_HOST_ARCH)+#define MLKEM_ALLOW_UNALIGNED_OP 1++-- Little-endian conversion in `basement` and `memory` 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)+#define MLKEM_FORCE_LITTLE_ENDIAN_ARCH 1+#endif++#endif++import Control.Exception (assert)++#ifdef MLKEM_ALLOW_UNALIGNED_OP+import qualified Data.Memory.Endian as B+#endif++import Data.Bits+import Data.Word++#ifdef MLKEM_ALLOW_UNALIGNED_OP++-- our preferred word size+#if WORD_SIZE_IN_BITS == 64+type WordM = Word64+#else+type WordM = Word32+#endif++type WordLE = B.LE WordM++fromLE :: WordLE -> WordM+#ifdef MLKEM_FORCE_LITTLE_ENDIAN_ARCH+fromLE = B.unLE  -- unwrap constructor with no byte swapping+#else+fromLE = B.fromLE  -- byte swap if necessary+#endif++toLE :: WordM -> WordLE+#ifdef MLKEM_FORCE_LITTLE_ENDIAN_ARCH+toLE = B.LE  -- wrap constructor with no byte swapping+#else+toLE = B.toLE  -- byte swap if necessary+#endif++#else++-- unaligned memory access is not allowed so we fallback to one byte at a time+-- and endianness does not matter++type WordM = Word8+type WordLE = WordM++fromLE :: WordLE -> WordM+fromLE = id++toLE :: WordM -> WordLE+toLE = id++#endif++wordBits :: Int+wordBits = finiteBitSize (0 :: WordM)++wordBytes :: Int+wordBytes = div wordBits 8++assertMultM :: Int -> a -> a+assertMultM n = assert (n .&. mask == 0)+  where mask = wordBytes - 1
src/Marking.hs view
@@ -14,7 +14,11 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilyDependencies #-} module Marking-    ( SecurityMarking(..), Classified(..), Leak(..)+    ( SecurityMarking(..), Classified(..), Leak(..), index+    , Marking.toNormalForm, unsafeCast+#ifdef ML_KEM_TESTING+    , Marking.toList+#endif     ) where  import Basement.Monad@@ -58,20 +62,13 @@     type SecureBlock marking = (block :: Type -> Type) | block -> marking      create :: PrimType ty => CountOf ty -> (Offset ty -> ty) -> SecureBlock marking ty-    index :: PrimType ty => SecureBlock marking ty -> Offset ty -> ty-    map :: (PrimType a, PrimType b) => (a -> b) -> SecureBlock marking a -> SecureBlock marking b     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))-    unsafeCast :: PrimType b => SecureBlock marking a -> SecureBlock marking b     unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty) -    toNormalForm :: SecureBlock marking ty -> () #ifdef ML_KEM_TESTING     eq :: PrimType ty => SecureBlock marking ty -> SecureBlock marking ty -> Bool     showsPrec :: (PrimType ty, Show ty) => Int -> SecureBlock marking ty -> ShowS-    fromList :: PrimType ty => [ty] -> SecureBlock marking ty-    replicate :: PrimType ty => CountOf ty -> ty -> SecureBlock marking ty-    toList :: PrimType ty => SecureBlock marking ty -> [ty]     lengthBlock :: PrimType ty => SecureBlock marking ty -> CountOf ty #endif @@ -85,21 +82,13 @@      create = Block.create     {-# INLINE create #-}-    index = blockIndex-    map = Block.map-    {-# INLINE map #-}     new _ = Block.new     thaw = Block.thaw-    unsafeCast = Block.unsafeCast     unsafeFreeze = Block.unsafeFreeze -    toNormalForm = Basement.NormalForm.toNormalForm #ifdef ML_KEM_TESTING     eq = (==)     showsPrec = Prelude.showsPrec-    fromList = Basement.Compat.IsList.fromList-    replicate = Block.replicate-    toList = Basement.Compat.IsList.toList     lengthBlock = Block.length #endif @@ -114,21 +103,13 @@      create = ScrubbedBlock.create     {-# INLINE create #-}-    index = ScrubbedBlock.index-    map = ScrubbedBlock.map-    {-# INLINE map #-}     new _ = ScrubbedBlock.new     thaw = ScrubbedBlock.thaw-    unsafeCast = ScrubbedBlock.unsafeCast     unsafeFreeze = ScrubbedBlock.unsafeFreeze -    toNormalForm = Basement.NormalForm.toNormalForm #ifdef ML_KEM_TESTING     eq = (==)     showsPrec = Prelude.showsPrec-    fromList = ScrubbedBlock.fromList-    replicate = ScrubbedBlock.replicate-    toList = ScrubbedBlock.toList     lengthBlock = ScrubbedBlock.length #endif @@ -137,3 +118,27 @@     {-# INLINE unsafeCreate #-}     lengthBytes = B.length     copyByteArrayToPtr = B.copyByteArrayToPtr+++-- for some functions we use the fact that Block and SecureBlock have the same+-- representation and implementation++unwrap :: SecureBlock marking a -> Block a+unwrap = unsafeCoerce++wrap :: Block b -> SecureBlock marking b+wrap = unsafeCoerce++index :: PrimType ty => SecureBlock marking ty -> Offset ty -> ty+index = blockIndex . unwrap++#ifdef ML_KEM_TESTING+toList :: PrimType ty => SecureBlock marking ty -> [ty]+toList = Basement.Compat.IsList.toList . unwrap+#endif++toNormalForm :: SecureBlock marking ty -> ()+toNormalForm = Basement.NormalForm.toNormalForm . unwrap++unsafeCast :: PrimType b => SecureBlock marking a -> SecureBlock marking b+unsafeCast = wrap . Block.unsafeCast . unwrap
src/Math.hs view
@@ -15,7 +15,8 @@     ( Add(..), Mul(..), MulAdd(..), BiMul(..), BiMulAdd(..)     ) where -infixl 7 .*, ..*+infixl 7 .*+infixr 7 ..* infixl 6 .+, .-  class Add a where
src/ScrubbedBlock.hs view
@@ -9,16 +9,11 @@ -- {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-}-{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE UnboxedTuples #-} module ScrubbedBlock-    ( ScrubbedBlock, create, foldZipWith, index-    , ScrubbedBlock.length, ScrubbedBlock.map, new-    , thaw, unsafeCast, unsafeFreeze, Block.withMutablePtr-#ifdef ML_KEM_TESTING-    , ScrubbedBlock.fromList, ScrubbedBlock.replicate, ScrubbedBlock.toList-#endif+    ( ScrubbedBlock, create, foldZipWith, ScrubbedBlock.length+    , new, thaw, unsafeFreeze, Block.withMutablePtr     , erasePtr     ) where @@ -26,11 +21,7 @@ import Basement.Block.Mutable (mutableLengthBytes, unsafeCopyBytesRO) import Basement.Compat.Primitive -#ifdef ML_KEM_TESTING-import Basement.Compat.IsList-#endif import Basement.Monad-import Basement.NormalForm import Basement.PrimType import Basement.Types.OffsetSize @@ -41,20 +32,20 @@  import Foreign.Ptr (Ptr) -import Block (blockIndex, blockWrite)+import Block (blockWrite) import qualified Block  #if MIN_VERSION_base(4,19,0)-import GHC.Base (Int(I#), setAddrRange#)+import GHC.Base (setAddrRange#) import GHC.Exts (Ptr(Ptr)) #else import Data.Memory.PtrMethods (memSet) #endif-import GHC.Base (IO(IO), setByteArray#)-import GHC.Exts (getSizeofMutableByteArray#, mkWeak#)+import GHC.Base (IO(IO), Int(I#), setByteArray#)+import GHC.Exts (mkWeak#)  newtype ScrubbedBlock ty = ScrubbedBlock (Block ty)-    deriving (Eq, Show, NormalForm)+    deriving (Eq, Show)  create :: PrimType ty        => CountOf ty@@ -76,77 +67,47 @@     Block.foldZipWith f c a b {-# INLINE foldZipWith #-} -index :: PrimType ty => ScrubbedBlock ty -> Offset ty -> ty-index (ScrubbedBlock b) = blockIndex b- length :: PrimType ty => ScrubbedBlock ty -> CountOf ty length (ScrubbedBlock b) = Block.length b -map :: (PrimType a, PrimType b) => (a -> b) -> ScrubbedBlock a -> ScrubbedBlock b-map f (ScrubbedBlock b) =-    create (CountOf n) $ \(Offset i) -> f (blockIndex b (Offset i))-  where-    CountOf n = Block.length b-{-# INLINE map #-}- 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+thaw (ScrubbedBlock !b) = do     mb <- new (Block.length b)     unsafeCopyBytesRO mb 0 b 0 (mutableLengthBytes mb)     return mb -unsafeCast :: PrimType b => ScrubbedBlock a -> ScrubbedBlock b-unsafeCast (ScrubbedBlock b) = ScrubbedBlock (Block.unsafeCast b)- unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty) unsafeFreeze mb = Block.unsafeFreeze mb >>= scrubbed -#ifdef ML_KEM_TESTING-replicate :: PrimType ty => CountOf ty -> ty -> ScrubbedBlock ty-replicate n e = create n (const e) -fromList :: PrimType ty => [ty] -> ScrubbedBlock ty-fromList elems = runST $ do-    mb <- new (CountOf len)-    go mb 0 elems-    unsafeFreeze mb-  where-    !len = Prelude.length elems--    go !mb !i list = case list of-        []     -> return ()-        (x:xs) -> blockWrite mb i x >> go mb (i + 1) xs--toList :: PrimType ty => ScrubbedBlock ty -> [ty]-toList (ScrubbedBlock b) = Basement.Compat.IsList.toList b-#endif-- {- internal -}  assertPinned :: Block ty -> a -> a assertPinned mb = assert (isPinned mb == Pinned)  scrubbed :: PrimMonad prim => Block ty -> prim (ScrubbedBlock ty)-scrubbed b = assertPinned b $ unsafePrimFromIO $ do-    addBlockFinalizer b (scrub $ Block.unsafeCast b)-    return (ScrubbedBlock b)+scrubbed b = assertPinned b $ unsafePrimFromIO $+    scheduleBlockScrubbing b >> return (ScrubbedBlock b) +scheduleBlockScrubbing :: Block ty -> IO ()+scheduleBlockScrubbing b = addBlockFinalizer b (scrub $ Block.unsafeCast b)+{-# NOINLINE scheduleBlockScrubbing #-}+ scrub :: Block Word8 -> IO ()-scrub b = Block.unsafeThaw b >>= erase+scrub b = Block.unsafeThaw b >>= erase len+  where CountOf len = Block.length b  addBlockFinalizer :: Block ty -> IO () -> IO () addBlockFinalizer (Block barr) (IO finalizer) = IO $ \s ->    case mkWeak# barr () finalizer s of { (# s1, _ #) -> (# s1, () #) } -erase :: MutableBlock ty RealWorld -> IO ()-erase (MutableBlock mbarr) = IO $ \s1 ->-    case getSizeofMutableByteArray# mbarr s1 of-        (# s2, len #) -> case setByteArray# mbarr 0# len 0# s2 of-            s3 -> (# s3, () #)+erase :: Int -> MutableBlock Word8 RealWorld -> IO ()+erase (I# len) (MutableBlock 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)
src/SecureBlock.hs view
@@ -8,11 +8,10 @@ -- {-# LANGUAGE CPP #-} module SecureBlock-    ( SecureBlock, create, index, iterModify, Marking.map, new, thaw+    ( SecureBlock, create, index, iterModify, new, thaw     , unsafeCast, unsafeFreeze, toNormalForm #ifdef ML_KEM_TESTING-    , eq, Marking.showsPrec, fromList, Marking.replicate, toList-    , SecureBlock.length+    , eq, Marking.showsPrec, toList, SecureBlock.length #endif     ) where
src/Vector.hs view
@@ -23,7 +23,9 @@  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@@ -38,22 +40,47 @@ #endif import Data.Proxy +import Iterate import Math  newtype Vector (n :: Nat) a = Vector { unVector :: Array a }-    deriving (Eq, Show, Functor, NormalForm)+    deriving (Eq, Show, NormalForm) +instance Functor (Vector n) where+    fmap = mapVector+    {-# INLINE fmap #-}+ instance (Add a, KnownNat n) => Add (Vector n a) where     zero = create (const zero)+    {-# INLINE zero #-}     (.+) = Vector.zipWith (.+)+    {-# INLINE (.+) #-}     (.-) = Vector.zipWith (.-)-    neg (Vector a) = Vector (fmap neg a)+    {-# INLINE (.-) #-}+    neg = mapVector neg+    {-# INLINE neg #-}  create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a-create = Vector . Array.create (CountOf sz)-  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+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))+  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+{-# INLINE [1] 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 #-}++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 #-}++mapVector :: (a -> b) -> Vector n a -> Vector n b+mapVector f = Vector <$> fmap f . unVector+{-# INLINE [1] mapVector #-}+ arrayIndex :: Array a -> Offset a -> a #ifdef ML_KEM_TESTING arrayIndex = Array.index@@ -69,20 +96,24 @@ index = arrayIndex . unVector  concatMap :: Monoid b => (a -> b) -> Vector n a -> b-concatMap f = mconcat . Prelude.map f . toList . unVector+concatMap f = mconcat . mapToList f {-# INLINE concatMap #-} +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+ 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) ->         f (arrayIndex a (Offset i)) (arrayIndex b (Offset i))   where     CountOf sa = Array.length a-{-# INLINE zipWith #-}+{-# INLINE [1] zipWith #-}  fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c fold1ZipWith f g (Vector a) (Vector !b) =-    foldl' ff gg [1 .. sa - 1]+    foldl' ff gg (offsetsFrom 1 sa)   where     ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i))     gg = g (arrayIndex a 0) (arrayIndex b 0)@@ -90,7 +121,7 @@ {-# INLINE fold1ZipWith #-}  foldIndexWith :: (c -> Offset a -> a -> c) -> c -> Vector n a -> c-foldIndexWith f c (Vector a) = foldl' g c [0 .. sa - 1]+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@@ -98,3 +129,14 @@  toNormalForm :: NFData a => Vector n a -> () toNormalForm = Array.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+  #-}
tests/Tests.hs view
@@ -34,6 +34,8 @@  import GHC.TypeNats +import Foreign.Ptr (plusPtr)+ import Auxiliary import BlockN (BlockN) import Builder (Builder)@@ -121,15 +123,38 @@ byteEncodeBytes :: Int -> BlockN Sec 256 Word16 -> Bytes byteEncodeBytes d = runBytes . byteEncode d -byteEncodeBytes1 :: BlockN Sec 256 Word16 -> Bytes-byteEncodeBytes1 = runBytes . byteEncode1--byteEncodeBytes12 :: Tq Sec -> Bytes-byteEncodeBytes12 = runBytes . byteEncode12- runBytes :: Builder Sec -> Bytes runBytes = Builder.run . leak +runUnaligned :: Builder Sec -> UnalignedBytes+runUnaligned = Builder.runRelaxed . leak++newtype UnalignedBytes = UnalignedBytes (B.View Bytes)+    deriving (Eq,Ord,Show)++instance Semigroup UnalignedBytes where+    UnalignedBytes a <> UnalignedBytes b =+        B.allocAndFreeze (na + nb) $ \p -> do+            B.copyByteArrayToPtr a p+            B.copyByteArrayToPtr b (p `plusPtr` na)+      where+        na = B.length a+        nb = B.length b++instance Monoid UnalignedBytes where+    mempty = B.convert (B.empty :: Bytes)++instance B.ByteArrayAccess UnalignedBytes where+    length (UnalignedBytes v) = B.length v+    withByteArray (UnalignedBytes v) = B.withByteArray v++instance B.ByteArray UnalignedBytes where+    allocRet n f = do+        let build ba = UnalignedBytes (B.dropView ba offset)+        (a, ba) <- B.allocRet (n + offset) $ \p -> f (p `plusPtr` offset)+        return (a, build ba)+      where offset = 1+ #endif  data P = forall a. (ParamSet a, Show a) => P (Proxy a)@@ -287,12 +312,20 @@         , testGroup "encoding"             [ testProperty "byteEncode . byteDecode == id" $ \(D d) -> do                 b <- arbitraryBytes (32 * d)-                return (b === byteEncodeBytes d (byteDecode d b))+                return (b === runBytes (byteEncode d (byteDecode d b)))+            , testProperty "byteEncode . byteDecode == id (unaligned)" $ \(D d) -> do+                b <- arbitraryBytes (32 * d)+                return (B.convert b === runUnaligned (byteEncode d (byteDecode d b)))             , testProperty "byteEncode1 . byteDecode1 == id" $ do                 b <- arbitraryBytes 32-                return (b === byteEncodeBytes1 (byteDecode1 b))+                return (b === runBytes (byteEncode1 (byteDecode1 b)))+            , testProperty "byteEncode1 . byteDecode1 == id (unaligned)" $ do+                b <- arbitraryBytes 32+                return (B.convert b === runUnaligned (byteEncode1 (byteDecode1 b)))             , testProperty "byteDecode12 . byteEncode12 == id" $ \(PolyNTT p) ->-                p === byteDecode12 (byteEncodeBytes12 p)+                p === byteDecode12 (runBytes $ byteEncode12 p)+            , testProperty "byteDecode12 . byteEncode12 == id (unaligned)" $ \(PolyNTT p) ->+                p === byteDecode12 (runUnaligned $ byteEncode12 p)             , testProperty "byteEncode 8" $ \x ->                 B.replicate 256 x === byteEncodeBytes 8 (BlockN.replicate $ fromIntegral x)             , testCase "byteEncode 1 (zeros)" $