packages feed

mlkem (empty) → 0.1.0.0

raw patch · 27 files changed

+3146/−0 lines, 27 filesdep +aesondep +basedep +basement

Dependencies added: aeson, base, basement, bytestring, criterion, crypton, cryptonite, deepseq, directory, memory, mlkem, process, tasty, tasty-hunit, tasty-quickcheck, text, zlib

Files

+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Changelog for `mlkem`++## 0.1.0.0 - 2025-11-02++* First version. Released on an unsuspecting world.
+ LICENSE view
@@ -0,0 +1,26 @@+Copyright 2025 Olivier Chéron++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++1.  Redistributions of source code must retain the above copyright notice, this+    list of conditions and the following disclaimer.++2.  Redistributions in binary form must reproduce the above copyright notice,+    this list of conditions and the following disclaimer in the documentation+    and/or other materials provided with the distribution.++3.  Neither the name of the copyright holder nor the names of its contributors+    may be used to endorse or promote products derived from this software+    without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,38 @@+# ML-KEM++_Module-Lattice-based Key-Encapsulation Mechanism_ implemented in Haskell.++See [FIPS 203](https://csrc.nist.gov/pubs/fips/203/final).++Example session:++```haskell+> import Crypto.PubKey.ML_KEM+> import Data.Proxy+> let params = Proxy :: Proxy ML_KEM_768+> (encapKey, decapKey) <- generate params+> (sharedKey, ciphertext) <- encapsulate encapKey+> let sharedKey' = decapsulate decapKey ciphertext+> sharedKey == sharedKey'+True+```++## Notes++The library does its best to destroy secrets and intermediate buffers from+memory after use, despite the implementation in functional style.  This relies+on finalization by the garbage collector and is not guaranteed to run before+the program exits.  Also, depending on optimizations applied, lambdas may+capture variables and move them to the heap.  This could theoretically include+machine words containing secret information that would not then be destroyed.++Best performance is obtained with the LLVM code generator.++## Testing++The test suite executes all NIST test vectors but necessary files are not+included in the package to limit its size.  Instead, two files are downloaded+from the project repository during execution, and this relies on commands `sh`+and `curl` to run the script `tests/get-vectors.sh`.  If not applicable to your+environment, please execute the same steps manually.  It will be needed only+the first time.
+ benchs/Bench.hs view
@@ -0,0 +1,52 @@+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+module Main (main) where++import Criterion.Main++import Crypto.Random+import Crypto.PubKey.ML_KEM++import Data.ByteArray (Bytes)+import Data.Proxy++data KEM = forall a . ParamSet a => KEM (Proxy a)++kems :: [(String, KEM)]+kems =+    [ ("ML-KEM-512",  KEM (Proxy :: Proxy ML_KEM_512))+    , ("ML-KEM-768",  KEM (Proxy :: Proxy ML_KEM_768))+    , ("ML-KEM-1024", KEM (Proxy :: Proxy ML_KEM_1024))+    ]++doBench :: (String, KEM) -> Benchmark+doBench (name, KEM p) = bgroup name+    [ bench "generate" $ perRunEnv setupGenerate (return . runGenerate)+    , bench "encapsulate" $ perRunEnv setupEncap (return . runEncap)+    , bench "encapsulate (batch)" $ perBatchEnv (const setupEncap) (return . runEncap)+    , bench "decapsulate" $ perRunEnv setupDecap (return . runDecap)+    , bench "decapsulate (batch)" $ perBatchEnv (const setupDecap) (return . runDecap)+    ]+  where+    gen32 = getRandomBytes 32 :: IO Bytes++    setupGenerate = (,) <$> gen32 <*> gen32+    runGenerate = uncurry (generateWith p)++    setupEncap = do+        (ek, _) <- generate p+        m <- gen32+        return (ek, m)+    runEncap (ek, m) = encapsulateWith ek m++    runDecap (dk, c) = decapsulate dk c+    setupDecap = do+        (ek, dk) <- generate p+        (_, c) <- encapsulate ek+        return (dk, c)++main :: IO ()+main = defaultMain+    [ bgroup "mlkem" $ map doBench kems+    ]
+ mlkem.cabal view
@@ -0,0 +1,188 @@+cabal-version: 2.2++-- This file has been generated from package.yaml by hpack version 0.38.0.+--+-- see: https://github.com/sol/hpack++name:           mlkem+version:        0.1.0.0+synopsis:       Module-Lattice-based Key-Encapsulation Mechanism+description:    Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM) implemented in+                Haskell.+category:       Crypto+homepage:       https://codeberg.org/ocheron/hs-mlkem#readme+bug-reports:    https://codeberg.org/ocheron/hs-mlkem/issues+author:         Olivier Chéron+maintainer:     olivier.cheron@gmail.com+copyright:      2025 Olivier Chéron+license:        BSD-3-Clause+license-file:   LICENSE+build-type:     Simple+extra-source-files:+    README.md+    tests/get-vectors.sh+extra-doc-files:+    CHANGELOG.md++source-repository head+  type: git+  location: https://codeberg.org/ocheron/hs-mlkem++flag use_crypton+  description: Use crypton instead of cryptonite+  manual: True+  default: False++library+  exposed-modules:+      Crypto.PubKey.ML_KEM+  other-modules:+      Auxiliary+      Block+      BlockN+      Builder+      ByteArrayST+      Crypto+      Internal+      K_PKE+      Marking+      Math+      Matrix+      ScrubbedBlock+      SecureBlock+      SecureBytes+      Vector+      Paths_mlkem+  autogen-modules:+      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+  build-depends:+      base >=4.7 && <5+    , basement >=0.0.8+    , deepseq+    , memory+  default-language: Haskell2010+  if flag(use_crypton)+    build-depends:+        crypton+  else+    build-depends:+        cryptonite >=0.26++test-suite mlkem-test+  type: exitcode-stdio-1.0+  main-is: Tests.hs+  other-modules:+      EncapDecap+      KeyGen+      Util+      Vectors+      Paths_mlkem+  autogen-modules:+      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+  build-depends:+      aeson+    , base >=4.7 && <5+    , basement >=0.0.8+    , bytestring+    , deepseq+    , directory+    , memory+    , mlkem+    , process+    , tasty+    , tasty-hunit+    , tasty-quickcheck+    , text+    , zlib+  default-language: Haskell2010+  if flag(use_crypton)+    build-depends:+        crypton+  else+    build-depends:+        cryptonite >=0.26++test-suite mlkem-test-full+  type: exitcode-stdio-1.0+  main-is: Tests.hs+  other-modules:+      Auxiliary+      Block+      BlockN+      Builder+      ByteArrayST+      Crypto+      Crypto.PubKey.ML_KEM+      Internal+      K_PKE+      Marking+      Math+      Matrix+      ScrubbedBlock+      SecureBlock+      SecureBytes+      Vector+      EncapDecap+      KeyGen+      Util+      Vectors+      Paths_mlkem+  autogen-modules:+      Paths_mlkem+  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+  cpp-options: -DML_KEM_TESTING+  build-depends:+      aeson+    , base >=4.7 && <5+    , basement >=0.0.8+    , bytestring+    , deepseq+    , directory+    , memory+    , process+    , tasty+    , tasty-hunit+    , tasty-quickcheck+    , text+    , zlib+  default-language: Haskell2010+  if flag(use_crypton)+    build-depends:+        crypton+  else+    build-depends:+        cryptonite >=0.26++benchmark mlkem-bench+  type: exitcode-stdio-1.0+  main-is: Bench.hs+  other-modules:+      Paths_mlkem+  autogen-modules:+      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+  build-depends:+      base >=4.7 && <5+    , basement >=0.0.8+    , criterion+    , deepseq+    , memory+    , mlkem+  default-language: Haskell2010+  if flag(use_crypton)+    build-depends:+        crypton+  else+    build-depends:+        cryptonite >=0.26
+ src/Auxiliary.hs view
@@ -0,0 +1,585 @@+-- |+-- Module      : Auxiliary+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- ML-KEM auxiliary functions+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}+module Auxiliary+    ( Zq, Rq, Tq, (..+), (..-)+    , ntt, nttInv, rcompress, rdecompress+    , byteEncode, byteDecode, byteEncode12, byteDecode12+    , byteEncode1, byteDecode1, sampleNTT, samplePolyCBD+#ifdef ML_KEM_TESTING+    , compress, decompress+    , bitRev7, fromZq, toZq, fromCoeffs, toCoeffs+#endif+    ) where++import Basement.NormalForm+import Basement.PrimType+import Basement.Types.OffsetSize++import Crypto.Hash.Algorithms++import Data.ByteArray (ByteArrayAccess, Bytes, View)+import qualified Data.ByteArray as B+import qualified Data.Memory.Endian as B++import Control.DeepSeq (NFData(..))+import Control.Monad+import Control.Monad.ST++import Data.Bits+import Data.Proxy+import Data.Word++import GHC.TypeNats++import Foreign.Ptr (Ptr, plusPtr)+import Foreign.Storable (pokeByteOff)++import Block (blockIndex)+import BlockN (BlockN, MutableBlockN)+import Builder (Builder)+import Crypto (BlockDigest)+import Marking (Classified, SecurityMarking(..), Leak(..))+import SecureBlock (SecureBlock)+import SecureBytes (SecureBytes)+import qualified BlockN+import qualified Builder+import qualified ByteArrayST as ST+import qualified Crypto+import Math++type N = 256++n :: Int+n = 256++q :: Integer+q = 3329++q16 :: Word16+q16 = fromInteger q++q32 :: Word32+q32 = fromInteger q++q64 :: Word64+q64 = fromInteger q++bitRev7 :: Word8 -> Word8+bitRev7 b =+    (b `unsafeShiftR` 6 .&. 1) .|.+    (b `unsafeShiftR` 5 .&. 1) `unsafeShiftL` 1 .|.+    (b `unsafeShiftR` 4 .&. 1) `unsafeShiftL` 2 .|.+    (b `unsafeShiftR` 3 .&. 1) `unsafeShiftL` 3 .|.+    (b `unsafeShiftR` 2 .&. 1) `unsafeShiftL` 4 .|.+    (b `unsafeShiftR` 1 .&. 1) `unsafeShiftL` 5 .|.+    (b .&. 1) `unsafeShiftL` 6++-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞+reduceSimple :: Word16 -> Word16+reduceSimple x = (mask .&. x) .|. (complement mask .&. subtracted)+  where+    subtracted = x - q16+    mask = negate (subtracted `unsafeShiftR` 15)+{-# INLINE reduceSimple #-}++-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞² + 𝑞+reduce :: Word32 -> Word16+reduce x = reduceSimple (fromIntegral remainder)+  where+    p = fromIntegral x * ((1 `unsafeShiftL` 24) `div` q64)+    quotient = fromIntegral (p `unsafeShiftR` 24)+    remainder = x - quotient * q32+{-# INLINE reduce #-}++newtype Zq = Zq Word16+#ifdef ML_KEM_TESTING+    deriving (Eq, Show)+#else+    deriving Eq+#endif++instance PrimType Zq where+    type PrimSize Zq = 2+    primSizeInBytes _ = primSizeInBytes (Proxy :: Proxy Word16)+    {-# INLINE primSizeInBytes #-}+    primShiftToBytes _ = primShiftToBytes (Proxy :: Proxy Word16)+    {-# INLINE primShiftToBytes #-}+    primBaUIndex ba (Offset i) = Zq (primBaUIndex ba (Offset i))+    {-# INLINE primBaUIndex #-}+    primMbaURead mba (Offset i) = Zq <$> primMbaURead mba (Offset i)+    {-# INLINE primMbaURead #-}+    primMbaUWrite mba (Offset i) (Zq a) = primMbaUWrite mba (Offset i) a+    {-# INLINE primMbaUWrite #-}+    primAddrIndex addr (Offset i) = Zq (primAddrIndex addr (Offset i))+    {-# INLINE primAddrIndex #-}+    primAddrRead addr (Offset i) = Zq <$> primAddrRead addr (Offset i)+    {-# INLINE primAddrRead #-}+    primAddrWrite addr (Offset i) (Zq a) = primAddrWrite addr (Offset i) a+    {-# INLINE primAddrWrite #-}++instance Add Zq where+    zero = Zq 0+    Zq a .+ Zq b = Zq $ reduceSimple (a + b)+    Zq a .- Zq b = Zq $ reduceSimple (a + q16 - b)+    neg (Zq a) = Zq $ reduceSimple (q16 - a)++instance Mul Zq where+    one = Zq 1+    Zq a .* Zq b = Zq $ reduce (fromIntegral a * fromIntegral b)++#ifdef ML_KEM_TESTING+instance MulAdd Zq where+    mulAdd (Zq a) (Zq b) (Zq c) = Zq $ reduce $+        fromIntegral a * fromIntegral b + fromIntegral c++instance BiMul Zq Zq where+    (..*) = (.*)++instance BiMulAdd Zq Zq where+    biMulAdd = mulAdd++fromZq :: Zq -> Word16+fromZq (Zq a) = a+#endif++toZq :: Word16 -> Zq+toZq = Zq . reduce . fromIntegral++newtype Rq marking = Rq (BlockN marking N Zq)+#ifdef ML_KEM_TESTING+    deriving (Eq, Show)+#endif++instance Classified marking => Add (Rq marking) where+    zero = Rq zero+    Rq a .+ Rq b = Rq (a .+ b)+    Rq a .- Rq b = Rq (a .- b)+    neg (Rq a) = Rq (neg a)+    {-# SPECIALIZE instance Add (Rq Sec) #-}+    {-# SPECIALIZE instance Add (Rq Pub) #-}++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++(..-) :: Rq Pub -> Rq Sec -> Rq Sec+Rq a ..- Rq b = Rq $ BlockN.zipWith (.-) a b++instance Leak Rq++#ifdef ML_KEM_TESTING+fromCoeffs :: [Zq] -> Maybe (Rq Sec)+fromCoeffs = fmap Rq . BlockN.fromList++toCoeffs :: Rq Sec -> [Zq]+toCoeffs (Rq a) = BlockN.toList a+#endif++newtype Tq marking = Tq (BlockN marking N Zq)+#ifdef ML_KEM_TESTING+    deriving (Eq, Show)+#endif++instance Classified marking => Add (Tq marking) where+    zero = Tq zero+    Tq a .+ Tq b = Tq (a .+ b)+    Tq a .- Tq b = Tq (a .- b)+    neg (Tq a) = Tq (neg a)+    {-# SPECIALIZE instance Add (Tq Sec) #-}+    {-# SPECIALIZE instance Add (Tq Pub) #-}++instance Leak Tq++instance Classified marking => NFData (Tq marking) where+    rnf (Tq a) = toNormalForm a++instance BiMul (Tq Pub) (Tq Sec) where+    (..*) = multiplyNTTs++instance BiMulAdd (Tq Pub) (Tq Sec) where+    biMulAdd = multiplyNTTsAdd++#ifdef ML_KEM_TESTING+instance Mul (Tq Sec) where+    one = Tq $ BlockN.create $ \(Offset i) -> if even i then one else zero+    (.*) = (..*) . leak++instance MulAdd (Tq Sec) where+    mulAdd = biMulAdd . leak+#endif++instance Crypto.ConstEqW (Tq Sec) where+    constEqW (Tq a) (Tq b) = Crypto.constEqW+        (BlockN.unsafeCast a :: SecureBlock Sec Word)+        (BlockN.unsafeCast b :: SecureBlock Sec Word)++instance Crypto.ConstEqW (Tq Pub) where+    constEqW (Tq a) (Tq b) = Crypto.constEqW+        (BlockN.unsafeCast a :: SecureBlock Pub Word)+        (BlockN.unsafeCast b :: SecureBlock Pub Word)++-- 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+  where+    outer !b !i len = when (len >= 2) $ inner b i len 0++    inner !b !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 !b !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 #-}++-- 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+    BlockN.iterModify (\x -> x .* Zq 3303) b+    Rq <$> BlockN.unsafeFreeze b+  where+    outer !b !i len = when (len <= 128) $ inner b i len 0++    inner !b !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 !b !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)++-- 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+  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+            a1 = BlockN.index f (ii + 1)+            b0 = BlockN.index g ii+            b1 = BlockN.index g (ii + 1)+            (c0, c1) = baseCaseMultiply a0 a1 b0 b1 (BlockN.index gamma i)+        BlockN.write b ii c0+        BlockN.write b (ii + 1) c1+        loop b (i + 1)++-- Computes the product of two degree-one polynomials with respect to a quadratic modulus+baseCaseMultiply :: Zq -> Zq -> Zq -> Zq -> Zq -> (Zq, Zq)+baseCaseMultiply (Zq a0) (Zq a1) (Zq b0) (Zq b1) (Zq g) = (Zq c0, Zq c1)+  where+    x `mul` y = fromIntegral x * fromIntegral y+    b1g = reduce (b1 `mul` g)+    !c0 = reduce (a0 `mul` b0 + a1 `mul` b1g)+    !c1 = reduce (a0 `mul` b1 + a1 `mul` b0)++-- 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+  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+            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+        loop b (i + 1)++-- baseCaseMultiply then add a third term+baseCaseMultiplyAdd :: Zq -> Zq -> Zq -> Zq -> Zq -> Zq -> Zq -> (Zq, Zq)+baseCaseMultiplyAdd (Zq a0) (Zq a1) (Zq b0) (Zq b1) (Zq c0) (Zq c1) (Zq g) = (Zq d0, Zq d1)+  where+    x `mul` y = fromIntegral x * fromIntegral y+    b1g = reduce (b1 `mul` g)+    !d0 = reduce (fromIntegral c0 + a0 `mul` b0 + a1 `mul` b1g)+    !d1 = reduce (fromIntegral c1 + a0 `mul` b1 + a1 `mul` b0)++-- Values of 17 ^ BitRev7(𝑖) mod 𝑞 for 𝑖 ∈ {0, … , 127}+zetaPowBitRev :: BlockN Pub 128 Zq+zetaPowBitRev = runST $ do+    out <- BlockN.new (Proxy :: Proxy Pub)+    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++-- Compress a field element with 𝑑 < 12+compress :: Int -> Zq -> Word16+compress d (Zq x) = fromIntegral $+    ((fromIntegral x `unsafeShiftL` d + qHalf) * factor) `unsafeShiftR` 34+  where+    qHalf = (q64 + 1) `unsafeShiftR` 1+    factor = (1 `unsafeShiftL` 34) `div` q64+{-# INLINE compress #-}++-- Decompress a field element with 𝑑 < 12+decompress :: Int -> Word16 -> Zq+decompress d y = Zq $ fromIntegral (x2d `unsafeShiftR` d)+  where x2d = fromIntegral y * q32 + (1 `unsafeShiftL` (d - 1))+{-# INLINE decompress #-}++-- 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 #-}++-- 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 #-}++-- 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+  where+    runXof !b !xofLen !pos !j = case someNatVal (8 * fromIntegral xofLen) of+        SomeNat proxy -> do+            let bytes = Crypto.unBlockDigest (doHash proxy)+            loop b xofLen bytes pos j++    loop !b !xofLen !bytes !pos j+        | j == 256 = return ()+        | pos >= Offset xofLen = runXof b (xofLen + 56 * 3) pos j+        | otherwise = do+            let c0 = fromIntegral $ blockIndex bytes pos+                c1 = fromIntegral $ blockIndex bytes (pos + 1)+                c2 = fromIntegral $ blockIndex bytes (pos + 2)+                d1 = c0 + (c1 .&. 0xF) `unsafeShiftL` 8+                d2 = (c1 `unsafeShiftR` 4) + (c2 `unsafeShiftL` 4)+            j2 <- poke b j d1+            when (j2 < 256) $ poke b j2 d2 >>= loop b xofLen bytes (pos + 3)++    poke b j d+        | d < q16 = BlockN.write b j (Zq d) >> return (j + 1)+        | otherwise = return j++    doHash :: KnownNat bitlen => proxy bitlen -> BlockDigest (SHAKE128 bitlen)+    doHash _ = Crypto.hashToBlock input++    input :: SecureBytes Pub+    !input = B.unsafeCreate (len + 2) $ \d -> do+        B.copyByteArrayToPtr seed d+        pokeByteOff d len x+        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++peekWordPos :: Ptr WordLE -> BitPos -> ST s Word+peekWordPos a bp = B.unLE <$> ST.peekElemOff a (wordOff bp)++pokeWordPos :: Ptr WordLE -> BitPos -> Word -> ST s ()+pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . B.LE++newtype BitPos = BitPos Int++zeroPos :: BitPos+zeroPos = BitPos 0++wordOff :: BitPos -> Int+wordOff (BitPos p) = div p wordBits++bitPos :: BitPos -> Int+bitPos (BitPos p) = p .&. (wordBits - 1)++availPos :: Int -> BitPos -> Int+availPos requested (BitPos p) = min available requested+  where available = wordBits - (p .&. (wordBits - 1))++nextPos :: Int -> BitPos -> (Int, BitPos)+nextPos requested (BitPos p) = (howMany, BitPos $ p + howMany)+  where howMany = availPos requested (BitPos p)++getMask :: Int -> Word+getMask howMany+    | howMany >= wordBits = maxBound+    | otherwise = (1 `unsafeShiftL` howMany) - 1++-- 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+  where+    loop :: Ptr WordLE -> MutableBlockN Sec N Zq s -> Offset Zq -> BitPos -> ST s ()+    loop !p !f !i !bp = when (i < Offset n) $ do+        (xs, bp') <- getBits p bp 0 (fromIntegral eta)+        (ys, bp'') <- getBits p bp' 0 (fromIntegral eta)+        BlockN.write f i (Zq xs .- Zq ys)+        loop p f (i + 1) bp''++    getBits :: Ptr WordLE -> BitPos -> Word16 -> Int -> ST s (Word16, BitPos)+    getBits !p !bp !acc !j+        | j == 0    = return (acc, bp)+        | otherwise = do+            x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp+            let (howMany, bp') = nextPos j bp+                bits = x .&. getMask howMany+            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+  where+    outer :: Ptr WordLE -> BitPos -> Word -> Int -> ST s ()+    outer !b !bp !o pos = when (pos < n) $+        inner b pos bp o (BlockN.index f (Offset pos)) d++    inner :: Ptr WordLE -> Int -> BitPos -> Word -> 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'+      where+        (howMany, bp') = nextPos j bp+        x = fromIntegral a .&. getMask howMany+        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+  where+    loop :: Ptr WordLE -> Word -> Int -> ST s ()+    loop !b !o pos+        | pos == n = return ()+        | bitPos bp + 1 < wordBits = loop b o' (pos + 1)+        | otherwise = pokeWordPos b bp o' >> loop b 0 (pos + 1)+      where+        bp = BitPos pos+        x = fromIntegral (a .&. 1)+        o' = o .|. (x `unsafeShiftL` bitPos bp)+        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++-- 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+  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++    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)+        | 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'++    get :: Ptr WordLE -> BitPos -> Int -> ST s Word+    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 :: Int -> View Bytes -> BlockN Sec N Word16 #-}+{-# SPECIALIZE byteDecode :: Int -> View Bytes -> BlockN Pub N Word16 #-}++-- 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+  where+    outer :: MutableBlockN Sec N Word16 s -> Ptr WordLE -> Int -> ST s ()+    outer !f !p i = when (i < n) $ do+        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 !f !p !acc !i j+        | j == wordBits = outer f 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)++-- byteDecode with 𝑑=12 and conversion to the field+byteDecode12 :: Classified marking => ByteArrayAccess ba => ba -> Tq marking+byteDecode12 = Tq . BlockN.map toZq . byteDecode 12
+ src/Block.hs view
@@ -0,0 +1,93 @@+-- |+-- Module      : Block+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- An array of primitive (unlifted) elements.  This module currently exposes+-- the implementation from basement and fixes a lack of inlining in the+-- @create@ function.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+module Block+    ( Block, MutableBlock, blockIndex, blockRead, blockWrite+    , create, foldZipWith, iterModify, Block.length, Block.map+    , 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)+import Basement.Block.Mutable (MutableBlock)+import qualified Basement.Block as Block hiding (create, map)+import qualified Basement.Block.Mutable as Block+import Basement.Monad+import Basement.PrimType+import Basement.Types.OffsetSize++import Control.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 ()+#ifdef ML_KEM_TESTING+blockIndex = Block.index+blockRead = Block.read+blockWrite = Block.write+#else+blockIndex = Block.unsafeIndex+blockRead = Block.unsafeRead+blockWrite = Block.unsafeWrite+#endif++create :: PrimType ty+       => CountOf ty+       -> (Offset ty -> ty)+       -> Block ty+create n initializer = runST $ do+    mb <- Block.new n+    loop mb 0+    Block.unsafeFreeze mb+  where+    loop !mb i+        | 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)+           -> MutableBlock ty (PrimState prim)+           -> prim ()+iterModify f ma = loop 0+  where+    !sz = Block.mutableLength ma+    loop i+        | i .==# sz = pure ()+        | otherwise = Block.unsafeRead ma i >>= \x -> Block.unsafeWrite ma i (f x) >> loop (i+1)+{-# 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) $+    loop c 0+  where+    CountOf sa = Block.length a+    CountOf sb = Block.length b++    loop !acc i+        | i == sa = acc+        | otherwise = do+            let va = Block.unsafeIndex a (Offset i)+            let vb = Block.unsafeIndex b (Offset i)+            loop (f acc va vb) (i + 1)+{-# INLINE foldZipWith #-}
+ src/BlockN.hs view
@@ -0,0 +1,123 @@+-- |+-- Module      : BlockN+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- A secure block with length at type level+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+module BlockN+    ( BlockN, MutableBlockN, create, index, iterModify, BlockN.map+    , BlockN.new, BlockN.read, BlockN.thaw, BlockN.unsafeCast+    , BlockN.unsafeFreeze, BlockN.write, BlockN.zipWith+#ifdef ML_KEM_TESTING+    , BlockN.fromList, BlockN.replicate, BlockN.toList+#endif+    ) where++import Basement.Monad+import Basement.Nat+import Basement.NormalForm+import Basement.PrimType+import Basement.Types.OffsetSize++import Data.Proxy++import Block (MutableBlock, blockRead, blockWrite)+import Marking (Classified, SecurityMarking)+import SecureBlock (SecureBlock)+import qualified SecureBlock+import Math++newtype BlockN marking (n :: Nat) a = BlockN { unBlockN :: SecureBlock marking a }++#ifdef ML_KEM_TESTING+instance (Classified marking, PrimType a) => Eq (BlockN marking n a) where+    BlockN a == BlockN b = SecureBlock.eq a b++instance (Classified marking, PrimType a, Show a) => Show (BlockN marking n a) where+    showsPrec d = SecureBlock.showsPrec d . unBlockN+#endif++instance Classified marking => 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+    zero = create (const zero)+    {-# INLINE zero #-}+    (.+) = BlockN.zipWith (.+)+    {-# INLINE (.+) #-}+    (.-) = BlockN.zipWith (.-)+    {-# INLINE (.-) #-}+    neg = BlockN.map neg+    {-# INLINE neg #-}++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 = 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)++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+  where+    a = SecureBlock.fromList elems+    !sz = fromIntegral $ natVal (Proxy :: Proxy n)++toList :: (Classified marking, PrimType a) => BlockN marking n a -> [a]+toList = SecureBlock.toList . unBlockN+#endif++create :: forall marking n ty. (Classified marking, KnownNat n, PrimType ty)+       => (Offset ty -> ty)+       -> BlockN marking n ty+create initializer = BlockN $ SecureBlock.create (CountOf sz) initializer+  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+{-# INLINE map #-}++iterModify :: (PrimType ty, PrimMonad prim)+           => (ty -> ty)+           -> MutableBlockN marking n ty (PrimState prim)+           -> prim ()+iterModify f = SecureBlock.iterModify f . unMutableBlockN+{-# INLINE iterModify #-}++zipWith :: (Classified ma, Classified mb, 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 = SecureBlock.unsafeCast . unBlockN++read :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> prim a+read = blockRead . unMutableBlockN++write :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> a -> prim ()+write = blockWrite . unMutableBlockN++new :: forall proxy marking n a prim. (Classified marking, KnownNat n, PrimMonad prim, PrimType a) => proxy marking -> prim (MutableBlockN marking n a (PrimState prim))+new prx = MutableBlockN <$> SecureBlock.new prx (CountOf sz)+  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)++thaw :: (Classified marking, PrimMonad prim, PrimType a) => 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
+ src/Builder.hs view
@@ -0,0 +1,83 @@+-- |+-- Module      : Builder+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Eliminate intermediate allocations when concatenating several buffers+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE RankNTypes #-}+module Builder+    ( Builder, builderLength, bytes, copyBuilderToPtr, create, run, runRelaxed+    , runToBlock, unsafeCreate+    ) where++import Basement.Types.OffsetSize++import Data.ByteArray (ByteArray)++import Control.Monad.ST+import Control.Monad.ST.Unsafe++import Data.Semigroup+import Data.Word++import Foreign.Ptr (Ptr, castPtr, plusPtr)++import Block (Block)+import Marking (Classified, Leak(..), SecurityMarking(..))+import SecureBytes (SecureBytes)+import qualified Block+import qualified ByteArrayST as ST+import qualified SecureBytes++data Builder (marking :: SecurityMarking) = Builder+    { builderLength :: {-# UNPACK #-} !Int+    , copyBuilderToPtr :: forall a s. Ptr a -> ST s ()+    }++instance Semigroup (Builder marking) where+    b1 <> b2  = create (n1 + n2) $ \p ->+        copyBuilderToPtr b1 p >> copyBuilderToPtr b2 (p `plusPtr` n1)+      where+        n1 = builderLength b1+        n2 = builderLength b2++instance Monoid (Builder marking) where+    mempty = empty+    mconcat builders = create n (`loop` 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++instance Leak Builder++bytes :: Classified marking => SecureBytes marking -> Builder marking+bytes b = unsafeCreate (SecureBytes.length b) (SecureBytes.copyByteArrayToPtr b)++create :: Int -> (forall s. Ptr a -> ST s ()) -> Builder marking+create n f = Builder n (f . castPtr)+{-# INLINE create #-}++empty :: Builder marking+empty = Builder 0 $ \_ -> return ()++run :: Classified marking => Builder marking -> SecureBytes marking+run b = SecureBytes.unsafeCreate (builderLength b) (copyBuilderToPtr b)++runRelaxed :: ByteArray ba => Builder Pub -> ba+runRelaxed b = ST.unsafeCreate (builderLength b) (copyBuilderToPtr b)++runToBlock :: Builder Pub -> Block Word8+runToBlock b = runST $ do+    mb <- Block.newPinned (CountOf $ builderLength b)+    Block.withMutablePtr mb (copyBuilderToPtr b)+    Block.unsafeFreeze mb++unsafeCreate :: Int -> (Ptr a -> IO ()) -> Builder marking+unsafeCreate n f = create n (unsafeIOToST . f)+{-# INLINE unsafeCreate #-}
+ src/ByteArrayST.hs view
@@ -0,0 +1,42 @@+-- |+-- Module      : ByteArrayST+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Byte array primitives in the @ST@ monad instead of @IO@+--+{-# LANGUAGE RankNTypes #-}+module ByteArrayST+    ( unsafeCreate, withByteArray+    , peek, peekElemOff, pokeElemOff, pokeByteOff+    ) where++import Data.ByteArray (ByteArray, ByteArrayAccess)+import qualified Data.ByteArray as B++import Control.Monad.ST+import Control.Monad.ST.Unsafe++import Foreign.Ptr (Ptr)+import Foreign.Storable (Storable)+import qualified Foreign.Storable as S++unsafeCreate :: ByteArray ba => Int -> (forall s. Ptr p -> ST s ()) -> ba+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)+{-# INLINE withByteArray #-}++peek :: Storable a => Ptr a -> ST s a+peek = unsafeIOToST . S.peek++peekElemOff :: Storable a => Ptr a -> Int -> ST s a+peekElemOff a = unsafeIOToST . S.peekElemOff a++pokeElemOff :: Storable a => Ptr a -> Int -> a -> ST s ()+pokeElemOff a off = unsafeIOToST . S.pokeElemOff a off++pokeByteOff :: Storable a => Ptr a -> Int -> a -> ST s ()+pokeByteOff a off = unsafeIOToST . S.pokeByteOff a off
+ src/Crypto.hs view
@@ -0,0 +1,214 @@+-- |+-- Module      : Crypto+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Crypto-related utilities like the ML-KEM hash and PRF functions, or more+-- general concerns like constant-time equality and selection.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE ScopedTypeVariables #-}+module Crypto+    ( ConstEqW(..), BoolW, andW, toBool, constSelectBytes, snoc, append+    , prf, h, j, g, BlockDigest, unBlockDigest, hashToBlock+    ) where++import Crypto.Hash (Context)+import Crypto.Hash.Algorithms+import Crypto.Hash.IO++import Control.Exception (assert)+import Control.Monad+import Control.Monad.ST++import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes)+import qualified Data.ByteArray as B++import Data.Bits+import Data.Word++import GHC.TypeNats++import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (pokeByteOff)++import Block (Block)+import Builder (Builder)+import ScrubbedBlock (ScrubbedBlock)+import Vector (Vector)+import qualified Block+import qualified Builder+import qualified ByteArrayST as ST+import qualified ScrubbedBlock+import qualified Vector++newtype BoolW = BoolW Word++#ifdef ML_KEM_TESTING+instance Show BoolW where+    showsPrec d = showsPrec d . toBool+#endif++toBool :: BoolW -> Bool+toBool (BoolW mask) = mask /= 0++falseW, trueW :: BoolW+falseW = BoolW 0+trueW = BoolW maxBound++andW :: BoolW -> BoolW -> BoolW+andW (BoolW a) (BoolW b) = BoolW (a .&. b)++bitsW :: Int+bitsW = let BoolW x = falseW in finiteBitSize x++bytesW :: Int+bytesW = div bitsW 8++eqW :: Word -> Word -> BoolW+eqW a b = isZeroW (a `xor` b)+  where+    isZeroW x = BoolW $ msbW (complement x .&. (x - 1))+    msbW x = negate (x `unsafeShiftR` (bitsW - 1))++assertMultW :: Int -> a -> a+assertMultW n = assert (n .&. mask == 0)+  where mask = bytesW - 1++class ConstEqW a where+    constEqW :: a -> a -> BoolW++instance ConstEqW a => ConstEqW (Vector n a) where+    constEqW =+        Vector.fold1ZipWith (\mask x y -> mask `andW` constEqW x y) constEqW++instance ConstEqW (Block Word) where+    constEqW a b+        | Block.length a /= Block.length b = falseW+        | otherwise = Block.foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b++instance ConstEqW (ScrubbedBlock Word) where+    constEqW a b+        | ScrubbedBlock.length a /= ScrubbedBlock.length b = falseW+        | otherwise = ScrubbedBlock.foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b++instance ConstEqW Bytes where+    constEqW = bytesConstEqW++instance ConstEqW ScrubbedBytes where+    constEqW = bytesConstEqW++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++foldZipWith :: (ByteArrayAccess bs1, ByteArrayAccess bs2)+            => (c -> Word -> Word -> c) -> c -> bs1 -> bs2 -> c+foldZipWith f c a b = assert (sa == sb) $ assertMultW sa $ assertMultW sb $+    runST $ ST.withByteArray a $ \pa -> ST.withByteArray b $ \pb ->+        loop (pa :: Ptr Word) (pb :: Ptr Word) c 0+  where+    !sa = B.length a+    !sb = B.length b++    loop !pa !pb !acc i+        | i == sa = return acc+        | otherwise = do+            va <- ST.peek pa+            vb <- ST.peek pb+            loop (pa `plusPtr` bytesW) (pb `plusPtr` bytesW) (f acc va vb) (i + bytesW)+{-# INLINE foldZipWith #-}++zipWith :: (Word -> Word -> Word) -> ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes+zipWith f a b = assert (sa == sb) $ assertMultW sa $ assertMultW sb $+    ST.unsafeCreate sa $ \out ->+        ST.withByteArray a $ \pa -> ST.withByteArray b $ \pb ->+            loop out pa pb 0+  where+    !sa = B.length a+    !sb = B.length b++    loop :: Ptr Word -> Ptr Word -> Ptr Word -> Int -> ST s ()+    loop !out !pa !pb i = when (i < sa) $ do+        va <- ST.peek pa+        vb <- ST.peek pb+        ST.pokeByteOff out i $ f va vb+        loop out (pa `plusPtr` bytesW) (pb `plusPtr` bytesW) (i + bytesW)+{-# INLINE zipWith #-}++constSelectBytes :: BoolW -> ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes+constSelectBytes (BoolW !mask) = Crypto.zipWith f+  where f yes no = (mask .&. yes) .|. (complement mask .&. no)++-- This version of snoc accepts a more general input and uses internally a call+-- to copyByteArrayToPtr, so it does not need a trampoline when the input is+-- backed by Block Word8+snoc :: ByteArrayAccess a => a -> Word8 -> ScrubbedBytes+snoc a b =+    B.allocAndFreeze (na + 1) $ \p -> do+        B.copyByteArrayToPtr a p+        pokeByteOff p na b+  where na = B.length a+{-# INLINE snoc #-}++-- This version of append is more polymorphic and requires no trampoline when+-- fed with an input backed by Block Word8.+append :: (ByteArrayAccess a, ByteArrayAccess b) => a -> b -> ScrubbedBytes+append a 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+{-# INLINE append #-}++prf :: ByteArrayAccess s => Word -> s -> Word8 -> ScrubbedBytes+prf !eta s !b = case someNatVal (8 * 64 * fromIntegral eta) of+    SomeNat proxy -> unDigest (doHash proxy)+  where+    doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen)+    doHash _ = hash (snoc s b)++h :: ByteArrayAccess s => s -> Bytes+h = Builder.run . hashWith SHA3_256++j :: ScrubbedBytes -> ScrubbedBytes+j = Builder.run . hashWith (SHAKE256 :: SHAKE256 256)++g  :: ByteArray ba => ScrubbedBytes -> (ba, B.View ScrubbedBytes)+g c = (B.convert $ B.takeView ab 32, B.dropView ab 32)+  where ab = Builder.run $ hashWith SHA3_512 c++-- Override cryptonite types and hashing functions.+--+-- Standard type Digest is a newtype over an unpinned Block Word8, which+-- requires a trampoline to implement most Ptr access to the underlying byte+-- array.  Instead we re-implement here the Digest type over ScrubbedBytes as+-- well as pinned Block backends, to avoid trampoline costs.  Additionnally+-- we use the mutable API to avoid copying the hashing Context in between+-- steps init/update/finalize and then clear the content.++newtype Digest a = Digest { unDigest :: ScrubbedBytes }+newtype BlockDigest a = BlockDigest { unBlockDigest :: Block Word8 }++hash :: forall a ba. (HashAlgorithm a, ByteArrayAccess ba) => ba -> Digest a+hash = Digest . Builder.run . hashWith (undefined :: a)++hashToBlock :: forall a. HashAlgorithm a => Bytes -> BlockDigest a+hashToBlock = BlockDigest . Builder.runToBlock . hashWith (undefined :: a)++hashWith :: forall marking a ba. (HashAlgorithm a, ByteArrayAccess ba) => a -> ba -> Builder marking+hashWith a ba = Builder.unsafeCreate (hashDigestSize a) $ \dig -> do+    ctx <- hashMutableInit+    hashMutableUpdate (ctx :: MutableContext a) ba+    B.withByteArray ctx $ \pctx -> do+        hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig+        ScrubbedBlock.erasePtr (B.length ctx) pctx
+ src/Crypto/PubKey/ML_KEM.hs view
@@ -0,0 +1,94 @@+-- |+-- Module      : Crypto.PubKey.ML_KEM+-- License     : BSD-3-Clause+-- Maintainer  : Olivier Chéron <olivier.cheron@gmail.com>+-- Stability   : provisional+-- Portability : unknown+--+-- Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM), defined+-- in <https://csrc.nist.gov/pubs/fips/203/final FIPS 203>.+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeFamilies #-}+module Crypto.PubKey.ML_KEM+    ( EncapsulationKey, DecapsulationKey, Ciphertext, SharedSecret+    -- * Operations+    , generate, generateWith, encapsulate, encapsulateWith, decapsulate+    -- * Parameter sets+    , ParamSet, ML_KEM_512, ML_KEM_768, ML_KEM_1024+    -- * Conversions and checks+    , Decode(..), Encode(..)+    , toPublic, checkKeyPair+    ) where++import Crypto.Random++import Data.ByteArray (ByteArrayAccess, ScrubbedBytes)+import qualified Data.ByteArray as B++import Internal++-- | ML-KEM-512 (security category 1)+data ML_KEM_512  = ML_KEM_512  deriving Show+-- | ML-KEM-768 (security category 3)+data ML_KEM_768  = ML_KEM_768  deriving Show+-- | ML-KEM-1024 (security category 5)+data ML_KEM_1024 = ML_KEM_1024 deriving Show++instance ParamSet ML_KEM_512 where+    type K ML_KEM_512 = 2+    getParams _ = Params 3 2 10 4+instance ParamSet ML_KEM_768 where+    type K ML_KEM_768 = 3+    getParams _ = Params 2 2 10 4+instance ParamSet ML_KEM_1024 where+    type K ML_KEM_1024 = 4+    getParams _ = Params 2 2 11 5++-- | Generate an ML-KEM key pair from a random seed.+generate :: (ParamSet a, MonadRandom m)+         => proxy a -> m (EncapsulationKey a, DecapsulationKey a)+generate p = do+    d <- getRandomBytes 32+    z <- getRandomBytes 32+    return (Internal.keyGen p (d :: ScrubbedBytes) z)++-- | Generate an ML-KEM key pair from the specified seed (d, z).  Length of+-- inputs must be 32 bytes.+generateWith :: (ParamSet a, ByteArrayAccess d, ByteArrayAccess z)+             => proxy a -> d -> z -> Maybe (EncapsulationKey a, DecapsulationKey a)+generateWith p d z+    | B.length d /= 32 = Nothing+    | B.length z /= 32 = Nothing+    | otherwise = Just $ Internal.keyGen p d (B.convert z)++-- | Generate a shared secret key and an associated ciphertext using randomness.+encapsulate :: (ParamSet a, MonadRandom m)+            => EncapsulationKey a -> m (SharedSecret a, Ciphertext a)+encapsulate ek = do+    m <- getRandomBytes 32+    return (Internal.encaps ek (m :: ScrubbedBytes))++-- | Generate a shared secret key and an associated ciphertext using a+-- specified random input.  This byte array must be 32 bytes and not repeated+-- with other encapsulations.  For testing purposes.+encapsulateWith :: (ParamSet a, ByteArrayAccess m)+                => EncapsulationKey a -> m -> Maybe (SharedSecret a, Ciphertext a)+encapsulateWith ek m+    | B.length m /= 32 = Nothing+    | otherwise = Just $ Internal.encaps ek m++-- | Return the shared secret for a given ciphertext.  Does implicit rejection+-- in the event the ciphertext or encapsulation key have been tampered with.+decapsulate :: ParamSet a => DecapsulationKey a -> Ciphertext a -> SharedSecret a+decapsulate = Internal.decaps++-- | Try to detect corruptions in a pair of keys.  Note that this does not+-- fully guarantee that the key pair was properly generated.  Returns @True@+-- when the key pair is found valid.+checkKeyPair :: (ParamSet a, MonadRandom m)+             => (EncapsulationKey a, DecapsulationKey a) -> m Bool+checkKeyPair (ek, dk) = do+    m <- getRandomBytes 32+    let (kk, ct) = Internal.encaps ek (m :: ScrubbedBytes)+        kk' = Internal.decaps dk ct+    return (kk' == kk)
+ src/Internal.hs view
@@ -0,0 +1,188 @@+-- |+-- Module      : Internal+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- ML-KEM main internal algorithms+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE TypeFamilies #-}+module Internal+    ( ParamSet(..), Params(..), Encode(..), Decode(..)+    , DecapsulationKey, EncapsulationKey, Ciphertext, SharedSecret+    , keyGen, toPublic, encaps, decaps+    ) where++import Basement.Nat++import Control.DeepSeq (NFData(..))+import Control.Monad++import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes)+import qualified Data.ByteArray as B++import qualified Builder+import qualified Crypto+import qualified K_PKE as K+import K_PKE (Params(..))+import Marking (Leak(..))++-- | The class of ML-KEM parameter sets.+class KnownNat (K a) => ParamSet a where+    type K a :: Nat+    getParams :: proxy a -> Params (K a)++-- | Utility class to serialize ML-KEM objects to byte arrays.+class Encode obj where+    -- | Serializes an object to a sequence of bytes.+    encode :: ByteArray ba => obj a -> ba++-- | Utility class to deserialize ML-KEM objects from byte arrays.+class Decode obj where+    -- | Deserializes an object from a sequence of bytes.+    decode :: (ParamSet a, ByteArrayAccess ba) => proxy a -> ba -> Maybe (obj a)++-- | An ML-KEM decapsulation key, aka private key.+data DecapsulationKey a = DK (K.DecryptionKey (K a)) (K.EncryptionKey (K a)) Bytes ScrubbedBytes++-- | An ML-KEM encapsulation key, aka public key.+data EncapsulationKey a = EK Bytes (K.EncryptionKey (K a))++-- | The ciphertext produced by the encapsulation function and consumed by the+-- decapsulation function.+newtype Ciphertext a = C Bytes deriving (Eq, ByteArrayAccess)++-- | A shared secret returned by the encapsulation and decapsulation functions.+-- Length is 32 bytes for all defined parameter sets.+newtype SharedSecret a = S ScrubbedBytes deriving ByteArrayAccess++instance Eq (DecapsulationKey a) where+    DK dk1 ek1 h1 z1 == DK dk2 ek2 h2 z2 = Crypto.toBool $+        Crypto.constEqW dk1 dk2 `Crypto.andW`+        Crypto.constEqW ek1 ek2 `Crypto.andW`+        Crypto.constEqW h1 h2 `Crypto.andW`+        Crypto.constEqW z1 z2++instance Eq (EncapsulationKey a) where+    EK _ ek1 == EK _ ek2 = Crypto.toBool $ Crypto.constEqW ek1 ek2++instance Eq (SharedSecret a) where+    S a == S b = Crypto.toBool $ Crypto.constEqW a b++instance Show (DecapsulationKey a) where+#ifdef ML_KEM_TESTING+    showsPrec d dk = showParen (d > 10) $+        showString "DecapsulationKey " . showsPrec 11 (encode dk :: Bytes)+#else+    showsPrec _ _ = showString "DecapsulationKey"+#endif++instance Show (EncapsulationKey a) where+    showsPrec d ek = showParen (d > 10) $+        showString "EncapsulationKey " . showsPrec 11 (encode ek :: Bytes)++instance Show (Ciphertext a) where+    showsPrec d (C ct) = showParen (d > 10) $+        showString "Ciphertext " . showsPrec 11 ct++instance Show (SharedSecret a) where+#ifdef ML_KEM_TESTING+    showsPrec d (S kk) = showParen (d > 10) $+        showString "SharedSecret " . showsPrec 11 kk+#else+    showsPrec _ _ = showString "SharedSecret"+#endif++instance NFData (DecapsulationKey a) where+    rnf (DK dk ek h z) = rnf dk `seq` rnf ek `seq` rnf h `seq` rnf z++instance NFData (EncapsulationKey a) where+    rnf (EK _ ek) = rnf ek  -- h omitted because just for caching++instance NFData (Ciphertext a) where+    rnf (C c) = rnf c++instance NFData (SharedSecret a) where+    rnf (S kk) = rnf kk++instance Encode EncapsulationKey where+    encode (EK _ ek) = Builder.runRelaxed $ K.ekEncode ek++instance Decode EncapsulationKey where+    decode p input = EK (Crypto.h input) <$> K.ekDecode params input+      where params = getParams p++instance Encode DecapsulationKey where+    encode (DK dk ek h z) = Builder.runRelaxed $+        leak (K.dkEncode dk) <> K.ekEncode ek <> Builder.bytes h <> leak (Builder.bytes z)++instance Decode DecapsulationKey where+    decode p input = do+        -- decapsulation key type check:+        guard (B.length input == 768 * k + 96)+        let dks = B.view input 0 (384 * k)+            eks = B.view input (384 * k) (384 * k + 32)+            !h = B.convert $ B.view input (768 * k + 32) 32+        -- hash check:+        guard (Crypto.toBool $ Crypto.constEqW h (Crypto.h eks))+        let !dk = K.dkDecode dks+        !ek <- K.ekDecode params eks+        let !z = B.convert $ B.view input (768 * k + 64) 32+        return (DK dk ek h z)+      where+        params = getParams p+        k = K.dimension params++instance Decode Ciphertext where+    decode p input+        -- ciphertext type check:+        | B.length input == 32 * (du * k + dv) = Just (C $ B.convert input)+        | otherwise = Nothing+      where+        params@Params{..} = getParams p+        k = K.dimension params++instance Decode SharedSecret where+    decode _ input+        | B.length input == 32 = Just (S $ B.convert input)+        | otherwise = Nothing++-- Uses randomness to generate an encapsulation key and a corresponding decapsulation key+keyGen :: (ParamSet a, ByteArrayAccess d) => proxy a -> d -> ScrubbedBytes -> (EncapsulationKey a, DecapsulationKey a)+keyGen p d z = (EK h ek, DK dk ek h z)+  where+    params = getParams p+    (ek, dk) = K.keyGen params d+    h = Crypto.h $ Builder.run (K.ekEncode ek)++-- | Returns the encapsulation key embedded in the given decapsulation key.+-- Note that they may not necessarily match when the decapsulation key was+-- decoded from an untrusted source.+toPublic :: DecapsulationKey a -> EncapsulationKey a+toPublic (DK _ ek h _) = EK h ek++-- Uses the encapsulation key and randomness to generate a key and an associated ciphertext+encaps :: (ParamSet a, ByteArrayAccess m) => EncapsulationKey a -> m -> (SharedSecret a, Ciphertext a)+encaps p@(EK h ek) m = (S kk, C c)+  where+    params = getParams p+    (kk, r) = Crypto.g (m `Crypto.append` h)+    c = K.encrypt params ek m r++-- Uses the decapsulation key to produce a shared secret key from a ciphertext+decaps :: ParamSet a => DecapsulationKey a -> Ciphertext a -> SharedSecret a+decaps p@(DK dk ek h z) (C c) = S $+    Crypto.constSelectBytes+        (Crypto.constEqW c c')            -- condition+        kk'                               -- when equal+        (Crypto.j (z `Crypto.append` c))  -- when different+  where+    params    = getParams p+    m'        = K.decrypt params dk c+    (kk', r') = Crypto.g (m' `Crypto.append` h)+    c'        = K.encrypt params ek m' r'
+ src/K_PKE.hs view
@@ -0,0 +1,143 @@+-- |+-- Module      : K_PKE+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- The K-PKE component scheme+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE RecordWildCards #-}+module K_PKE+    ( Params(..), dimension, keyGen, encrypt, decrypt+    , DecryptionKey, dkEncode, dkDecode+    , EncryptionKey, ekEncode, ekDecode+    ) where++import Basement.Nat+import Basement.Types.OffsetSize++import Control.DeepSeq (NFData(..))+import Control.Monad++import Data.ByteArray (ByteArrayAccess, Bytes, ScrubbedBytes)+import qualified Data.ByteArray as B++import Unsafe.Coerce++import Auxiliary (Rq, Tq, (..+), (..-))+import Builder (Builder)+import Marking (SecurityMarking(..), Leak(..))+import Vector (Vector)+import qualified Auxiliary as Aux+import qualified Crypto+import qualified Builder+import Math+import qualified Matrix+import qualified Vector++data Params (k :: Nat) = Params+    { eta1 :: {-# UNPACK #-} !Word+    , eta2 :: {-# UNPACK #-} !Word+    , du :: {-# UNPACK #-} !Int+    , dv :: {-# UNPACK #-} !Int+    }++dimension :: KnownNat k => Params k -> Int+dimension = fromIntegral . natVal++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++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)) }++instance Crypto.ConstEqW (DecryptionKey k) where+    constEqW a b = Crypto.constEqW (dkS a) (dkS b)++instance Crypto.ConstEqW (EncryptionKey k) where+    constEqW a b = Crypto.constEqW (ekT a) (ekT b) `Crypto.andW` Crypto.constEqW (ekRho a) (ekRho b)++instance NFData (DecryptionKey k) where+    rnf = Vector.toNormalForm . dkS++instance NFData (EncryptionKey k) where+    rnf ek = Vector.toNormalForm (ekT ek) `seq` rnf (ekRho ek)+    -- ekA omitted because just for caching++ekEncode :: EncryptionKey k -> Builder Pub+ekEncode ek = Vector.concatMap Aux.byteEncode12 (ekT ek) <> Builder.bytes (ekRho ek)++ekDecode :: (KnownNat k, ByteArrayAccess ba) => Params k -> ba -> Maybe (EncryptionKey k)+ekDecode params input = do+    -- type check:+    guard (B.length input == 384 * k + 32)+    let !tt = Vector.create $ \i -> Aux.byteDecode12 (view384 i)+        !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)+    Just EncryptionKey { ekT = tt, ekRho = rho, ekA = aa }+  where+    k = dimension params+    view384 (Offset i) = B.view input (384 * i) 384++dkEncode :: DecryptionKey k -> Builder Sec+dkEncode = Vector.concatMap Aux.byteEncode12 . dkS++dkDecode :: (KnownNat k, ByteArrayAccess ba) => ba -> DecryptionKey k+dkDecode input = do+    let !dk = Vector.create $ \i -> Aux.byteDecode12 (view384 i)+     in DecryptionKey { dkS = dk }+  where+    view384 (Offset i) = B.view input (384 * i) 384++-- 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))+    !ss = Aux.ntt <$> s+    ee  = Aux.ntt <$> e+    !tt = leakVec $ Matrix.mulw aa ss ee+    ek  = EncryptionKey { ekT = tt, ekRho = rho, ekA = aa }+    dk  = DecryptionKey { dkS = ss }++-- Uses the encryption key to encrypt a plaintext message using the randomness 𝑟+encrypt :: (KnownNat k, ByteArrayAccess m, ByteArrayAccess r) => Params k -> EncryptionKey k -> m -> r -> Bytes+encrypt params@Params{..} ek m r = Builder.run (c1 <> c2)+  where+    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))+    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+    c1  = Vector.concatMap (Aux.byteEncode du . Aux.rcompress du) u+    c2  = Aux.byteEncode dv (Aux.rcompress dv v)++-- Uses the decryption key to decrypt a ciphertext+decrypt :: KnownNat k => Params k -> DecryptionKey k -> Bytes -> ScrubbedBytes+decrypt params@Params{..} dk c = Builder.run m+  where+    k  = dimension params+    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)+    m  = Aux.byteEncode1 (Aux.rcompress 1 w)
+ src/Marking.hs view
@@ -0,0 +1,139 @@+-- |+-- Module      : Marking+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Infrastructure that associates a security marking at type level to all+-- buffers created by the library.  This determines which buffers need the+-- scrubbed (Sec) or regular (Pub) variants.+--+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilyDependencies #-}+module Marking+    ( SecurityMarking(..), Classified(..), Leak(..)+    ) where++import Basement.Monad+import Basement.NormalForm+import Basement.PrimType+import Basement.Types.OffsetSize+#ifdef ML_KEM_TESTING+import qualified Basement.Compat.IsList+#endif++import Control.Monad.ST++import Data.ByteArray (Bytes, ScrubbedBytes)+import qualified Data.ByteArray as B++import Data.Kind++import Foreign.Ptr (Ptr)++import Unsafe.Coerce++import Block (Block, MutableBlock, blockIndex)+import ScrubbedBlock (ScrubbedBlock)+import qualified Block+import qualified ByteArrayST as ST+import qualified ScrubbedBlock++data SecurityMarking = Sec | Pub  -- secret or public information++-- Transformation called only at expected location in the LWE problem, after+-- adding noise to secret information.+--+-- Block and ScrubbedBlock have the same representation, we can force coercion+-- from Sec to Pub even though the block will be actually scrubbed.  This is+-- simpler than copying to a real non-scrubbed block.+class Leak t where+    leak :: t Sec -> t Pub+    leak = unsafeCoerce++class Classified (marking :: SecurityMarking) where+    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++    type SecureBytes marking = bytes | bytes -> marking+    unsafeCreate :: Int -> (forall s. Ptr a -> ST s ()) -> SecureBytes marking+    lengthBytes :: SecureBytes marking -> Int+    copyByteArrayToPtr :: SecureBytes marking -> Ptr a -> IO ()++instance Classified Pub where+    type SecureBlock Pub = Block++    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++    type SecureBytes Pub = Bytes+    unsafeCreate = ST.unsafeCreate+    {-# INLINE unsafeCreate #-}+    lengthBytes = B.length+    copyByteArrayToPtr = B.copyByteArrayToPtr++instance Classified Sec where+    type SecureBlock Sec = ScrubbedBlock++    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++    type SecureBytes Sec = ScrubbedBytes+    unsafeCreate = ST.unsafeCreate+    {-# INLINE unsafeCreate #-}+    lengthBytes = B.length+    copyByteArrayToPtr = B.copyByteArrayToPtr
+ src/Math.hs view
@@ -0,0 +1,40 @@+-- |+-- Module      : Math+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Type classes that define additive and multiplicative operations, as well as+-- a multiply-then-add operation that will often optimize chaining.+--+-- The module also defines a non-homogenous multiplication that combines+-- typically a public operand (left) with a secret operand (right), producing a+-- secret output.+--+{-# LANGUAGE MultiParamTypeClasses #-}+module Math+    ( Add(..), Mul(..), MulAdd(..), BiMul(..), BiMulAdd(..)+    ) where++infixl 7 .*, ..*+infixl 6 .+, .-++class Add a where+    zero :: a+    (.+) :: a -> a -> a+    (.-) :: a -> a -> a+    neg :: a -> a++class Add a => Mul a where+    one :: a+    (.*) :: a -> a -> a++class Mul a => MulAdd a where+    -- invariant: mulAdd a b c == a .* b .+ c+    mulAdd :: a -> a -> a -> a++class Add a => BiMul b a where+    (..*) :: b -> a -> a++class BiMul b a => BiMulAdd b a where+    -- invariant: biMulAdd a b c == a ..* b .+ c+    biMulAdd :: b -> a -> a -> a
+ src/Matrix.hs view
@@ -0,0 +1,46 @@+-- |+-- Module      : Matrix+-- License     : BSD-3-Clause+-- 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.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+module Matrix+    ( create, mulw, muly, mulz+#ifdef ML_KEM_TESTING+    , transpose+#endif+    ) where++import Basement.Nat+import Basement.Types.OffsetSize++import Math+import Vector (Vector)+import qualified Vector++create :: (KnownNat m, KnownNat n) => (Offset ty -> Offset (Vector n ty) -> ty) -> Vector m (Vector n ty)+create f = Vector.create $ \j -> Vector.create (`f` j)+{-# INLINE create #-}++index :: Vector m (Vector n ty) -> Offset ty -> Offset (Vector n ty) -> ty+index a i j = Vector.index (Vector.index a j) i++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++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) (..*)++#ifdef ML_KEM_TESTING+transpose :: (KnownNat m, KnownNat n) => Vector m (Vector n ty) -> Vector n (Vector m ty)+transpose a = create $ \(Offset j) (Offset i) -> index a (Offset i) (Offset j)+#endif
+ src/ScrubbedBlock.hs view
@@ -0,0 +1,158 @@+-- |+-- Module      : ScrubbedBlock+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- A block that is always pinned in memory and automatically erased by a+-- finalizer when not referenced anymore.  Same pattern as ScrubbedBytes from+-- package memory but for blocks.+--+{-# 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+    , erasePtr+    ) where++import Basement.Block (Block(..), MutableBlock(..), isPinned)+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++import Control.Exception (assert)+import Control.Monad.ST++import Data.Word++import Foreign.Ptr (Ptr)++import Block (blockIndex, blockWrite)+import qualified Block++#if MIN_VERSION_base(4,19,0)+import GHC.Base (Int(I#), setAddrRange#)+import GHC.Exts (Ptr(Ptr))+#else+import Data.Memory.PtrMethods (memSet)+#endif+import GHC.Base (IO(IO), setByteArray#)+import GHC.Exts (getSizeofMutableByteArray#, mkWeak#)++newtype ScrubbedBlock ty = ScrubbedBlock (Block ty)+    deriving (Eq, Show, NormalForm)++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) =+    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+    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)++scrub :: Block Word8 -> IO ()+scrub b = Block.unsafeThaw b >>= erase++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, () #)++erasePtr :: Int -> Ptr Word8 -> IO ()+#if MIN_VERSION_base(4,19,0)+erasePtr (I# n) (Ptr addr) = IO $ \s1 ->+    case setAddrRange# addr n 0# s1 of+        s2 -> (# s2, () #)+#else+erasePtr n ptr = memSet ptr 0 n+#endif
+ src/SecureBlock.hs view
@@ -0,0 +1,35 @@+-- |+-- Module      : SecureBlock+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Use a type-level annotation to decide between a scrubbed block or a regular+-- block+--+{-# LANGUAGE CPP #-}+module SecureBlock+    ( SecureBlock, create, index, iterModify, Marking.map, new, thaw+    , unsafeCast, unsafeFreeze, toNormalForm+#ifdef ML_KEM_TESTING+    , eq, Marking.showsPrec, fromList, Marking.replicate, toList+    , SecureBlock.length+#endif+    ) where++import Basement.Monad+import Basement.PrimType+#ifdef ML_KEM_TESTING+import Basement.Types.OffsetSize+#endif++import Block (MutableBlock)+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+length = lengthBlock+#endif
+ src/SecureBytes.hs view
@@ -0,0 +1,16 @@+-- |+-- Module      : SecureBytes+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Use a type-level annotation to decide between a scrubbed byte array or+-- a regular byte array+--+module SecureBytes+    ( SecureBytes, unsafeCreate, SecureBytes.length, copyByteArrayToPtr+    ) where++import Marking++length :: Classified marking => SecureBytes marking -> Int+length = lengthBytes
+ src/Vector.hs view
@@ -0,0 +1,100 @@+-- |+-- Module      : Vector+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- A vector of lifted elements with the vector dimension at type level.+-- Currently backed by type t'Array' from basement.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+module Vector+    ( Vector, Vector.concatMap+    , Vector.fold1ZipWith, Vector.foldIndexWith, Vector.toNormalForm+    , Vector.create, Vector.index+#ifdef ML_KEM_TESTING+    , Vector.replicateM, Vector.zipWith+#endif+    ) where++import Basement.BoxedArray (Array)+import qualified Basement.BoxedArray as Array+import Basement.Compat.IsList+import Basement.Nat+import Basement.NormalForm+import Basement.Types.OffsetSize++import Control.DeepSeq (NFData(..))+#ifdef ML_KEM_TESTING+import Control.Monad+#endif++#if !(MIN_VERSION_base(4,20,0))+import Data.List (foldl')+#endif+import Data.Proxy++import Math++newtype Vector (n :: Nat) a = Vector { unVector :: Array a }+    deriving (Eq, Show, Functor, NormalForm)++instance (Add a, KnownNat n) => Add (Vector n a) where+    zero = create (const zero)+    (.+) = Vector.zipWith (.+)+    (.-) = Vector.zipWith (.-)+    neg (Vector a) = Vector (fmap neg a)++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)+{-# INLINE create #-}++arrayIndex :: Array a -> Offset a -> a+#ifdef ML_KEM_TESTING+arrayIndex = Array.index++replicateM :: forall n m a. (KnownNat n, Applicative m) => m a -> m (Vector n a)+replicateM f = Vector . fromList <$> Control.Monad.replicateM sz f+  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+#else+arrayIndex = Array.unsafeIndex+#endif++index :: Vector n a -> Offset a -> a+index = arrayIndex . unVector++concatMap :: Monoid b => (a -> b) -> Vector n a -> b+concatMap f = mconcat . Prelude.map f . toList . unVector+{-# INLINE concatMap #-}++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 #-}++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]+  where+    ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i))+    gg = g (arrayIndex a 0) (arrayIndex b 0)+    CountOf !sa = Array.length a+{-# INLINE fold1ZipWith #-}++foldIndexWith :: (c -> Offset a -> a -> c) -> c -> Vector n a -> c+foldIndexWith f c (Vector a) = foldl' g c [0 .. sa - 1]+  where+    g x i = f x (Offset i) (arrayIndex a (Offset i))+    CountOf !sa = Array.length a+{-# INLINE foldIndexWith #-}++toNormalForm :: NFData a => Vector n a -> ()+toNormalForm = Array.foldl' (\acc x -> acc `seq` rnf x) () . unVector
+ tests/EncapDecap.hs view
@@ -0,0 +1,127 @@+-- |+-- Module      : EncapDecap+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Types that are specific to encapsulation/decapsulation test vectors+--+{-# LANGUAGE OverloadedStrings #-}+module EncapDecap+    ( TestGroup(..), Test(..), TestGroupPayload(..)+    , EncapsulationExt(..), DecapsulationExt(..)+    , EncapsulationKeyCheckExt(..), DecapsulationKeyCheckExt(..)+    ) where++import Data.Aeson+import Data.Aeson.Types+import Data.ByteString (ByteString)++import Util++data TestGroup = TestGroup+    { tgId :: Int+    , testType :: String+    , parameterSet :: String+    , function :: String+    , payload :: TestGroupPayload+    } deriving Show++data TestGroupPayload+    = FunctionEncapsulation [Test EncapsulationExt]+    | FunctionDecapsulation [Test DecapsulationExt]+    | FunctionEncapsulationKeyCheck [Test EncapsulationKeyCheckExt]+    | FunctionDecapsulationKeyCheck [Test DecapsulationKeyCheckExt]+    deriving Show++parsePayload :: Object -> Parser TestGroupPayload+parsePayload o = do+    fn <- o .: "function"+    case fn of+        "encapsulation" -> FunctionEncapsulation <$> (o .: "tests")+        "decapsulation" -> FunctionDecapsulation <$> (o .: "tests")+        "encapsulationKeyCheck" -> FunctionEncapsulationKeyCheck <$> (o .: "tests")+        "decapsulationKeyCheck" -> FunctionDecapsulationKeyCheck <$> (o .: "tests")+        unknown -> fail ("parsePayload: unknown function " ++ unknown)++instance FromJSON TestGroup where+    parseJSON = withObject "TestGroup" $ \o -> TestGroup+        <$> o .: "tgId"+        <*> o .: "testType"+        <*> o .: "parameterSet"+        <*> o .: "function"+        <*> parsePayload o++data Test ext = Test+    { tcId :: Int+    , deferred :: Bool+    , tcExt :: ext+    } deriving Show++class TestExt ext where+    parseExt :: Object -> Parser ext++instance TestExt ext => FromJSON (Test ext) where+    parseJSON = withObject "Test" $ \o -> Test+        <$> o .: "tcId"+        <*> o .: "deferred"+        <*> parseExt o++data EncapsulationExt = EncapsulationExt+    { ekEnc :: ByteString+    , dkEnc :: ByteString+    , cEnc :: ByteString+    , kEnc :: ByteString+    , mEnc :: ByteString+    } deriving Show++instance TestExt EncapsulationExt where+    parseExt o = EncapsulationExt+        <$> o .:: "ek"+        <*> o .:: "dk"+        <*> o .:: "c"+        <*> o .:: "k"+        <*> o .:: "m"++data DecapsulationExt = DecapsulationExt+    { ekDec :: ByteString+    , dkDec :: ByteString+    , cDec :: ByteString+    , kDec :: ByteString+    , reasonDec :: String+    } deriving Show++instance TestExt DecapsulationExt where+    parseExt o = DecapsulationExt+        <$> o .:: "ek"+        <*> o .:: "dk"+        <*> o .:: "c"+        <*> o .:: "k"+        <*> o .:  "reason"++data DecapsulationKeyCheckExt = DecapsulationKeyCheckExt+    { passedDkc :: Bool+    , ekDkc :: ByteString+    , dkDkc :: ByteString+    , reasonDkc :: String+    } deriving Show++instance TestExt DecapsulationKeyCheckExt where+    parseExt o = DecapsulationKeyCheckExt+        <$> o .:  "testPassed"+        <*> o .:: "ek"+        <*> o .:: "dk"+        <*> o .:  "reason"++data EncapsulationKeyCheckExt = EncapsulationKeyCheckExt+    { passedEkc :: Bool+    , ekEkc :: ByteString+    , dkEkc :: ByteString+    , reasonEkc :: String+    } deriving Show++instance TestExt EncapsulationKeyCheckExt where+    parseExt o = EncapsulationKeyCheckExt+        <$> o .:  "testPassed"+        <*> o .:: "ek"+        <*> o .:: "dk"+        <*> o .:  "reason"
+ tests/KeyGen.hs view
@@ -0,0 +1,48 @@+-- |+-- Module      : KeyGen+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Types that are specific to key-generation test vectors+--+{-# LANGUAGE OverloadedStrings #-}+module KeyGen+    ( TestGroup(..), Test(..)+    ) where++import Data.Aeson+import Data.ByteString (ByteString)++import Util++data TestGroup = TestGroup+    { tgId :: Int+    , testType :: String+    , parameterSet :: String+    , tests :: [Test]+    } deriving Show++instance FromJSON TestGroup where+    parseJSON = withObject "TestGroup" $ \o -> TestGroup+        <$> o .: "tgId"+        <*> o .: "testType"+        <*> o .: "parameterSet"+        <*> o .: "tests"++data Test = Test+    { tcId :: Int+    , deferred :: Bool+    , d :: ByteString+    , z :: ByteString+    , ek :: ByteString+    , dk :: ByteString+    } deriving Show++instance FromJSON Test where+    parseJSON = withObject "Test" $ \o -> Test+        <$> o .: "tcId"+        <*> o .: "deferred"+        <*> o .:: "d"+        <*> o .:: "z"+        <*> o .:: "ek"+        <*> o .:: "dk"
+ tests/Tests.hs view
@@ -0,0 +1,468 @@+-- |+-- Module      : Main+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- The ML-KEM test suite.  Can be instanciated twice, with and without the+-- @ML_KEM_TESTING@ macro to run property testing with assertions enabled in+-- the internal modules.+--+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE KindSignatures #-}+module Main (main) where++import Data.ByteArray (Bytes)+import qualified Data.ByteArray as B++import Control.Monad++import Data.Maybe (isJust, fromJust)+import Data.Proxy++import GHC.IO.Exception (IOErrorType(..))++import System.Directory (doesFileExist)+import System.IO.Error (catchIOError, mkIOError)+import System.Process (readProcess)++#ifdef ML_KEM_TESTING++import Data.Bits+import Data.Word++import GHC.TypeNats++import Auxiliary+import BlockN (BlockN)+import Builder (Builder)+import Marking (Leak(..), SecurityMarking(..))+import Math+import Matrix+import Vector+import qualified BlockN+import qualified Builder+#endif++import Crypto.PubKey.ML_KEM as Lib++import Test.Tasty+import Test.Tasty.HUnit+import Test.Tasty.QuickCheck++import qualified EncapDecap+import qualified KeyGen+import qualified Vectors++#ifdef ML_KEM_TESTING++newtype Bit7 = Bit7 Word8 deriving Show++instance Arbitrary Bit7 where+#if (MIN_VERSION_tasty_quickcheck(0,10,2))+    arbitrary = Bit7 <$> chooseBoundedIntegral (0, 127)+#else+    arbitrary = Bit7 <$> choose (0, 127)+#endif++newtype FE = FE { unFE :: Zq} deriving Show++instance Arbitrary FE where+#if (MIN_VERSION_tasty_quickcheck(0,10,2))+    arbitrary = FE . toZq <$> chooseBoundedIntegral (0, 3328)+#else+    arbitrary = FE . toZq <$> choose (0, 3328)+#endif++newtype Poly = Poly (Rq Sec) deriving Show++instance Arbitrary Poly where+    arbitrary = do+      coeffs <- map unFE <$> vectorOf 256 arbitrary+      let a = fromJust (fromCoeffs coeffs)+      return (Poly a)++newtype PolyNTT = PolyNTT (Tq Sec) deriving Show++instance Arbitrary PolyNTT where+    arbitrary = (\(Poly f) -> PolyNTT (ntt f)) <$> arbitrary++newtype D = D Int deriving Show++instance Arbitrary D where+    arbitrary = D <$> choose (0, 12)++data Dim = forall (n :: Nat). KnownNat n => Dim (Proxy n)++instance Show Dim where+    show (Dim n) = show n++instance Arbitrary Dim where+    arbitrary = toDim <$> choose (1, 9)++toDim :: Int -> Dim+toDim n = case someNatVal (fromIntegral n) of SomeNat p -> Dim p++type VElem = Zq  -- test with any ring but Tq would also work here++arbitraryVector :: KnownNat n => proxy n -> Gen (Vector n VElem)+arbitraryVector _ = Vector.replicateM (unFE <$> arbitrary)++arbitraryMatrix :: (KnownNat m, KnownNat n) => proxy n -> proxy m -> Gen (Vector n (Vector m VElem))+arbitraryMatrix _ m = Vector.replicateM (arbitraryVector m)++arbitraryBytes :: Int -> Gen Bytes+arbitraryBytes n = B.pack <$> vectorOf n arbitrary++byteDecodeBytes :: Int -> Bytes -> BlockN Sec 256 Word16+byteDecodeBytes = byteDecode++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++#endif++data P = forall a. (ParamSet a, Show a) => P (Proxy a)++instance Show P where+    show (P p) = show p++instance Arbitrary P where+    arbitrary = elements+        [ P (Proxy :: Proxy ML_KEM_512)+        , P (Proxy :: Proxy ML_KEM_768)+        , P (Proxy :: Proxy ML_KEM_1024)+        ]++toP :: String -> P+toP "ML-KEM-512" =  P (Proxy :: Proxy ML_KEM_512)+toP "ML-KEM-768" =  P (Proxy :: Proxy ML_KEM_768)+toP "ML-KEM-1024" = P (Proxy :: Proxy ML_KEM_1024)+toP paramSet      = error ("unknown parameter set " ++ paramSet)++withVectors :: (IO () -> TestTree) -> TestTree+withVectors = withResource alloc free+  where+    scriptPath = "tests/get-vectors.sh"+    free _ = return ()+    alloc = do+        keyGenExists <- doesFileExist "tests/keyGen.json.gz"+        encapDecapExists <- doesFileExist "tests/encapDecap.json.gz"+        unless (keyGenExists && encapDecapExists) $ catchIOError+            (void $ readProcess "/bin/sh" [scriptPath] "")+            (\e ->+                let msg = "Could not download test vectors, you will need to run the script `" +++                            scriptPath ++ "' manually. Script failure was: " ++ show e+                 in ioError (mkIOError OtherError msg Nothing Nothing)+            )++keyGenVectors :: (String -> IO ()) -> Assertion+keyGenVectors step = do+    step "Reading test vectors ..."+    file <- Vectors.readJson "tests/keyGen.json.gz"+    forM_ (Vectors.testGroups file) $ \group -> do+        let paramSet = KeyGen.parameterSet group+        step paramSet+        case toP paramSet of+            P p -> forM_ (KeyGen.tests group) $ \t -> do+                let tcId = KeyGen.tcId t+                    eks = Lib.encode ek+                    dks = Lib.encode dk+                    (ek, dk) = fromJust $ Lib.generateWith p (KeyGen.d t) (KeyGen.z t)+                assertEqual ("ek mismatch for tcId=" ++ show tcId) (KeyGen.ek t) eks+                assertEqual ("dk mismatch for tcId=" ++ show tcId) (KeyGen.dk t) dks++encapDecapVectors :: (String -> IO ()) -> Assertion+encapDecapVectors step = do+    step "Reading test vectors ..."+    file <- Vectors.readJson "tests/encapDecap.json.gz"+    forM_ (Vectors.testGroups file) $ \group -> do+        let paramSet = EncapDecap.parameterSet group+        step (paramSet ++ " (" ++ EncapDecap.function group ++ ")")+        case toP paramSet of+            P p -> case EncapDecap.payload group of+                EncapDecap.FunctionEncapsulation tests ->+                    forM_ tests (testEncapsulation p)+                EncapDecap.FunctionDecapsulation tests ->+                    forM_ tests (testDecapsulation p)+                EncapDecap.FunctionEncapsulationKeyCheck tests ->+                    forM_ tests (testEncapsulationKeyCheck p)+                EncapDecap.FunctionDecapsulationKeyCheck tests ->+                    forM_ tests (testDecapsulationKeyCheck p)+  where+    ensureEk = id :: f (EncapsulationKey a) -> f (EncapsulationKey a)+    ensureDk = id :: f (DecapsulationKey a) -> f (DecapsulationKey a)+    testEncapsulation p test = do+        let tcId = EncapDecap.tcId test+            ext = EncapDecap.tcExt test+            ek = fromJust $ Lib.decode p (EncapDecap.ekEnc ext)+            k' = fromJust $ Lib.decode p (EncapDecap.kEnc ext)+            c' = fromJust $ Lib.decode p (EncapDecap.cEnc ext)+            (k, c) = fromJust $ Lib.encapsulateWith ek (EncapDecap.mEnc ext)+        assertEqual ("k mismatch for tcId=" ++ show tcId) k' k+        assertEqual ("c mismatch for tcId=" ++ show tcId) c' c+    testDecapsulation p test = do+        let tcId = EncapDecap.tcId test+            ext = EncapDecap.tcExt test+            dk = fromJust $ Lib.decode p (EncapDecap.dkDec ext)+            c  = fromJust $ Lib.decode p (EncapDecap.cDec ext)+            k' = fromJust $ Lib.decode p (EncapDecap.kDec ext)+            k = Lib.decapsulate dk c+        assertEqual ("k mismatch for tcId=" ++ show tcId) k' k+    testEncapsulationKeyCheck p test = do+        let tcId = EncapDecap.tcId test+            ext = EncapDecap.tcExt test+            mek = ensureEk $ Lib.decode p (EncapDecap.ekEkc ext)+        assertBool ("opposite outcome for tcId=" ++ show tcId)+            (EncapDecap.passedEkc ext == isJust mek)+    testDecapsulationKeyCheck p test = do+        let tcId = EncapDecap.tcId test+            ext = EncapDecap.tcExt test+            mdk = ensureDk $ Lib.decode p (EncapDecap.dkDkc ext)+        assertBool ("opposite outcome for tcId=" ++ show tcId)+            (EncapDecap.passedDkc ext == isJust mdk)++main :: IO ()+main = defaultMain $ testGroup "mlkem"+    [ withVectors $ \_ -> testGroup "vectors"+        [ testCaseSteps "keyGen" keyGenVectors+        , testCaseSteps "encapDecap" encapDecapVectors+        ]+    , testGroup "properties"+        [ testGroup "ML-KEM"+            [ testProperty "encapsulate/decapsulate" $ \(P p) -> ioProperty $ do+                (ek, dk) <- Lib.generate p+                (kk, c) <- Lib.encapsulate ek+                let kk' = Lib.decapsulate dk c+                return (kk === kk')+            , testProperty "encode/decode keys" $ \(P p) -> ioProperty $ do+                (ek, dk) <- Lib.generate p+                return $ conjoin+                    [ Just ek === Lib.decode p (Lib.encode ek :: Bytes)+                    , Just dk === Lib.decode p (Lib.encode dk :: Bytes)+                    ]+            , testProperty "convert/decode ciphertext and shared secret" $ \(P p) -> ioProperty $ do+                (ek, _) <- Lib.generate p+                (kk, c) <- Lib.encapsulate ek+                return $ conjoin+                    [ Just c === Lib.decode p (B.convert c :: Bytes)+                    , Just kk === Lib.decode p (B.convert kk :: Bytes)+                    ]+            , testProperty "toPublic" $ \(P p) -> ioProperty $ do+                (ek, dk) <- Lib.generate p+                return (ek === toPublic dk)+            , testProperty "checkKeyPair" $ \(P p) -> ioProperty $+                Lib.generate p >>= checkKeyPair+            ]+#ifdef ML_KEM_TESTING+        , testGroup "bitRev7"+            [ testCase "powers of two" $+                let powers = [1, 2, 4, 8, 16, 32, 64]+                 in reverse powers @=? map bitRev7 powers+            , testProperty "or" $ \(Bit7 a) (Bit7 b) ->+                bitRev7 (a .|. b) === bitRev7 a .|. bitRev7 b+            , testProperty "not" $ \(Bit7 a) ->+                let comp = xor 127+                 in bitRev7 (comp a) === comp (bitRev7 a)+            , testProperty "involutive" $ \(Bit7 a) ->+                a === bitRev7 (bitRev7 a)+            , testProperty "preserves bit count" $ \(Bit7 a) ->+                popCount a === popCount (bitRev7 a)+            ]+        , testGroup "compression"+            [ testProperty "compress . decompress == id" $ \(D d) -> do+                y <- choose (0, 2^d - 1)+                return (d < 12 ==> y === compress d (decompress d y))+            ]+        , testGroup "encoding"+            [ testProperty "byteEncode . byteDecode == id" $ \(D d) -> do+                b <- arbitraryBytes (32 * d)+                return (b === byteEncodeBytes d (byteDecode d b))+            , testProperty "byteEncode1 . byteDecode1 == id" $ do+                b <- arbitraryBytes 32+                return (b === byteEncodeBytes1 (byteDecode1 b))+            , testProperty "byteDecode12 . byteEncode12 == id" $ \(PolyNTT p) ->+                p === byteDecode12 (byteEncodeBytes12 p)+            , testProperty "byteEncode 8" $ \x ->+                B.replicate 256 x === byteEncodeBytes 8 (BlockN.replicate $ fromIntegral x)+            , testCase "byteEncode 1 (zeros)" $+                B.replicate 32 0 @=? byteEncodeBytes 1 (BlockN.replicate 0)+            , testCase "byteEncode 1 (ones)" $+                B.replicate 32 255 @=? byteEncodeBytes 1 (BlockN.replicate 1)+            , testProperty "byteDecode1 == byteDecode 1" $ do+                b <- arbitraryBytes 32+                return (byteDecodeBytes 1 b === byteDecode1 b)+            ]+        , testGroup "Zq"+            [ testProperty "toZq . fromZq == id " $ \(FE a) ->+                a === toZq (fromZq a)+            , testProperty "fromZq . toZq == id " $ \a ->+                mod a 3329 === fromZq (toZq a)+            , testCase "field order" $ zero @=? toZq 3329+            , testProperty "addition with zero" $ \(FE a) ->+                conjoin [ a === zero .+ a+                        , a === a .+ zero+                        ]+            , testProperty "addition associative" $ \(FE a) (FE b) (FE c) ->+                a .+ (b .+ c) === (a .+ b) .+ c+            , testProperty "addition commutative" $ \(FE a) (FE b) ->+                a .+ b === b .+ a+            , testProperty "substraction with zero" $ \(FE a) ->+                a === a .- zero+            , testProperty "substraction non-associative" $ \(FE a) (FE b) (FE c) ->+                a .- (b .- c) === (a .- b) .+ c+            , testProperty "substraction anti-commutative" $ \(FE a) (FE b) ->+                a .- b === neg (b .- a)+            , testProperty "negation" $ \(FE a) ->+                neg a === zero .- a+            , testProperty "double negation" $ \(FE a) ->+                a === neg (neg a)+            , testProperty "multiplication with zero" $ \(FE a) ->+                conjoin [ zero === zero .* a+                        , zero === a .* zero+                        ]+            , testProperty "multiplication with one" $ \(FE a) ->+                conjoin [ a === one .* a+                        , a === a .* one+                        ]+            , testProperty "multiplication associative" $ \(FE a) (FE b) (FE c) ->+                a .* (b .* c) === (a .* b) .* c+            , testProperty "multiplication commutative" $ \(FE a) (FE b) ->+                a .* b === b .* a+            , testProperty "multiplication distributive" $ \(FE a) (FE b) (FE c) ->+                conjoin [ (a .* b) .+ (a .* c) === a .* (b .+ c)+                        , (b .* a) .+ (c .* a) === (b .+ c) .* a+                        ]+            , testProperty "mulAdd" $ \(FE a) (FE b) (FE c) ->+                a .* b .+ c === mulAdd a b c+            ]+        , testGroup "Rq"+            [ testProperty "fromCoeffs . toCoeffs == id " $ \(Poly a) ->+                Just a === fromCoeffs (toCoeffs a)+            , testProperty "addition with zero" $ \(Poly a) ->+                conjoin [ a === zero .+ a+                        , a === a .+ zero+                        ]+            , testProperty "addition associative" $ \(Poly a) (Poly b) (Poly c) ->+                a .+ (b .+ c) === (a .+ b) .+ c+            , testProperty "addition commutative" $ \(Poly a) (Poly b) ->+                a .+ b === b .+ a+            , testProperty "substraction with zero" $ \(Poly a) ->+                a === a .- zero+            , testProperty "substraction non-associative" $ \(Poly a) (Poly b) (Poly c) ->+                a .- (b .- c) === (a .- b) .+ c+            , testProperty "substraction anti-commutative" $ \(Poly a) (Poly b) ->+                a .- b === neg (b .- a)+            , testProperty "negation" $ \(Poly a) ->+                neg a === zero .- a+            , testProperty "double negation" $ \(Poly a) ->+                a === neg (neg a)+            ]+        , testGroup "Tq"+            [ testProperty "nttInv . ntt == id" $ \(Poly a) ->+                a === nttInv (ntt a)+            , testProperty "addition with zero" $ \(PolyNTT a) ->+                conjoin [ a === zero .+ a+                        , a === a .+ zero+                        ]+            , testProperty "addition associative" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) ->+                a .+ (b .+ c) === (a .+ b) .+ c+            , testProperty "addition commutative" $ \(PolyNTT a) (PolyNTT b) ->+                a .+ b === b .+ a+            , testProperty "substraction with zero" $ \(PolyNTT a) ->+                a === a .- zero+            , testProperty "substraction non-associative" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) ->+                a .- (b .- c) === (a .- b) .+ c+            , testProperty "substraction anti-commutative" $ \(PolyNTT a) (PolyNTT b) ->+                a .- b === neg (b .- a)+            , testProperty "negation" $ \(PolyNTT a) ->+                neg a === zero .- a+            , testProperty "double negation" $ \(PolyNTT a) ->+                a === neg (neg a)+            , testProperty "multiplication with zero" $ \(PolyNTT a) ->+                conjoin [ zero === zero .* a+                        , zero === a .* zero+                        ]+            , testProperty "multiplication with one" $ \(PolyNTT a) ->+                conjoin [ a === one .* a+                        , a === a .* one+                        ]+            , testProperty "multiplication associative" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) ->+                a .* (b .* c) === (a .* b) .* c+            , testProperty "multiplication commutative" $ \(PolyNTT a) (PolyNTT b) ->+                a .* b === b .* a+            , testProperty "multiplication distributive" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) ->+                conjoin [ (a .* b) .+ (a .* c) === a .* (b .+ c)+                        , (b .* a) .+ (c .* a) === (b .+ c) .* a+                        ]+            , testProperty "mulAdd" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) ->+                a .* b .+ c === mulAdd a b c+            ]+        , testGroup "Vector"+            [ testProperty "addition with zero" $ \(Dim n) -> do+                a <- arbitraryVector n+                return $ conjoin+                    [ a === zero .+ a+                    , a === a .+ zero+                    ]+            , testProperty "addition associative" $ \(Dim n) -> do+                (a, b, c) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n+                return (a .+ (b .+ c) === (a .+ b) .+ c)+            , testProperty "addition commutative" $ \(Dim n) -> do+                (a, b) <- (,) <$> arbitraryVector n <*> arbitraryVector n+                return (a .+ b === b .+ a)+            , testProperty "substraction with zero" $ \(Dim n) -> do+                a <- arbitraryVector n+                return (a === a .- zero)+            , testProperty "substraction non-associative" $ \(Dim n) -> do+                (a, b, c) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n+                return (a .- (b .- c) === (a .- b) .+ c)+            , testProperty "substraction anti-commutative" $ \(Dim n) -> do+                (a, b) <- (,) <$> arbitraryVector n <*> arbitraryVector n+                return (a .- b === neg (b .- a))+            , testProperty "negation" $ \(Dim n) -> do+                a <- arbitraryVector n+                return (neg a === zero .- a)+            , testProperty "double negation" $ \(Dim n) -> do+                a <- arbitraryVector n+                return (a === neg (neg a))+            ]+        , testGroup "Matrix"+            [ testProperty "mulw distributive left" $ \(Dim n) (Dim m) -> do+                (a, b, u, v) <- (,,,) <$> arbitraryMatrix m n <*> arbitraryMatrix m n <*> arbitraryVector m <*> arbitraryVector n+                return (mulw (a .+ b) u v === mulw a u (mulw b u zero) .+ v)+            , testProperty "mulw distributive right" $ \(Dim n) (Dim m) -> do+                (a, u, v, w) <- (,,,) <$> arbitraryMatrix m n <*> arbitraryVector m <*> arbitraryVector m <*> arbitraryVector n+                return (mulw a (u .+ v) w === mulw a u (mulw a v w))+            , testProperty "muly definition" $ \(Dim n) (Dim m) -> do+                (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector n+                return ((a `muly` u) .+ v == mulw (transpose a) u v)+            , testProperty "muly distributive left" $ \(Dim n) (Dim m) -> do+                (a, b, u) <- (,,) <$> arbitraryMatrix n m <*> arbitraryMatrix n m <*> arbitraryVector m+                return ((a .+ b) `muly` u === (a `muly` u) .+ (b `muly` u))+            , 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+        ]+    ]
+ tests/Util.hs view
@@ -0,0 +1,41 @@+-- |+-- Module      : Util+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Utility to read hexadecimal byte arrays with aeson+--+{-# LANGUAGE CPP #-}+module Util+    ( (.::)+    ) where++import Data.Aeson+import Data.Aeson.Types+import Data.Char+import Data.ByteString (ByteString, cons, empty)+import Data.Text (Text, uncons)++#if !(MIN_VERSION_aeson(2,0,0))+type Key = Text+#endif++(.::) :: Object -> Key -> Parser ByteString+o .:: name = (o .: name) >>= fromBase16++fromBase16 :: Text -> Parser ByteString+fromBase16 t = case uncons t of+    Nothing      -> return empty+    Just (a, as) ->+        case uncons as of+            Nothing      -> fail "incomplete Base16"+            Just (b, bs) -> do+                ia <- fromHexDigit a+                ib <- fromHexDigit b+                let w  = fromIntegral (ia * 16 + ib)+                cons w <$> fromBase16 bs++fromHexDigit :: Char -> Parser Int+fromHexDigit c+    | isHexDigit c = return (digitToInt c)+    | otherwise    = fail "invalid hex digit"
+ tests/Vectors.hs view
@@ -0,0 +1,41 @@+-- |+-- Module      : Vectors+-- License     : BSD-3-Clause+-- Copyright   : (c) 2025 Olivier Chéron+--+-- Common implementation of ML-KEM test vectors+--+{-# LANGUAGE OverloadedStrings #-}+module Vectors+    ( VectorFile(..), readJson+    ) where++import Data.Aeson+import qualified Data.ByteString.Lazy as L++import qualified Codec.Compression.GZip as GZip++data VectorFile tg = VectorFile+    { vsId :: Int+    , algorithm :: String+    , mode :: String+    , revision :: String+    , isSample :: Bool+    , testGroups :: [tg]+    } deriving Show++instance FromJSON tg => FromJSON (VectorFile tg) where+    parseJSON = withObject "File" $ \o -> VectorFile+        <$> o .: "vsId"+        <*> o .: "algorithm"+        <*> o .: "mode"+        <*> o .: "revision"+        <*> o .: "isSample"+        <*> o .: "testGroups"++readJson :: FromJSON tg => FilePath -> IO (VectorFile tg)+readJson path = do+    bs <- L.readFile path+    case decode (GZip.decompress bs) of+        Just file -> return file+        _         -> fail "could not parse"
+ tests/get-vectors.sh view
@@ -0,0 +1,13 @@+#!/bin/sh++DESTDIR="`dirname "$0"`"++REF=commit/cbb097019deb25eb966932ee293d89417482d8bc+CURL=curl++for KEY in keyGen encapDecap; do+    FILENAME="$DESTDIR"/$KEY.json.gz+    URL=https://codeberg.org/ocheron/hs-mlkem/raw/$REF/tests/$KEY.json.gz++    "$CURL" --silent -o "$FILENAME" "$URL" || exit $?+done