mldsa (empty) → 0.1.0.0
raw patch · 33 files changed
+4820/−0 lines, 33 filesdep +aesondep +basedep +bytestring
Dependencies added: aeson, base, bytestring, criterion, crypton, cryptonite, deepseq, directory, filelock, memory, mldsa, primitive, process, ram, tasty, tasty-hunit, tasty-quickcheck, text, zlib
Files
- CHANGELOG.md +5/−0
- LICENSE +26/−0
- README.md +55/−0
- benchs/Bench.hs +67/−0
- mldsa.cabal +206/−0
- src/Auxiliary.hs +959/−0
- src/Base.hs +62/−0
- src/Block.hs +127/−0
- src/BlockN.hs +262/−0
- src/Builder.hs +99/−0
- src/ByteArrayST.hs +46/−0
- src/Crypto.hs +214/−0
- src/Crypto/PubKey/ML_DSA.hs +334/−0
- src/Equality.hs +25/−0
- src/Fusion.hs +102/−0
- src/Internal.hs +389/−0
- src/Iterate.hs +47/−0
- src/Machine.hs +92/−0
- src/Marking.hs +133/−0
- src/Math.hs +55/−0
- src/Matrix.hs +35/−0
- src/ScrubbedBlock.hs +90/−0
- src/SecureBlock.hs +18/−0
- src/SecureBytes.hs +16/−0
- src/Vector.hs +253/−0
- tests/KeyGen.hs +46/−0
- tests/SigExt.hs +57/−0
- tests/SigGen.hs +81/−0
- tests/SigVer.hs +79/−0
- tests/Tests.hs +739/−0
- tests/Util.hs +47/−0
- tests/Vectors.hs +41/−0
- tests/get-vectors.sh +13/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Changelog for `mldsa`++## 0.1.0.0 - 2026-07-05++* First version. Released on an unsuspecting world.
+ LICENSE view
@@ -0,0 +1,26 @@+Copyright 2026 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,55 @@+# ML-DSA++[](https://en.wikipedia.org/wiki/BSD_licenses)+[](https://haskell.org/)+[](https://hackage.haskell.org/package/mldsa)++_Module-Lattice-based Digital Signature Algorithm_ implemented in Haskell.++See [FIPS 204](https://csrc.nist.gov/pubs/fips/204/final).++Example session:++```haskell+> :set -XOverloadedStrings+> import Crypto.PubKey.ML_DSA+> import Data.ByteString+> import Data.Proxy+> let params = Proxy :: Proxy ML_DSA_65+> (pubKey, privKey) <- generate params+> sigma <- sign privKey ("original message" :: ByteString) defaultContext+> verify pubKey ("original message" :: ByteString) sigma defaultContext+True+> verify pubKey ("altered message" :: ByteString) sigma defaultContext+False+```++## 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.+Cautious users can run the benchmarks with info-table profiling and verify that+closures containing non pointers capture only non-secret variables like loop+indices or algorithm parameters.++Best performance is obtained with the LLVM code generator. On ARM, define+macro `__ARM_FEATURE_UNALIGNED` if unaligned access is supported by the target.++Randomness is provided either from explicit inputs or through a user-selected+instance of the `MonadRandom` type class from `crypton`. A good implementation+would combine multiple sources of entropy, reseed periodically, and protect its+internal state in memory. Deterministic signing with no source of randomness is+possible but is not advised.++## Testing++The test suite executes all NIST test vectors but necessary files are not+included in the package to limit its size. Instead, three 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,67 @@+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+module Main (main) where++import Criterion.Main++import Crypto.Random+import Crypto.PubKey.ML_DSA++import Data.ByteArray (Bytes)+import Data.Maybe (fromJust)+import Data.Proxy++data PS = forall a . ParamSet a => PS (Proxy a)++message :: Bytes+message = mempty++psList :: [(String, PS)]+psList =+ [ ("ML-DSA-44", PS (Proxy :: Proxy ML_DSA_44))+ , ("ML-DSA-65", PS (Proxy :: Proxy ML_DSA_65))+ , ("ML-DSA-87", PS (Proxy :: Proxy ML_DSA_87))+ ]++doBench :: (String, PS) -> Benchmark+doBench (name, PS p) = bgroup name+ [ bench "generate" $ perRunEnv setupGenerate (return . runGenerate)+ , bench "sign" $ perRunEnv setupSign runSign+ , bench "sign (batch)" $ perBatchEnv (const setupSign) runSign+ , bench "verify" $ perRunEnv setupVerify (return . runVerify)+ , bench "verify (batch)" $ perBatchEnv (const setupVerify) (return . runVerify)+ , env (generate p) $ bench "encoding (public)" . nf (encodeDecode p) . fst+ , env (generate p) $ bench "encoding (private)" . nf (encodeDecode p) . snd+ , env setupVerify $ bench "encoding (signature)" . nf (encodeDecode p) . snd+ ]+ where+ gen32 = getRandomBytes 32 :: IO Bytes++ setupGenerate = gen32+ runGenerate = fromJust . generateWith p++ setupSign = snd <$> generate p+ runSign sk = do+ -- signature time is highly variable due to the rejection loop, so we+ -- want diversity in batch runs and use fresh randomness+ rnd <- fromJust . randomness <$> gen32+ return $ signWith rnd sk message defaultContext++ runVerify (pk, sig) = verify pk message sig defaultContext+ setupVerify = do+ (pk, sk) <- generate p+ sig <- sign sk message defaultContext+ return (pk, sig)++ encodeDecode :: (Encode obj, Decode obj, Eq (obj a), ParamSet a)+ => proxy a -> obj a -> Bool+ encodeDecode prx k = decode prx (aligned $ encode k) == Just k++ aligned :: Bytes -> Bytes+ aligned = id++main :: IO ()+main = defaultMain+ [ bgroup "mldsa" $ map doBench psList+ ]
+ mldsa.cabal view
@@ -0,0 +1,206 @@+cabal-version: 2.2++-- This file has been generated from package.yaml by hpack version 0.39.1.+--+-- see: https://github.com/sol/hpack++name: mldsa+version: 0.1.0.0+synopsis: Module-Lattice-based Digital Signature Algorithm+description: Module-Lattice-based Digital Signature Algorithm (ML-DSA) implemented in+ Haskell.+category: Cryptography+homepage: https://codeberg.org/ocheron/hs-mldsa#readme+bug-reports: https://codeberg.org/ocheron/hs-mldsa/issues+author: Olivier Chéron+maintainer: olivier.cheron@gmail.com+copyright: 2026 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-mldsa++flag use_crypton+ description: Use crypton instead of cryptonite+ manual: True+ default: True++library+ exposed-modules:+ Crypto.PubKey.ML_DSA+ other-modules:+ Auxiliary+ Base+ Block+ BlockN+ Builder+ ByteArrayST+ Crypto+ Equality+ Fusion+ Internal+ Iterate+ Machine+ Marking+ Math+ Matrix+ ScrubbedBlock+ SecureBlock+ SecureBytes+ Vector+ Paths_mldsa+ autogen-modules:+ Paths_mldsa+ 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 -Wno-unticked-promoted-constructors -O2+ build-depends:+ base >=4.7 && <5+ , deepseq+ , primitive >=0.7.2+ default-language: Haskell2010+ if flag(use_crypton)+ build-depends:+ crypton >=1.1.1+ , ram+ else+ build-depends:+ cryptonite >=0.26+ , memory++test-suite mldsa-test+ type: exitcode-stdio-1.0+ main-is: Tests.hs+ other-modules:+ KeyGen+ SigExt+ SigGen+ SigVer+ Util+ Vectors+ Paths_mldsa+ autogen-modules:+ Paths_mldsa+ 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 -Wno-unticked-promoted-constructors -threaded -rtsopts -with-rtsopts=-N+ build-depends:+ aeson+ , base >=4.7 && <5+ , bytestring+ , deepseq+ , directory+ , filelock+ , mldsa+ , primitive >=0.7.2+ , process+ , tasty+ , tasty-hunit+ , tasty-quickcheck+ , text+ , zlib+ default-language: Haskell2010+ if flag(use_crypton)+ build-depends:+ crypton >=1.1.1+ , ram+ else+ build-depends:+ cryptonite >=0.26+ , memory++test-suite mldsa-test-full+ type: exitcode-stdio-1.0+ main-is: Tests.hs+ other-modules:+ Auxiliary+ Base+ Block+ BlockN+ Builder+ ByteArrayST+ Crypto+ Crypto.PubKey.ML_DSA+ Equality+ Fusion+ Internal+ Iterate+ Machine+ Marking+ Math+ Matrix+ ScrubbedBlock+ SecureBlock+ SecureBytes+ Vector+ KeyGen+ SigExt+ SigGen+ SigVer+ Util+ Vectors+ Paths_mldsa+ autogen-modules:+ Paths_mldsa+ 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 -Wno-unticked-promoted-constructors -fno-ignore-asserts -threaded -rtsopts -with-rtsopts=-N+ cpp-options: -DML_DSA_TESTING+ build-depends:+ aeson+ , base >=4.7 && <5+ , bytestring+ , deepseq+ , directory+ , filelock+ , primitive >=0.7.2+ , process+ , tasty+ , tasty-hunit+ , tasty-quickcheck+ , text+ , zlib+ default-language: Haskell2010+ if flag(use_crypton)+ build-depends:+ crypton >=1.1.1+ , ram+ else+ build-depends:+ cryptonite >=0.26+ , memory++benchmark mldsa-bench+ type: exitcode-stdio-1.0+ main-is: Bench.hs+ other-modules:+ Paths_mldsa+ autogen-modules:+ Paths_mldsa+ 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 -Wno-unticked-promoted-constructors -threaded -rtsopts -with-rtsopts=-N -with-rtsopts=-A48m+ build-depends:+ base >=4.7 && <5+ , criterion+ , deepseq+ , mldsa+ , primitive >=0.7.2+ default-language: Haskell2010+ if flag(use_crypton)+ build-depends:+ crypton >=1.1.1+ , ram+ else+ build-depends:+ cryptonite >=0.26+ , memory
+ src/Auxiliary.hs view
@@ -0,0 +1,959 @@+-- |+-- Module : Auxiliary+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- ML-DSA auxiliary functions+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UnboxedTuples #-}+module Auxiliary+ ( Zq, Rq, Tq+ , Norm(..), normFrom, countFrom, ntt, nttInv+ , simpleBitPack, simpleBitPack10, simpleBitUnpack10+ , bitPack, bitUnpack, bitPackSafe, bitUnpackSafe+ , hintBitPack, hintBitUnpack+ , sampleInBall, rejNttPoly, rejBoundedPoly+ , powerTwoRound+ , highBits, lowBits+ , Hints, makeHint, useHint+#ifdef ML_DSA_TESTING+ , Mq, toMontgomery, fromMontgomery, norm+ , bitRev8, fromZq, toZq, fromCoeffs, toCoeffs, fromBools, toBools+ , powerTwoRoundZq, decomposeZq+ , simpleBitUnpack+#endif+ ) where++import Crypto.Hash.Algorithms++import Data.ByteArray (ByteArrayAccess)+import qualified Data.ByteArray as B+import qualified Data.Memory.Endian as B++import Data.Primitive.Types (Prim(..))++import Control.DeepSeq (NFData(..))+import Control.Monad+import Control.Monad.ST++import Data.Bits+import Data.Int+import Data.Proxy+import Data.Word++#if MIN_VERSION_base(4,16,0)+import GHC.Exts as Exts (eqWord32#, int2Word#, wordToWord8#)+#else+import GHC.Exts as Exts (eqWord#, int2Word#)+#endif+import GHC.TypeNats+import GHC.Word (Word8(W8#), Word32(W32#))++import Foreign.Ptr (Ptr, plusPtr)+import Foreign.Storable (pokeByteOff)++import Base+import Block (blockIndex)+import BlockN (BlockN, MutableBlockN)+import Builder (Builder)+import Crypto (BlockDigest, Digest)+import Fusion+import Machine+import Marking (Classified, SecurityMarking(..))+import SecureBlock (SecureBlock)+import SecureBytes (SecureBytes)+import Vector (Vector)+import qualified BlockN+import qualified Builder+import qualified ByteArrayST as ST+import qualified Crypto+import qualified Vector+import Math++type N = 256++n :: Int+n = 256++q :: Integer+q = 8380417++q32 :: Word32+q32 = fromInteger q++q64 :: Word64+q64 = fromInteger q++bitRev8 :: Word8 -> Word8+bitRev8 = fromIntegral . rev8 . fromIntegral+ where+ rev1, rev2, rev4, rev8 :: Word -> Word+ rev1 b = b .&. 1+ rev2 b = rev1 (b `unsafeShiftR` 1) .|. rev1 b `unsafeShiftL` 1+ rev4 b = rev2 (b `unsafeShiftR` 2) .|. rev2 b `unsafeShiftL` 2+ rev8 b = rev4 (b `unsafeShiftR` 4) .|. rev4 b `unsafeShiftL` 4++unsafeShiftIR :: Word32 -> Int -> Word32+unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int32) `unsafeShiftR` s)+{-# INLINE unsafeShiftIR #-}++select32 :: Word32 -> Word32 -> Word32 -> Word32+select32 mask yes no = (mask .&. yes) .|. (complement mask .&. no)+{-# INLINE select32 #-}++max32 :: Word32 -> Word32 -> Word32+max32 a b = let cond = (b - a) `unsafeShiftIR` 31 in select32 cond a b+{-# INLINE max32 #-}++min32 :: Word32 -> Word32 -> Word32+min32 a b = let cond = (a - b) `unsafeShiftIR` 31 in select32 cond a b+{-# INLINE min32 #-}++-- Reduction 𝑥 mod 5 for 0 ≤ 𝑥 < 15+mod5 :: Word32 -> Word32+mod5 x = let quotient = (x * 205) `unsafeShiftR` 10 in x - 5 * quotient++-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞+reduceSimple :: Word32 -> Word32+reduceSimple x = select32 mask x subtracted+ where+ subtracted = x - q32+ mask = subtracted `unsafeShiftIR` 31+{-# INLINE reduceSimple #-}++#ifdef ML_DSA_TESTING+-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 1025𝑞+reduce :: Word64 -> Word32+reduce x = reduceSimple $ fromIntegral (x - t * q64)+ where t = x `unsafeShiftR` 23+{-# INLINE reduce #-}+#endif++-- Computes 𝑎 . 2^-32 mod 𝑞+reduceMontgomery :: Word64 -> Word32+reduceMontgomery a = reduceSimple (fromIntegral r)+ where+ t = (fromIntegral a :: Word32) * 4236238847 -- inverse of -𝑞 modulo 2^32+ r = (a + fromIntegral t * q64) `unsafeShiftR` 32+{-# INLINE reduceMontgomery #-}++newtype Zq = Zq Word32+#ifdef ML_DSA_TESTING+ deriving (Eq, Show)+#else+ deriving Eq+#endif++instance Prim Zq where+ sizeOf# (Zq a) = sizeOf# a+ {-# INLINE sizeOf# #-}+ alignment# (Zq a) = alignment# a+ {-# INLINE alignment# #-}+#if MIN_VERSION_primitive(0,9,0)+ sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word32)+ {-# INLINE sizeOfType# #-}+ alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word32)+ {-# INLINE alignmentOfType# #-}+#endif+ indexByteArray# ba i = Zq (indexByteArray# ba i)+ {-# INLINE indexByteArray# #-}+ readByteArray# mba i s =+ case readByteArray# mba i s of+ (# s', a #) -> (# s', Zq a #)+ {-# INLINE readByteArray# #-}+ writeByteArray# mba i (Zq a) = writeByteArray# mba i a+ {-# INLINE writeByteArray# #-}+ setByteArray# mba i len (Zq a) = setByteArray# mba i len a+ {-# INLINE setByteArray# #-}+ indexOffAddr# addr i = Zq (indexOffAddr# addr i)+ {-# INLINE indexOffAddr# #-}+ readOffAddr# addr i s =+ case readOffAddr# addr i s of+ (# s', a #) -> (# s', Zq a #)+ {-# INLINE readOffAddr# #-}+ writeOffAddr# addr i (Zq a) = writeOffAddr# addr i a+ {-# INLINE writeOffAddr# #-}+ setOffAddr# addr i len (Zq a) = setOffAddr# addr i len a+ {-# INLINE setOffAddr# #-}++instance PrimSized Zq where+ type PrimSize Zq = 4++instance Add Zq where+ zero = Zq 0+ Zq a .+ Zq b = Zq $ reduceSimple (a + b)+ Zq a .- Zq b = Zq $ reduceSimple (a + q32 - b)+ neg (Zq a) = Zq $ reduceSimple (q32 - a)++#ifdef ML_DSA_TESTING+instance Mul Zq where+ one = Zq 1+ a .* b = fromMontgomery (toMontgomery a .* toMontgomery b)++toZq :: Word32 -> Zq+toZq = Zq . reduce . fromIntegral+#endif++fromZq :: Zq -> Word32+fromZq (Zq a) = a++newtype Mq = Mq Word32+#ifdef ML_DSA_TESTING+ deriving (Eq, Show)+#else+ deriving Eq+#endif++instance Prim Mq where+ sizeOf# (Mq a) = sizeOf# a+ {-# INLINE sizeOf# #-}+ alignment# (Mq a) = alignment# a+ {-# INLINE alignment# #-}+#if MIN_VERSION_primitive(0,9,0)+ sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word32)+ {-# INLINE sizeOfType# #-}+ alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word32)+ {-# INLINE alignmentOfType# #-}+#endif+ indexByteArray# ba i = Mq (indexByteArray# ba i)+ {-# INLINE indexByteArray# #-}+ readByteArray# mba i s =+ case readByteArray# mba i s of+ (# s', a #) -> (# s', Mq a #)+ {-# INLINE readByteArray# #-}+ writeByteArray# mba i (Mq a) = writeByteArray# mba i a+ {-# INLINE writeByteArray# #-}+ setByteArray# mba i len (Mq a) = setByteArray# mba i len a+ {-# INLINE setByteArray# #-}+ indexOffAddr# addr i = Mq (indexOffAddr# addr i)+ {-# INLINE indexOffAddr# #-}+ readOffAddr# addr i s =+ case readOffAddr# addr i s of+ (# s', a #) -> (# s', Mq a #)+ {-# INLINE readOffAddr# #-}+ writeOffAddr# addr i (Mq a) = writeOffAddr# addr i a+ {-# INLINE writeOffAddr# #-}+ setOffAddr# addr i len (Mq a) = setOffAddr# addr i len a+ {-# INLINE setOffAddr# #-}++instance PrimSized Mq where+ type PrimSize Mq = 4++instance Add Mq where+ zero = Mq 0+ Mq a .+ Mq b = Mq $ reduceSimple (a + b)+ Mq a .- Mq b = Mq $ reduceSimple (a + q32 - b)+ neg (Mq a) = Mq $ reduceSimple (q32 - a)++instance Mul Mq where+ one = Mq 4193792 -- toMongomery (Zq 1)+ Mq a .* Mq b = Mq $ reduceMontgomery (fromIntegral a * fromIntegral b)++instance MulAdd Mq where+ mulAdd a b c = a .* b .+ c++#ifdef ML_DSA_TESTING+instance BiMul Mq Mq where+ (..*) = (.*)++instance BiMulAdd Mq Mq where+ biMulAdd = mulAdd+#endif++toMontgomery :: Zq -> Mq+toMontgomery (Zq a) = Mq $ reduceMontgomery $ fromIntegral a * 2365951 -- R^2 mod 𝑞++fromMontgomery :: Mq -> Zq+fromMontgomery (Mq m) = Zq $ reduceMontgomery $ fromIntegral m++newtype Rq marking = Rq (BlockN marking N Zq)+#ifdef ML_DSA_TESTING+ deriving (Eq, Show, NFData)+#else+ deriving NFData+#endif++instance Classified marking => Add (Rq marking) where+ zero = Rq zero+ Rq a .+ Rq b = Rq (a .+ b)+ {-# INLINE (.+) #-}+ Rq a .- Rq b = Rq (a .- b)+ {-# INLINE (.-) #-}+ neg (Rq a) = Rq (neg a)+ {-# INLINE neg #-}++instance Crypto.ConstEqW (Rq Sec) where+ constEqW (Rq a) (Rq b) = Crypto.constEqW+ (BlockN.unsafeCast a :: SecureBlock Sec Word)+ (BlockN.unsafeCast b :: SecureBlock Sec Word)++instance Crypto.ConstEqW (Rq Pub) where+ constEqW (Rq a) (Rq b) = Crypto.constEqW+ (BlockN.unsafeCast a :: SecureBlock Pub Word)+ (BlockN.unsafeCast b :: SecureBlock Pub Word)++#ifdef ML_DSA_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 Mq)+#ifdef ML_DSA_TESTING+ deriving (Eq, Show, NFData)+#else+ deriving NFData+#endif++instance Classified marking => Add (Tq marking) where+ zero = Tq zero+ Tq a .+ Tq b = Tq (a .+ b)+ {-# INLINE (.+) #-}+ Tq a .- Tq b = Tq (a .- b)+ {-# INLINE (.-) #-}+ neg (Tq a) = Tq (neg a)+ {-# INLINE neg #-}++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)++instance BiMul (Tq marking) (Tq Sec) where+ Tq a ..* Tq b = Tq $ BlockN.zipWithEqPrimSizeR (.*) a b+ {-# INLINE (..*) #-}++instance BiMulAdd (Tq marking) (Tq Sec) where+ biMulFold = multiplyNTTsFold+ {-# INLINE biMulFold #-}++#ifdef ML_DSA_TESTING+instance Mul (Tq Sec) where+ one = Tq $ BlockN.replicate one+ (.*) = (..*)++instance MulAdd (Tq Sec) where+ mulAdd = biMulAdd+#endif++newtype Norm = Norm { getNorm :: Word32 }+#ifdef ML_DSA_TESTING+ deriving (Eq, Ord, Show, Num)+#endif++instance Semigroup Norm where+ Norm a <> Norm b = Norm (max32 a b)++instance Monoid Norm where+ mempty = Norm 0++#ifdef ML_DSA_TESTING+norm :: Rq marking -> Norm+norm = normFrom mempty+#endif++normFrom :: Norm -> Rq marking -> Norm+normFrom (Norm m0) (Rq w) = Norm $ BlockN.foldl' f m0 w+ where+ absZq (Zq z) = min32 z (q32 - z)+ f m = max32 m . absZq+{-# INLINE normFrom #-}++countFrom :: Word -> Hints -> Word+countFrom b (Hints w) = let f acc (H h) = acc + fromIntegral h in BlockN.foldl' f b w+{-# INLINE countFrom #-}++-- Computes the NTT representation of the given polynomial+ntt :: Classified marking => Rq marking -> Tq marking+ntt (Rq a) = Tq $ BlockN.runThaw b mutNtt+ where b = BlockN.mapEqPrimSize toMontgomery a+{-# INLINE ntt #-}++mutNtt :: MutableBlockN marking N Mq s -> ST s ()+mutNtt !b = outer 1 128+ where+ outer !m len = when (len >= 1) $ inner m len 0++ inner !m !len start+ | start < 256 = do+ let z = BlockN.index zetaPowBitRev m -- 1753 ^ bitRev8 m+ loop z (start + len) len start+ inner (m + 1) len (start + offsetShiftL 1 len)+ | otherwise = outer m (offsetShiftR 1 len)++ loop !z end len j =+ when (j < end) $ do+ t <- (z .*) <$> BlockN.read b (j + len)+ x <- BlockN.read b j+ BlockN.write b (j + len) (x .- t)+ BlockN.write b j (x .+ t)+ loop z end len (j + 1)+{-# NOINLINE mutNtt #-}++-- Computes the polynomial that corresponds to the given NTT representation+nttInv :: Classified marking => Tq marking -> Rq marking+nttInv (Tq a) = Rq $ BlockN.mapEqPrimSize fromMontgomery $ BlockN.runThaw a mutNttInv+{-# INLINE nttInv #-}++mutNttInv :: MutableBlockN marking N Mq s -> ST s ()+mutNttInv !b = do+ outer 255 1+ BlockN.iterModify (\x -> x .* Mq 16382) b -- toMontgomery (Zq 8347681)+ where+ outer !m len = when (len < 256) $ inner m len 0++ inner !m !len start+ | start < 256 = do+ let z = BlockN.index zetaPowBitRev m -- 1753 ^ bitRev8 m+ loop z (start + len) len start+ inner (m - 1) len (start + offsetShiftL 1 len)+ | otherwise = outer m (offsetShiftL 1 len)++ loop !z 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) (z .* (x .- t))+ loop z end len (j + 1)+{-# NOINLINE mutNttInv #-}++multiplyNTTsFold :: Foldable t => Tq Sec -> t (Tq marking, Tq Sec) -> Tq Sec+multiplyNTTsFold (Tq c) = Tq . BlockN.runFold c (uncurry multiplyNTTsAdd)+{-# INLINE multiplyNTTsFold #-}++-- Multiply then add a third term+multiplyNTTsAdd :: Tq marking -> Tq Sec -> MutableBlockN Sec N Mq s -> ST s ()+multiplyNTTsAdd (Tq !f) (Tq !g) bb = loop bb 0+ where+ loop :: MutableBlockN Sec N Mq s -> Offset Mq -> ST s ()+ loop !b i = when (i < 256) $ do+ c <- BlockN.read b i+ BlockN.write b i $ mulAdd (BlockN.index f i) (BlockN.index g i) c+ loop b (i + 1)++-- Values of 1753 ^ BitRev8(𝑖) mod 𝑞 for 𝑖 ∈ {0, … , 255}+zetaPowBitRev :: BlockN Pub 256 Mq+zetaPowBitRev = BlockN.runNew (Proxy :: Proxy Pub) $ \out ->+ foldM_ (loop out) one offsets+ where+ offsets = Prelude.map (fromIntegral . bitRev8) [0 .. 255]+ loop b acc i = BlockN.write b i acc >> return (Mq 2091667 .* acc)+ -- toMontgomery (Zq 1753)++newtype H = H Word8+#ifdef ML_DSA_TESTING+ deriving (Eq, Show)+#else+ deriving Eq+#endif++instance Prim H where+ sizeOf# (H a) = sizeOf# a+ {-# INLINE sizeOf# #-}+ alignment# (H a) = alignment# a+ {-# INLINE alignment# #-}+#if MIN_VERSION_primitive(0,9,0)+ sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word8)+ {-# INLINE sizeOfType# #-}+ alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word8)+ {-# INLINE alignmentOfType# #-}+#endif+ indexByteArray# ba i = H (indexByteArray# ba i)+ {-# INLINE indexByteArray# #-}+ readByteArray# mba i s =+ case readByteArray# mba i s of+ (# s', a #) -> (# s', H a #)+ {-# INLINE readByteArray# #-}+ writeByteArray# mba i (H a) = writeByteArray# mba i a+ {-# INLINE writeByteArray# #-}+ setByteArray# mba i len (H a) = setByteArray# mba i len a+ {-# INLINE setByteArray# #-}+ indexOffAddr# addr i = H (indexOffAddr# addr i)+ {-# INLINE indexOffAddr# #-}+ readOffAddr# addr i s =+ case readOffAddr# addr i s of+ (# s', a #) -> (# s', H a #)+ {-# INLINE readOffAddr# #-}+ writeOffAddr# addr i (H a) = writeOffAddr# addr i a+ {-# INLINE writeOffAddr# #-}+ setOffAddr# addr i len (H a) = setOffAddr# addr i len a+ {-# INLINE setOffAddr# #-}++instance PrimSized H where+ type PrimSize H = 1++-- boolean-less comparison, to avoid code duplication with pattern matching+notEqualH :: Word32 -> Word32 -> H+notEqualH (W32# x#) (W32# y#) = H $+#if MIN_VERSION_base(4,16,0)+ 1 - W8# (Exts.wordToWord8# (Exts.int2Word# (Exts.eqWord32# x# y#)))+#else+ 1 - W8# (Exts.int2Word# (Exts.eqWord# x# y#))+#endif++newtype Hints = Hints (BlockN Sec N H)+#ifdef ML_DSA_TESTING+ deriving (Eq, Show, NFData)+#else+ deriving NFData+#endif++instance Crypto.ConstEqW Hints where+ constEqW (Hints a) (Hints b) = Crypto.constEqW+ (BlockN.unsafeCast a :: SecureBlock Sec Word)+ (BlockN.unsafeCast b :: SecureBlock Sec Word)++#ifdef ML_DSA_TESTING+fromBools :: [Bool] -> Maybe Hints+fromBools = fmap Hints . BlockN.fromList . map (H . fromIntegral . fromEnum)++toBools :: Hints -> [Bool]+toBools (Hints a) = map (\(H h) -> toEnum (fromIntegral h)) (BlockN.toList a)+#endif++peekWordPos :: Ptr WordLE -> BitPos -> ST s WordM+peekWordPos a bp = fromLE <$> ST.peekElemOff a (wordOff bp)++pokeWordPos :: Ptr WordLE -> BitPos -> WordM -> ST s ()+pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . toLE++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 -> WordM+getMask howMany+ | howMany >= wordBits = maxBound+ | otherwise = (1 `unsafeShiftL` howMany) - 1+ -- branch useful only when processing one byte at a time due to+ -- architecture not supporting unaligned memory access++-- Encodes an array of 𝑑-bit integers for 1 ≤ 𝑑 ≤ 22+simpleBitPack :: Int -> BlockN marking N Word32 -> Builder marking+simpleBitPack d w = Builder.create (32 * d) (runSimpleBitPack d w)+{-# INLINE simpleBitPack #-}++runSimpleBitPack :: Int -> BlockN marking N Word32 -> Ptr WordLE -> ST s ()+runSimpleBitPack !d !w dst = loop dst 0 zeroPos 0 (get 0) d+ where+ get = BlockN.index w . Offset+ {-# INLINE get #-}++ loop !b !pos !bp !o !a j+ | j == 0, pos' == n = return ()+ | j == 0 = loop b pos' bp o (get pos') d+ | bitPos bp + howMany < wordBits = loop b pos bp' o' a' j'+ | otherwise = pokeWordPos b bp o' >> loop b pos bp' 0 a' j'+ where+ pos' = pos + 1+ (howMany, bp') = nextPos j bp+ x = fromIntegral a .&. getMask howMany+ o' = o .|. (x `unsafeShiftL` bitPos bp)+ a' = a `unsafeShiftR` howMany+ j' = j - howMany++-- simpleBitPack with 𝑑=10 after conversion from a polynomial and division by 2^13+simpleBitPack10 :: Rq Pub -> Builder Pub+simpleBitPack10 (Rq w) = simpleBitPack 10 (BlockN.mapEqPrimSize f w)+ where f (Zq z) = z `unsafeShiftR` 13+{-# INLINE simpleBitPack10 #-}++-- Decodes an array of 𝑑-bit integers for 1 ≤ 𝑑 ≤ 22+simpleBitUnpack :: (Classified marking, ByteArrayAccess ba) => Int -> ba -> BlockN marking N Word32+simpleBitUnpack d b = BlockN.runNew Proxy $ mutSimpleBitUnpack d b+{-# INLINE simpleBitUnpack #-}++mutSimpleBitUnpack :: ByteArrayAccess ba => Int -> ba -> MutableBlockN marking N Word32 s -> ST s ()+mutSimpleBitUnpack !d b !w = ST.withByteArray b $ \p -> outer p zeroPos 0+ where+ outer !p !bp i = when (i < Offset n) $ inner p i bp 0 0++ inner !p !i !bp !v j+ | j == d = BlockN.write w i v >> outer p bp (i + 1)+ | otherwise = do+ let (howMany, bp') = nextPos (d - j) bp+ y <- get p bp howMany+ let v' = v .|. (fromIntegral y `unsafeShiftL` j)+ j' = j + howMany+ inner p i bp' v' j'++ get :: Ptr WordLE -> BitPos -> Int -> ST s WordM+ get p bp howMany = do+ x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp+ return (x .&. getMask howMany)++-- simpleBitUnpack with 𝑑=10 and conversion to a polynomial with multiplication by 2^13+simpleBitUnpack10 :: ByteArrayAccess ba => ba -> Rq Pub+simpleBitUnpack10 = Rq . BlockN.mapEqPrimSize f . simpleBitUnpack 10+ where f z = Zq (z `unsafeShiftL` 13)+{-# INLINE simpleBitUnpack10 #-}++-- bitPack when upper bound is a power of two+bitPackSafe :: Classified marking => Int -> Rq marking -> Builder marking+bitPackSafe d = simpleBitPack d . unmirror (1 `unsafeShiftL` (d - 1))+{-# INLINE bitPackSafe #-}++-- bitUnpack when upper bound is a power of two+bitUnpackSafe :: (Classified marking, ByteArrayAccess ba) => Int -> ba -> Rq marking+bitUnpackSafe d b = mirror (1 `unsafeShiftL` (d - 1)) (simpleBitUnpack d b)+{-# INLINE bitUnpackSafe #-}++-- Encodes a polynomial 𝑤 into a byte string+bitPack :: Word32 -> Int -> Rq Sec -> Builder Sec+bitPack !m d = simpleBitPack d . unmirror m++-- Reverses the procedure bitPack+bitUnpack :: ByteArrayAccess ba => Word32 -> Int -> ba -> Maybe (Rq Sec)+bitUnpack !m d b+ | valid unpacked = Just $ mirror m unpacked+ | otherwise = Nothing+ where+ unpacked = simpleBitUnpack d b+ valid = Crypto.toBool . BlockN.foldl' f Crypto.trueW+ f acc z = acc `Crypto.andW` inRange z+ inRange z = Crypto.ltW (fromIntegral z) (fromIntegral (2 * m + 1))++mirror :: Classified marking => Word32 -> BlockN marking N Word32 -> Rq marking+mirror m w = Rq $ BlockN.mapEqPrimSize (\z -> Zq m .- Zq z) w+{-# INLINE mirror #-}++unmirror :: Classified marking => Word32 -> Rq marking -> BlockN marking N Word32+unmirror m (Rq w) = BlockN.mapEqPrimSize (\z -> fromZq (Zq m .- z)) w+{-# INLINE unmirror #-}++-- Encodes a polynomial vector 𝐡 with binary coefficients into a byte string+hintBitPack :: forall k. KnownNat k => Int -> Vector k Hints -> Builder Sec+hintBitPack !omega !h = Builder.create (omega + k) $ outer 0 0+ where+ k = fromIntegral $ natVal (Proxy :: Proxy k)++ outer :: Int -> Offset Word8 -> Ptr Word8 -> ST s ()+ outer i !idx !p+ | i == k = end p idx+ | otherwise = inner p i (Vector.index h $ Offset i) idx 0++ inner :: Ptr Word8 -> Int -> Hints -> Offset Word8 -> Int -> ST s ()+ inner !p !i e@(Hints !w) !idx j+ | j == n = mark p i idx >> outer (i + 1) idx p+ | BlockN.index w (Offset j) == H 0 = inner p i e idx (j + 1)+ | otherwise = set p idx j >> inner p i e (idx + 1) (j + 1)++ set :: Ptr Word8 -> Offset Word8 -> Int -> ST s ()+ set p (Offset idx) j = ST.pokeElemOff p idx (fromIntegral j)++ mark :: Ptr Word8 -> Int -> Offset Word8 -> ST s ()+ mark p i (Offset idx) = ST.pokeElemOff p (omega + i) (fromIntegral idx)++ end :: Ptr Word8 -> Offset Word8 -> ST s ()+ end !p (Offset idx) =+ when (idx < omega) $ ST.fillBytes (p `plusPtr` idx) 0 (omega - idx)+{-# INLINE hintBitPack #-}++-- Reverses the procedure hintBitPack+hintBitUnpack :: (KnownNat k, ByteArrayAccess ba) => Int -> ba -> Maybe (Vector k Hints)+hintBitUnpack !omega !y = Vector.createMaybeAccum outer cont 0+ where+ outer :: Int -> Offset Hints -> Maybe (Int, Hints)+ outer idx (Offset i) = runST $ ST.withByteArray y $ \p -> do+ m <- fromIntegral <$> ST.peekElemOff p (omega + i)+ if m < idx || m > omega+ then return Nothing+ else new >>= inner p m idx++ inner :: Ptr Word8 -> Int -> Int -> MutableBlockN Sec N H s -> ST s (Maybe (Int, Hints))+ inner !p !m !first !mb = go first+ where+ go idx+ | idx >= m = unsafeFreezeF mb >>= \b ->+ return $ Just (idx, Hints b)+ | idx > first = do+ val <- ST.peekElemOff p idx+ val' <- ST.peekElemOff p (idx - 1)+ if val' >= val+ then return Nothing+ else set mb val >> go (idx + 1)+ | otherwise = ST.peekElemOff p idx >>= \val ->+ set mb val >> go (idx + 1)++ new :: ST s (MutableBlockN Sec N H s)+ new = newF >>= \mb -> BlockN.erase mb >> return mb++ set :: MutableBlockN Sec N H s -> Word8 -> ST s ()+ set mb val = BlockN.write mb (fromIntegral val) (H 1)++ cont :: forall r. Int -> r -> Maybe r+ cont idx h = runST $ ST.withByteArray y $ \p -> go p idx+ where+ go :: Ptr Word8 -> Int -> ST s (Maybe r)+ go !p i+ | i >= omega = return (Just h)+ | otherwise = do+ val <- ST.peekElemOff p i+ case val of+ 0 -> go p (i + 1)+ _ -> return Nothing+{-# INLINE hintBitUnpack #-}++-- Samples a polynomial 𝑐 ∈ 𝑅 with coefficients from {-1, 0, 1} and+-- Hamming weight 𝜏 ≤ 64+sampleInBall :: Int -> SecureBytes Sec -> Rq Sec+sampleInBall tau rho = Rq $+ BlockN.runNew (Proxy :: Proxy Sec) $ \c -> runXof c 136+ where+ runXof :: MutableBlockN Sec N Zq s -> Int -> ST s ()+ runXof !c !xofLen = case someNatVal (fromIntegral (8 * xofLen)) of+ SomeNat proxy -> do+ let bytes = Crypto.unDigest (doHash proxy)+ BlockN.erase c+ loop c xofLen bytes 8 (256 - tau)++ loop :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Int -> ST s ()+ loop !c !xofLen !bytes !pos i+ | i == n = return ()+ | pos >= xofLen = runXof c (xofLen + 136)+ | j > Offset i = loop c xofLen bytes (pos + 1) i+ | otherwise = do+ BlockN.read c j >>= BlockN.write c (Offset i)+ let idx = i + tau - 256+ iByte = idx `unsafeShiftR` 3+ iBit = idx .&. 7+ b = B.index bytes iByte `unsafeShiftR` iBit+ BlockN.write c j (Zq $ if even b then 1 else q32 - 1)+ loop c xofLen bytes (pos + 1) (i + 1)+ where+ j = fromIntegral $ B.index bytes pos++ doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen)+ doHash _ = Crypto.hash rho++-- Samples a polynomial ∈ 𝑇𝑞+rejNttPoly :: SecureBytes Pub -> Word8 -> Word8 -> Tq Pub+rejNttPoly rho !s !r = Tq $+ BlockN.runNew (Proxy :: Proxy Pub) $ \b -> runXof b (280 * 3) 0 0+ where+ runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of+ SomeNat proxy -> do+ 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 s0 = fromIntegral $ blockIndex bytes pos+ s1 = fromIntegral $ blockIndex bytes (pos + 1)+ s2 = fromIntegral $ blockIndex bytes (pos + 2)+ z = s0 + (s1 `unsafeShiftL` 8)+ + ((s2 .&. 0x7F) `unsafeShiftL` 16)+ poke b j z >>= loop b xofLen bytes (pos + 3)++ poke b j z+ | z < q32 = BlockN.write b j (toMontgomery $ Zq z) >> 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 rho d+ pokeByteOff d len s+ pokeByteOff d (len + 1) r+ len = B.length rho++-- Samples an element 𝑎 ∈ 𝑅 with coefficients in [-𝜂, 𝜂] computed via rejection+-- sampling from 𝜌+rejBoundedPoly :: SecureBytes Sec -> Int -> Word16 -> Rq Sec+rejBoundedPoly rho !eta !r = Rq $+ BlockN.runNew (Proxy :: Proxy Sec) $ \b -> runXof b 408 0 0+ where+ runXof :: MutableBlockN Sec N Zq s -> Int -> Int -> Offset Zq -> ST s ()+ runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of+ SomeNat proxy -> case eta of+ 2 -> loop2 b xofLen bytes pos j+ 4 -> loop4 b xofLen bytes pos j+ _ -> error "rejBoundedPoly: unexpected eta value"+ where bytes = Crypto.unDigest (doHash proxy)++ loop2 :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Offset Zq -> ST s ()+ loop2 !b !xofLen !bytes !pos j+ | j == 256 = return ()+ | pos >= xofLen = runXof b (xofLen + 136) pos j+ | otherwise = do+ let z = fromIntegral $ B.index bytes pos+ z0 = z .&. 0xF+ z1 = z `unsafeShiftR` 4+ j2 <- poke2 b j z0+ when (j2 < 256) $ poke2 b j2 z1 >>= loop2 b xofLen bytes (pos + 1)++ loop4 :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Offset Zq -> ST s ()+ loop4 !b !xofLen !bytes !pos j+ | j == 256 = return ()+ | pos >= xofLen = runXof b (xofLen + 136) pos j+ | otherwise = do+ let z = fromIntegral $ B.index bytes pos+ z0 = z .&. 0xF+ z1 = z `unsafeShiftR` 4+ j2 <- poke4 b j z0+ when (j2 < 256) $ poke4 b j2 z1 >>= loop4 b xofLen bytes (pos + 1)++ -- rejection sampling from {-2, … , 2}+ poke2 b j z+ | z < 15 = BlockN.write b j (Zq 2 .- Zq (mod5 z)) >> return (j + 1)+ | otherwise = return j++ -- rejection sampling from {-4, … , 4}+ poke4 b j z+ | z < 9 = BlockN.write b j (Zq 4 .- Zq z) >> return (j + 1)+ | otherwise = return j++ doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen)+ doHash _ = Crypto.hash input++ input :: SecureBytes Sec+ !input = B.unsafeCreate (len + 2) $ \d -> do+ B.copyByteArrayToPtr rho d+ pokeByteOff d len (B.toLE r)+ len = B.length rho++powerTwoRound :: Rq Sec -> (Rq Pub, Rq Pub)+powerTwoRound (Rq !a) = runST $ do+ mb <- newF+ mc <- newF+ loop mb mc 0+ b <- unsafeFreezeF mb+ c <- unsafeFreezeF mc+ return (Rq b, Rq c)+ where+ loop :: MutableBlockN Pub N Zq s -> MutableBlockN Pub N Zq s -> Int -> ST s ()+ loop !mb !mc i = when (i < n) $ do+ let r = BlockN.index a (Offset i)+ (r1, r0) = powerTwoRoundZq r+ BlockN.write mb (Offset i) r1+ BlockN.write mc (Offset i) r0+ loop mb mc (i + 1)++-- Decomposes 𝑟 into (𝑟1, 𝑟0) such that 𝑟 ≡ 𝑟1.2^𝑑 + 𝑟0 mod 𝑞 (for 𝑑=13)+powerTwoRoundZq :: Zq -> (Zq, Zq)+powerTwoRoundZq (Zq r) = (Zq r1, Zq r0)+ where+ r0' = r .&. 0x1fff+ r1' = r .&. 0xffffe000 -- keep multiplied by 2^13++ -- if r0 > 4096, increment r1 and subtract 8192 from r0+ cond = (4096 - r0') `unsafeShiftIR` 31+ r0 = select32 cond (q32 - 0x2000 + r0') r0'+ r1 = reduceSimple (r1' + (cond .&. 0x2000))++-- Decomposes 𝑟 into (𝑟1, 𝑟0) such that 𝑟 ≡ 𝑟1.(2𝛾2) + 𝑟0 mod 𝑞+decomposeZq :: Word32 -> Zq -> (Word32, Zq)+decomposeZq gamma2 r@(Zq rr) = (r1, Zq r0)+ where+ r1 = highZq gamma2 r+ r0' = rr - r1 * (2 * gamma2)+ !r0 = r0' + ((r0' `unsafeShiftIR` 31) .&. q32)++highZq :: Word32 -> Zq -> Word32+highZq gamma2 (Zq r) = case gamma2 of+ 95232 -> high190464 r+ _ -> high523776 r++high190464 :: Word32 -> Word32+high190464 r = r1 `xor` (((43 - r1) `unsafeShiftIR` 31) .&. r1)+ where+ r' = (r + 127) `unsafeShiftR` 7+ r1 = (r' * 11275 + (1 `unsafeShiftL` 23)) `unsafeShiftR` 24;++high523776 :: Word32 -> Word32+high523776 r = r1 .&. 15+ where+ r' = (r + 127) `unsafeShiftR` 7+ r1 = (r' * 1025 + (1 `unsafeShiftL` 21)) `unsafeShiftR` 22;++-- Returns 𝑟1 from the output of decomposeZq 𝑟+highBits :: Word32 -> Rq Sec -> BlockN Sec N Word32+highBits gamma2 (Rq w) = BlockN.seq gamma2 $+ BlockN.mapEqPrimSize (highZq gamma2) w+{-# INLINE highBits #-}++-- Returns 𝑟0 from the output of decomposeZq 𝑟+lowBits :: Word32 -> Rq Sec -> Rq Sec+lowBits gamma2 (Rq w) = Rq $ BlockN.seq gamma2 $+ BlockN.mapEqPrimSize (snd . decomposeZq gamma2) w+{-# INLINE lowBits #-}++-- Computes hint bit indicating whether adding 𝑧 to 𝑟 alters the high bits of 𝑟+makeHint :: Word32 -> Rq Sec -> Rq Sec -> Hints+makeHint gamma2 (Rq z) (Rq r) = Hints $ BlockN.seq gamma2 $+ BlockN.zipWith (makeHintH gamma2) z r+{-# INLINE makeHint #-}++-- Returns the high bits of 𝑟 adjusted according to hint ℎ+useHint :: Word32 -> Hints -> Rq Sec -> BlockN Sec N Word32+useHint gamma2 (Hints h) (Rq r) = BlockN.seq gamma2 $+ BlockN.zipWithEqPrimSizeR (useHintH gamma2) h r+{-# INLINE useHint #-}++makeHintH :: Word32 -> Zq -> Zq -> H+makeHintH gamma2 z r = notEqualH r1 v1+ where+ r1 = highZq gamma2 r+ v1 = highZq gamma2 (r .+ z)++useHintH :: Word32 -> H -> Zq -> Word32+useHintH !gamma2 h r+ | h == H 0 = r1+ | positive r0 = decr gamma2 r1+ | otherwise = incr gamma2 r1+ where+ (r1, Zq r0) = decomposeZq gamma2 r++ -- useHintH already has a 6-way code path with just 3 possibilities for+ -- output times the match on gamma2, so for the other tests we do not use+ -- boolean expressions, otherwise code size would grow even more++ -- z > 0 and z <= (𝑞-1)/2+ positive z =+ let x = negate z+ m = (q32 - 1) `div` 2+ w = x .&. xor (m + x) 0x80000000+ in w < 0x80000000++ incr43, decr43 :: Word32 -> Word32+ decr43 x = 43 - incr43 (43 - x)+ incr43 x = (x + 1) .&. complement wrapAroundMask+ where wrapAroundMask = (42 - x) `unsafeShiftIR` 31++ incr 95232 z = incr43 z -- if z == 43 then 0 else z + 1+ incr _ z = (z + 1) .&. 15++ decr 95232 z = decr43 z -- if z == 0 then 43 else z - 1+ decr _ z = (z - 1) .&. 15+{-# INLINE useHintH #-}
+ src/Base.hs view
@@ -0,0 +1,62 @@+-- |+-- Module : Base+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Utilities similar to ones provided in basement but adapted for primitive+--+{-# LANGUAGE CPP #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE TypeFamilies #-}+module Base+ ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimSized(..)+ , PrimState, natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#)+#ifdef ML_DSA_TESTING+ , checkBounds+#endif+ ) where++import Control.Monad.Primitive++import Data.Primitive.Types (Prim)++import Data.Bits+import Data.Word++import GHC.TypeNats++type PrimType = Prim++class PrimType a => PrimSized a where+ type PrimSize a :: Nat++instance PrimSized Word32 where+ type PrimSize Word32 = 4++newtype CountOf ty = CountOf Int+ deriving (Show, Eq, Ord, Num)++newtype Offset ty = Offset Int+ deriving (Show, Eq, Ord, Num)++offsetShiftL :: Int -> Offset ty -> Offset ty2+offsetShiftL n (Offset o) = Offset (o `unsafeShiftL` n)++offsetShiftR :: Int -> Offset ty -> Offset ty2+offsetShiftR n (Offset o) = Offset (o `unsafeShiftR` n)++unsafePrimFromIO :: PrimMonad m => IO a -> m a+unsafePrimFromIO = unsafeIOToPrim++(.==#) :: Offset ty -> CountOf ty -> Bool+(.==#) (Offset ofs) (CountOf sz) = ofs == sz+{-# INLINE (.==#) #-}++#ifdef ML_DSA_TESTING+checkBounds :: CountOf ty -> Offset ty -> a -> a+checkBounds (CountOf sz) (Offset ofs) a+ | ofs >= 0 && ofs < sz = a+ | otherwise = error ("offset not in valid range: " ++ show ofs)+#endif
+ src/Block.hs view
@@ -0,0 +1,127 @@+-- |+-- Module : Block+-- License : BSD-3-Clause+-- Copyright : (c) 2025 Olivier Chéron+--+-- An array of primitive (unlifted) elements. This module exposes the+-- t'PrimArray' implementation from primitive but through an API similar to+-- basement @Block@.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}+module Block+ ( Block, MutableBlock, blockIndex, blockRead, blockWrite, erase+ , foldZipWith, Block.length, mutableContents, getMutableLength+ , Block.new, Block.newPinned, Block.thaw, thawPinned+ , Block.unsafeCast, unsafeCastMut, Block.unsafeFreeze+#ifdef ML_DSA_TESTING+ , Block.toList+#endif+ ) where++import Control.Monad.Primitive++import Data.Primitive.ByteArray+import Data.Primitive.PrimArray+import Data.Primitive.Types (sizeOf#)++import Control.Exception (assert)++import Foreign.Ptr (Ptr)++import Base++import GHC.Base (Int(I#), setByteArray#)++type Block = PrimArray+type MutableBlock ty s = MutablePrimArray s ty++blockIndex :: PrimType ty => Block ty -> Offset ty -> ty+blockRead :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> Offset ty -> prim ty+blockWrite :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> Offset ty -> ty -> prim ()+#ifdef ML_DSA_TESTING+blockIndex b off@(Offset i) =+ checkBounds (Block.length b) off $ indexPrimArray b i+blockRead mb off@(Offset i) = getSizeofMutablePrimArray mb >>= \sz ->+ checkBounds (CountOf sz) off $ readPrimArray mb i+blockWrite mb off@(Offset i) a = getSizeofMutablePrimArray mb >>= \sz ->+ checkBounds (CountOf sz) off $ writePrimArray mb i a+#else+blockIndex b (Offset i) = indexPrimArray b i+blockRead mb (Offset i) = readPrimArray mb i+blockWrite mb (Offset i) = writePrimArray mb i+#endif++erase :: (PrimMonad prim, PrimType ty) => CountOf ty -> MutableBlock ty (PrimState prim) -> prim ()+erase len@(CountOf n) = eraseBytes (I# (sizeOf# (toType len)) * n)+ where+ toType :: CountOf ty -> ty+ toType = undefined++eraseBytes :: PrimMonad prim => Int -> MutableBlock ty (PrimState prim) -> prim ()+eraseBytes (I# len) (MutablePrimArray mbarr) = primitive $ \s1 ->+ case setByteArray# mbarr 0# len 0# s1 of+ s2 -> (# s2, () #)++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+ sa = sizeofPrimArray a+ sb = sizeofPrimArray b++ loop !acc i+ | i == sa = acc+ | otherwise = do+ let va = indexPrimArray a i+ let vb = indexPrimArray b i+ loop (f acc va vb) (i + 1)+{-# INLINE foldZipWith #-}++length :: PrimType ty => Block ty -> CountOf ty+length = CountOf . sizeofPrimArray++getMutableLength :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> prim (CountOf ty)+getMutableLength = fmap CountOf . getSizeofMutablePrimArray++mutableContents :: MutableBlock ty s -> Ptr ty+mutableContents = mutablePrimArrayContents -- pinned only++new :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))+new (CountOf n) = newPrimArray n++newPinned :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim))+newPinned (CountOf n) = newPinnedPrimArray n++thaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))+thaw (PrimArray !barr) = unsafeSTToPrim $+ -- as optimization, combine both steps in a known monad and avoid+ -- round trip between byte length and element count+ thawByteArray ba 0 (sizeofByteArray ba) >>= \(MutableByteArray mbarr) ->+ return (MutablePrimArray mbarr)+ where ba = ByteArray barr++thawPinned :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim))+thawPinned (PrimArray !barr) = unsafeSTToPrim $ do+ let ba = ByteArray barr+ n = sizeofByteArray ba+ mb@(MutableByteArray mbarr) <- newPinnedByteArray n+ copyByteArray mb 0 ba 0 n+ return (MutablePrimArray mbarr)++#ifdef ML_DSA_TESTING+toList :: PrimType ty => Block ty -> [ty]+toList = primArrayToList+#endif++unsafeCast :: Block a -> Block b+unsafeCast (PrimArray b) = PrimArray b++unsafeCastMut :: MutableBlock a m -> MutableBlock b m+unsafeCastMut (MutablePrimArray mb) = MutablePrimArray mb++unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (Block ty)+unsafeFreeze = unsafeFreezePrimArray
+ src/BlockN.hs view
@@ -0,0 +1,262 @@+-- |+-- 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 RankNTypes #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ScopedTypeVariables #-}+module BlockN+ ( BlockN, MutableBlockN, BlockN.erase, BlockN.foldl', index, iterModify+ , iterSet, mapEqPrimSize, BlockN.read, BlockN.replicate, runNew, runThaw+ , runFold, BlockN.seq, BlockN.unsafeCast, BlockN.write, BlockN.zipWith+ , zipWithEqPrimSizeR+#ifdef ML_DSA_TESTING+ , create, BlockN.fromList, BlockN.toList+#endif+ ) where++import Control.DeepSeq (NFData(..))+import Control.Monad.ST++import Data.Proxy++import Base+import Block (MutableBlock, blockRead, blockWrite, erase, unsafeCastMut)+import Equality+import Fusion+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_DSA_TESTING+instance (Classified marking, Eq a, PrimType a) => Eq (BlockN marking n a) where+ BlockN a == BlockN b = SecureBlock.eq a b++instance (Classified marking, PrimType a, Show a) => Show (BlockN marking n a) where+ showsPrec d = SecureBlock.showsPrec d . unBlockN+#endif++instance NFData (BlockN marking n a) where+ rnf = SecureBlock.toNormalForm . unBlockN++instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where+ zero = BlockN.replicate zero+ {-# INLINE zero #-}+ (.+) = BlockN.zipWithEndoL (.+)+ {-# INLINE (.+) #-}+ (.-) = BlockN.zipWithEndoL (.-)+ {-# INLINE (.-) #-}+ neg = BlockN.mapEndo neg+ {-# INLINE neg #-}++newtype MutableBlockN (marking :: SecurityMarking) (n :: Nat) a m = MutableBlockN { unMutableBlockN :: MutableBlock a m }++instance (Classified marking, KnownNat n, PrimType a) => Fusion (BlockN marking n a) where+ type Mut (BlockN marking n a) s = MutableBlockN marking n a s+ newF = new Proxy+ thawF = thaw+ unsafeFreezeF = unsafeFreeze++-- Endomorphism specialization: a different implementation is substituted+-- wherever possible with rewrite rules. Identical input and output types give+-- a chance for a transformation to be fused on an existing mutable block.++{-# RULES+"mapEndo" [~2] forall f. BlockN.map f = mapEndo f+"zipWithEndoL" [~2] forall f. BlockN.zipWith f = zipWithEndoL f+ #-}++foldl' :: forall marking n a b. (KnownNat n, PrimType a) => (b -> a -> b) -> b -> BlockN marking n a -> b+foldl' f b (BlockN !a) = loop b 0+ where+ !sz = fromIntegral $ natVal (Proxy :: Proxy n)+ loop !acc i+ | i .==# sz = acc+ | otherwise = loop (acc `f` SecureBlock.index a i) (i + 1)+{-# INLINE foldl' #-}++index :: PrimType a => BlockN marking n a -> Offset a -> a+index = SecureBlock.index . unBlockN++replicate :: (Classified marking, KnownNat n, PrimType a) => a -> BlockN marking n a+replicate = create . const++#ifdef ML_DSA_TESTING+fromList :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => [a] -> Maybe (BlockN marking n a)+fromList elems+ | Prelude.length elems /= sz = Nothing+ | otherwise = Just $ runNew (Proxy :: Proxy marking) $ \mb -> go mb 0 elems+ where+ !sz = fromIntegral $ natVal (Proxy :: Proxy n)++ go !mb !i list = case list of+ [] -> return ()+ (x:xs) -> write mb i x >> go mb (i + 1) xs++toList :: PrimType a => BlockN marking n a -> [a]+toList = SecureBlock.toList . unBlockN+#endif++create :: (Classified marking, KnownNat n, PrimType ty)+ => (Offset ty -> ty)+ -> BlockN marking n ty+create initializer = runNew Proxy $ iterSet initializer+{-# INLINE create #-}++map :: (Classified marking, KnownNat n, PrimType a, PrimType b)+ => (a -> b) -> BlockN marking n a -> BlockN marking n b+map f (BlockN a) = BlockN.seq a $+ create $ \(Offset i) -> f (SecureBlock.index a (Offset i))+{-# INLINE [2] map #-}++mapEndo :: (Classified marking, KnownNat n, PrimType a)+ => (a -> a) -> BlockN marking n a -> BlockN marking n a+mapEndo = mapEqPrimSize+{-# INLINE mapEndo #-}++mapEqPrimSize :: (Classified marking, KnownNat n, EqPrimSize a b) => (a -> b) -> BlockN marking n a -> BlockN marking n b+mapEqPrimSize f = runContext . iterMapContext f . thawContext+{-# INLINE mapEqPrimSize #-}++erase :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)+ => MutableBlockN marking n ty (PrimState prim) -> prim ()+erase (MutableBlockN ma) = Block.erase sz ma+ where !sz = fromIntegral $ natVal (Proxy :: Proxy n)++iterModify :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)+ => (ty -> ty)+ -> MutableBlockN marking n ty (PrimState prim)+ -> prim ()+iterModify f = iterModifyIx (\_ x -> f x)+{-# INLINE iterModify #-}++iterModifyIx :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)+ => (Offset ty -> ty -> ty)+ -> MutableBlockN marking n ty (PrimState prim)+ -> prim ()+iterModifyIx f (MutableBlockN !ma) = loop 0+ where+ !sz = fromIntegral $ natVal (Proxy :: Proxy n)++ loop i+ | i .==# sz = pure ()+ | otherwise = blockRead ma i >>= \x -> blockWrite ma i (f i x) >> loop (i + 1)+{-# INLINE iterModifyIx #-}++iterSet :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim)+ => (Offset ty -> ty)+ -> MutableBlockN marking n ty (PrimState prim)+ -> prim ()+iterSet f (MutableBlockN !ma) = loop 0+ where+ !sz = fromIntegral $ natVal (Proxy :: Proxy n)++ loop i+ | i .==# sz = pure ()+ | otherwise = blockWrite ma i (f i) >> loop (i + 1)+{-# INLINE iterSet #-}++zipWith :: (Classified mc, KnownNat n, PrimType a, PrimType b, PrimType c)+ => (a -> b -> c) -> BlockN ma n a -> BlockN mb n b -> BlockN mc n c+zipWith f (BlockN a) (BlockN b) =+ BlockN.seq a $ BlockN.seq b $ create $ \(Offset i) ->+ f (SecureBlock.index a (Offset i)) (SecureBlock.index b (Offset i))+{-# INLINE [2] zipWith #-}++zipWithEndoL :: (Classified ma, KnownNat n, PrimType a, PrimType b)+ => (a -> b -> a) -> BlockN ma n a -> BlockN mb n b -> BlockN ma n a+zipWithEndoL = flip . zipWithEqPrimSizeR . flip+{-# INLINE zipWithEndoL #-}++zipWithEqPrimSizeR :: (Classified marking, KnownNat n, PrimType a, EqPrimSize b c)+ => (a -> b -> c) -> BlockN ma n a -> BlockN marking n b -> BlockN marking n c+zipWithEqPrimSizeR f a b = runContext (seqContext a (iterMapIxContext g (thawContext b)))+ where g = f . index a . Offset+{-# INLINE zipWithEqPrimSizeR #-}++unsafeCast :: BlockN marking n a -> SecureBlock marking b+unsafeCast = SecureBlock.unsafeCast . unBlockN++read :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> prim a+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)+{-# INLINE new #-}++runThaw :: (Classified marking, KnownNat n, PrimType a) => BlockN marking n a -> (forall s. MutableBlockN marking n a s -> ST s ()) -> BlockN marking n a+runThaw a f = runContext (modifyContext f (thawContext a))+{-# INLINE runThaw #-}++runNew :: (Classified marking, KnownNat n, PrimType a) => proxy marking -> (forall s. MutableBlockN marking n a s -> ST s ()) -> BlockN marking n a+runNew _ f = runContext (modifyContext f newContext)+{-# INLINE runNew #-}++runFold :: (Classified marking, KnownNat n, PrimType a, Foldable t) => BlockN marking n a -> (forall s. b -> MutableBlockN marking n a s -> ST s ()) -> t b -> BlockN marking n a+runFold a f = runContext . foldContext f (thawContext a)+{-# INLINE runFold #-}++seq :: (Classified marking, KnownNat n, PrimType a) => b -> BlockN marking n a -> BlockN marking n a+seq b a = runContext (seqContext b (thawContext a))+{-# INLINE seq #-}++thaw :: (Classified marking, PrimMonad prim) => BlockN marking n a -> prim (MutableBlockN marking n a (PrimState prim))+thaw = fmap MutableBlockN . SecureBlock.thaw . unBlockN++unsafeFreeze :: (Classified marking, PrimMonad prim) => MutableBlockN marking n a (PrimState prim) -> prim (BlockN marking n a)+unsafeFreeze = fmap BlockN . SecureBlock.unsafeFreeze . unMutableBlockN++unsafeMapIx :: forall marking n a b prim. (KnownNat n, EqPrimSize a b, PrimMonad prim) => (Int -> a -> b) -> MutableBlockN marking n a (PrimState prim) -> prim (MutableBlockN marking n b (PrimState prim))+unsafeMapIx f (MutableBlockN !ma) = MutableBlockN . ensureEqPrimSize witness <$> loop 0+ where+ witness = undefined :: a -> b+ !sz = fromIntegral $ natVal (Proxy :: Proxy n)+ loop i+ | i == sz = return (unsafeCastMut ma)+ | otherwise = do+ a <- blockRead ma (Offset i)+ blockWrite (unsafeCastMut ma) (Offset i) (f i a)+ loop (i + 1)+{-# INLINE unsafeMapIx #-}++--++iterMapContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b)+iterMapContext f = iterMapIxContext (\_ x -> f x)+{-# INLINE iterMapContext #-}++iterMapIxContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (Int -> a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b)+iterMapIxContext f = mapContext m+ where m = MapF { mapUpdate = unsafeMapIx f+ , mapInit = \x -> newF >>= \mb -> Prelude.seq x (iterSet (g x) mb) >> return mb+ }+ g x (Offset i) = f i (index x (Offset i))+{-# INLINE [1] iterMapIxContext #-}+++-- Fusion rules+--+-- "iterMapIxContext/iterMapIxContext" merges element-wise transformations as+-- single operations. For example @a .+ b .+ c@ becomes a single loop that+-- processes all input blocks in parallel and writes to the destination block.+--+-- "iterMapIxContext/seqContext" moves strictness annotations upstream so that+-- they do not prevent other rules from firing.++{-# RULES+"iterMapIxContext/seqContext" [~1] forall a f c. iterMapIxContext f (seqContext a c) = seqContext a (iterMapIxContext f c)+"iterMapIxContext/iterMapIxContext" [~1] forall f g c. iterMapIxContext f (iterMapIxContext g c) = iterMapIxContext (\i a -> f i (g i a)) c+ #-}
+ src/Builder.hs view
@@ -0,0 +1,99 @@+-- |+-- 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, promote, public+ , run, runRelaxed, runToBlock, secret, storable, unsafeCreate+ ) where++import Data.ByteArray (ByteArray, ByteArrayAccess)+import qualified Data.ByteArray as B++import Control.Monad.ST+import Control.Monad.ST.Unsafe++import Data.Semigroup+import Data.Word++import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (Storable(..))++import qualified GHC.Exts as Exts++import Base+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 s. Ptr Word8 -> 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 (getSize (Exts.inline builders)) $+ go (Exts.inline builders)+ where+ getSize = getSum . Prelude.mconcat . map (Sum . builderLength)+ go = foldr c (const $ return ())+ c b k = Exts.oneShot $ \p -> copyBuilderToPtr b p >> k (p `plusPtr` builderLength b)+ {-# INLINE mconcat #-}++instance Leak Builder++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 ()++public :: ByteArrayAccess ba => ba -> Builder marking+public b = unsafeCreate (B.length b) (B.copyByteArrayToPtr b)++promote :: Builder Pub -> Builder Sec+promote (Builder n f) = Builder n f++secret :: ByteArrayAccess ba => ba -> Builder Sec+secret = public++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)+ copyBuilderToPtr b (Block.mutableContents mb)+ Block.unsafeFreeze mb++storable :: Storable a => a -> Builder marking+storable !a = unsafeCreate (sizeOf a) (`poke` a)++unsafeCreate :: Int -> (Ptr a -> IO ()) -> Builder marking+unsafeCreate n f = create n (unsafeIOToST . f)+{-# INLINE unsafeCreate #-}
+ src/ByteArrayST.hs view
@@ -0,0 +1,46 @@+-- |+-- 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+ , fillBytes+ ) where++import Data.ByteArray (ByteArray, ByteArrayAccess)+import qualified Data.ByteArray as B++import Control.Monad.ST+import Control.Monad.ST.Unsafe++import Data.Word++import Foreign.Ptr (Ptr)+import qualified Foreign.Marshal.Utils as S+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 -> (Ptr p -> ST s a) -> ST s a+withByteArray b f = unsafeIOToST $ B.withByteArray b (unsafeSTToIO . f)+{-# INLINE withByteArray #-}++peek :: Storable a => Ptr a -> ST s a+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++fillBytes :: Ptr a -> Word8 -> Int -> ST s ()+fillBytes a v = unsafeIOToST . S.fillBytes a v
+ src/Crypto.hs view
@@ -0,0 +1,214 @@+-- |+-- Module : Crypto+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Crypto-related utilities like the ML-DSA hash functions, or more general+-- concerns like constant-time equality and selection.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE ScopedTypeVariables #-}+module Crypto+ ( ConstEqW(..), BoolW, trueW, andW, ltW, lteW, toBool+ , h, h64, h128+ , BlockDigest, unBlockDigest, hashToBlock+ , Digest, unDigest, hash+ ) where++import Crypto.Hash (Context)+import Crypto.Hash.Algorithms+import Crypto.Hash.IO++import Control.Exception (assert, mask_)+import Control.Monad.ST++import Data.ByteArray (ByteArrayAccess, Bytes, MemView(..), ScrubbedBytes)+import qualified Data.ByteArray as B++import Data.Bits+import Data.Word++import GHC.TypeNats++import Foreign.Marshal.Utils (fillBytes)+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (pokeByteOff)++import Block (Block)+import Builder (Builder)+import Marking (Classified)+import ScrubbedBlock (ScrubbedBlock)+import SecureBytes (SecureBytes)+import Vector (Vector)+import qualified Block+import qualified Builder+import qualified ByteArrayST as ST+import qualified ScrubbedBlock+import qualified Vector++unsafeShiftIR :: Word -> Int -> Word+unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int) `unsafeShiftR` s)+{-# INLINE unsafeShiftIR #-}++newtype BoolW = BoolW Word++#ifdef ML_DSA_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++notW :: BoolW -> BoolW+notW (BoolW a) = BoolW (complement a)++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++msbW :: Word -> Word+msbW x = x `unsafeShiftIR` (bitsW - 1)++eqW :: Word -> Word -> BoolW+eqW a b = isZeroW (a `xor` b)+ where isZeroW x = BoolW $ msbW (complement x .&. (x - 1))++ltW :: Word -> Word -> BoolW+ltW a b = BoolW $ msbW (a - b)++lteW :: Word -> Word -> BoolW+lteW a b = notW (ltW b a)++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 = 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 #-}++h :: (Classified marking, ByteArrayAccess (SecureBytes marking))+ => Int -> SecureBytes marking -> SecureBytes marking+h !len input = case someNatVal (fromIntegral (8 * len)) of+ SomeNat proxy -> Builder.run $ hashWith (alg proxy) input+ where+ alg :: proxy bitlen -> SHAKE256 bitlen+ alg _ = SHAKE256++h64 :: ByteArrayAccess (SecureBytes marking)+ => SecureBytes marking -> Builder marking+h64 = hashWith (SHAKE256 :: SHAKE256 512)++h128 :: ByteArrayAccess a+ => a -> Word8 -> Word8 -> (Bytes, ScrubbedBytes, ScrubbedBytes)+h128 xi !k !l = (a, b, c)+ where+ rho = B.takeView bs 32+ rho' = B.view bs 32 64+ kk = B.dropView bs 96++ !a = B.convert rho+ !b = B.convert rho'+ !c = B.convert kk++ bs :: ScrubbedBytes+ bs = Builder.run $ hashWith (SHAKE256 :: SHAKE256 1024) input++ input :: ScrubbedBytes+ input = B.unsafeCreate (len + 2) $ \p -> do+ B.copyByteArrayToPtr xi p+ pokeByteOff p len k+ pokeByteOff p (len + 1) l+ len = B.length xi++-- Override cryptonite/crypton types and hashing functions.+--+-- Standard type Digest is a newtype over an unpinned Block Word8, which+-- requires a trampoline to implement most Ptr access to the underlying byte+-- 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 ->+ hashMutableInit >>= \ctx -> mask_ $ do+ hashUpdateChunked (ctx :: MutableContext a) ba+ B.withByteArray ctx $ \pctx -> do+ hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig+ fillBytes pctx 0 (B.length ctx)++hashUpdateChunked :: (HashAlgorithm a, ByteArrayAccess ba) => MutableContext a -> ba -> IO ()+hashUpdateChunked ctx ba = B.withByteArray ba $ goChunked (B.length ba)+ where+ chunkSize = 1073741824 -- 1 GB++ goChunked remaining p+ | remaining > chunkSize = do+ hashMutableUpdate ctx (MemView p chunkSize)+ goChunked (remaining - chunkSize) (p `plusPtr` chunkSize)+ | otherwise = hashMutableUpdate ctx (MemView p remaining)
+ src/Crypto/PubKey/ML_DSA.hs view
@@ -0,0 +1,334 @@+-- |+-- Module : Crypto.PubKey.ML_DSA+-- License : BSD-3-Clause+-- Maintainer : Olivier Chéron <olivier.cheron@gmail.com>+-- Stability : provisional+-- Portability : unknown+--+-- Module-Lattice-based Digital Signature Algorithm (ML-DSA), defined+-- in <https://csrc.nist.gov/pubs/fips/204/final FIPS 204>.+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TypeFamilies #-}+module Crypto.PubKey.ML_DSA+ ( PublicKey, PrivateKey, Signature++ -- * Operations+ , generate, generateWith+ , Context, context, defaultContext, Randomness, randomness, deterministic++ -- ** Pure version+ --+ -- | This is the preferred and most common version of ML-DSA.+ , sign, signWith, verify++ -- ** Pre-hash version+ --+ -- | This version of ML-DSA can be used to sign or verify a message that is+ -- hashed externally. Typical use case is when the message is too large to+ -- fit into memory, as pre-hashing can be performed through an incremental+ -- API.+ --+ -- Note that signatures produced by HashML-DSA are incompatible with the+ -- "pure" version above. So both signing and verification must use the same+ -- version of ML-DSA. Keys can use either version indistinctly but it is+ -- advised to restrict them to one version only to avoid ambiguity.+ , PreHashAlgorithm, signDigest, signDigestWith, verifyDigest++ -- ** External µ version+ --+ -- | An internal signing and verfication API that operates on a message+ -- representative µ. For testing purpose only.+ , Mu, externalMu, signExternalMu, signExternalMuWith, verifyExternalMu++ -- ** Internal version+ --+ -- | An internal API that exposes algorithms @ML-DSA.Sign_internal@ and+ -- @ML-DSA.Verify_internal@. For testing purpose only.+ , signInternal, signInternalWith, verifyInternal++ -- * Parameter sets+ , ParamSet, ML_DSA_44, ML_DSA_65, ML_DSA_87++ -- * Conversions and checks+ , Decode(..), Encode(..), toPublic, checkKeyPair++ ) where++import Crypto.Hash (Digest)+import Crypto.Hash.Algorithms+import Crypto.Random++import Data.ByteArray (ByteArrayAccess, Bytes, ScrubbedBytes)+import qualified Data.ByteArray as B+import qualified Data.Memory.Endian as B++import Data.Word++import Builder+import Internal++-- | ML-DSA-44 (security category 2)+data ML_DSA_44 = ML_DSA_44 deriving Show+-- | ML-DSA-65 (security category 3)+data ML_DSA_65 = ML_DSA_65 deriving Show+-- | ML-DSA-87 (security category 5)+data ML_DSA_87 = ML_DSA_87 deriving Show++instance ParamSet ML_DSA_44 where+ type K ML_DSA_44 = 4+ type L ML_DSA_44 = 4+ getParams _ = Params 39 32 17 95232 6 2 3 78 80+instance ParamSet ML_DSA_65 where+ type K ML_DSA_65 = 6+ type L ML_DSA_65 = 5+ getParams _ = Params 49 48 19 261888 4 4 4 196 55+instance ParamSet ML_DSA_87 where+ type K ML_DSA_87 = 8+ type L ML_DSA_87 = 7+ getParams _ = Params 60 64 19 261888 4 2 3 120 75++-- | A byte string used during signature generation and verification for domain+-- separation. Always 255 bytes or less. Use 'defaultContext' unless specified+-- otherwise.+data Context = forall ba. ByteArrayAccess ba => Context ba++instance Eq Context where+ Context a == Context b = B.eq a b++instance Show Context where+ showsPrec d (Context b) = showParen (d > 10) $+ showString "Context " . showsPrec 11 (B.convert b :: Bytes)++instance ByteArrayAccess Context where+ length (Context ctx) = B.length ctx+ withByteArray (Context ctx) = B.withByteArray ctx+ copyByteArrayToPtr (Context ctx) = B.copyByteArrayToPtr ctx++-- | Smart constructor for a context value. Length of input must be 255 bytes+-- maximum.+context :: ByteArrayAccess ba => ba -> Maybe Context+context ctx+ | B.length ctx < 256 = Just $ Context ctx+ | otherwise = Nothing++-- | The default (empty) context.+defaultContext :: Context+defaultContext = Context (B.empty :: Bytes)++-- | A source of randomness to be used during signature generation.+-- Always 32 bytes.+--+-- The use of fresh randomness during signing helps mitigate side-channel+-- attacks.+data Randomness = forall ba. ByteArrayAccess ba => Randomness ba++instance Eq Randomness where+ Randomness a == Randomness b = B.constEq a b++instance Show Randomness where+#ifdef ML_DSA_TESTING+ showsPrec d (Randomness b) = showParen (d > 10) $+ showString "Randomness " . showsPrec 11 (B.convert b :: Bytes)+#else+ showsPrec _ _ = showString "Randomness"+#endif++instance ByteArrayAccess Randomness where+ length (Randomness _) = 32+ withByteArray (Randomness rnd) = B.withByteArray rnd+ copyByteArrayToPtr (Randomness rnd) = B.copyByteArrayToPtr rnd++-- | Smart constructor for randomness. Length of input must be 32 bytes.+randomness :: ByteArrayAccess ba => ba -> Maybe Randomness+randomness rnd+ | B.length rnd == 32 = Just $ Randomness rnd+ | otherwise = Nothing++-- | Enables a fully deterministic variant of the signing procedure. Can be+-- used if the signer has no access to a fresh source of randomness at signing+-- time. However this increases risks side-channel attacks, particularly fault+-- attacks.+deterministic :: Randomness+deterministic = Randomness (B.replicate 32 0 :: Bytes)++-- | Generate an ML-DSA key pair from a random seed.+generate :: (ParamSet a, MonadRandom m)+ => proxy a -> m (PublicKey a, PrivateKey a)+generate p = do+ xi <- getRandomBytes 32+ return $ Internal.keyGen p (xi :: ScrubbedBytes)++-- | Generate an ML-DSA key pair from the specified seed. Length of input+-- must be 32 bytes.+generateWith :: (ParamSet a, ByteArrayAccess seed)+ => proxy a -> seed -> Maybe (PublicKey a, PrivateKey a)+generateWith p xi+ | B.length xi /= 32 = Nothing+ | otherwise = Just $ Internal.keyGen p xi++-- | Generates an ML-DSA signature. Fresh randomness is acquired from a+-- 'MonadRandom' instance.+sign :: (ParamSet a, ByteArrayAccess message, MonadRandom m)+ => PrivateKey a -> message -> Context -> m (Signature a)+sign sk m ctx = do+ rnd <- getRandomBytes 32+ return $ Internal.sigGen sk (getPureM' m ctx) (rnd :: ScrubbedBytes)++-- | Generates an ML-DSA signature using explicit randomness.+signWith :: (ParamSet a, ByteArrayAccess message)+ => Randomness -> PrivateKey a -> message -> Context -> Signature a+signWith (Randomness rnd) sk m ctx = Internal.sigGen sk (getPureM' m ctx) rnd++-- | Verifies an ML-DSA signature. Returns @True@ when the signature is valid+-- for the message.+verify :: (ParamSet a, ByteArrayAccess message)+ => PublicKey a -> message -> Signature a -> Context -> Bool+verify pk m sig ctx = Internal.sigVer pk (getPureM' m ctx) sig++getPureM' :: (ByteArrayAccess message, ByteArrayAccess context)+ => message -> context -> ScrubbedBytes+getPureM' m ctx =+ let ctxLen = fromIntegral (B.length ctx) :: Word16+ in Builder.run $ Builder.storable (B.toBE ctxLen) <>+ Builder.public ctx <> Builder.public m++-- | Generates an HashML-DSA signature. Fresh randomness is acquired from a+-- 'MonadRandom' instance.+signDigest :: (ParamSet a, PreHashAlgorithm alg, MonadRandom m)+ => PrivateKey a -> Digest alg -> Context -> m (Signature a)+signDigest sk phm ctx = do+ rnd <- getRandomBytes 32+ return $ Internal.sigGen sk (getPreHashM' phm ctx) (rnd :: ScrubbedBytes)++-- | Generates a HashML-DSA signature using explicit randomness.+signDigestWith :: (ParamSet a, PreHashAlgorithm alg)+ => Randomness -> PrivateKey a -> Digest alg -> Context -> Signature a+signDigestWith (Randomness rnd) sk phm ctx =+ Internal.sigGen sk (getPreHashM' phm ctx) rnd++-- | Verifies a HashML-DSA signature. Returns @True@ when the signature is+-- valid for the message digest.+verifyDigest :: (ParamSet a, PreHashAlgorithm alg)+ => PublicKey a -> Digest alg -> Signature a -> Context -> Bool+verifyDigest pk phm sig ctx = Internal.sigVer pk (getPreHashM' phm ctx) sig++getPreHashM' :: (PreHashAlgorithm alg, ByteArrayAccess context)+ => Digest alg -> context -> ScrubbedBytes+getPreHashM' phm ctx =+ let ctxLen = 256 + fromIntegral (B.length ctx) :: Word16+ in Builder.run $ Builder.storable (B.toBE ctxLen) <> Builder.public ctx <>+ Builder.public (oid phm) <> Builder.public phm++-- | A message representative for ML-DSA. Always 64 bytes.+data Mu = forall ba. ByteArrayAccess ba => Mu ba++instance Eq Mu where+ Mu a == Mu b = B.eq a b++instance Show Mu where+ showsPrec d (Mu b) = showParen (d > 10) $+ showString "Mu " . showsPrec 11 (B.convert b :: Bytes)++instance ByteArrayAccess Mu where+ length (Mu _) = 64+ withByteArray (Mu mu) = B.withByteArray mu+ copyByteArrayToPtr (Mu mu) = B.copyByteArrayToPtr mu++-- | Smart constructor for a µ value. Length of input must be 64 bytes.+externalMu :: ByteArrayAccess ba => ba -> Maybe Mu+externalMu mu+ | B.length mu == 64 = Just $ Mu mu+ | otherwise = Nothing++-- | Generates a signature for a message representative µ value. Fresh+-- randomness is acquired from a 'MonadRandom' instance.+signExternalMu :: (ParamSet a, MonadRandom m)+ => PrivateKey a -> Mu -> m (Signature a)+signExternalMu sk (Mu mu) = do+ rnd <- getRandomBytes 32+ return $ Internal.sigGenMu sk (Builder.public mu) (rnd :: ScrubbedBytes)++-- | Generates a signature for a message representative µ value using explicit+-- randomness.+signExternalMuWith :: (ParamSet a)+ => Randomness -> PrivateKey a -> Mu -> Signature a+signExternalMuWith (Randomness rnd) sk (Mu mu) =+ Internal.sigGenMu sk (Builder.public mu) rnd++-- | Verifies a signature. Returns @True@ when the signature is valid for the+-- message representative µ.+verifyExternalMu :: ParamSet a+ => PublicKey a -> Mu -> Signature a -> Bool+verifyExternalMu pk (Mu mu) = Internal.sigVerMu pk (Builder.public mu)++-- | Internal procedure @ML-DSA.Sign_internal@. Generates a signature for a+-- formatted message. Fresh randomness is acquired from a 'MonadRandom'+-- instance.+signInternal :: (ParamSet a, ByteArrayAccess message, MonadRandom m)+ => PrivateKey a -> message -> m (Signature a)+signInternal sk m' = do+ rnd <- getRandomBytes 32+ return $ Internal.sigGen sk m' (rnd :: ScrubbedBytes)++-- | Internal procedure @ML-DSA.Sign_internal@. Generates a signature for a+-- formatted message using explicit randomness.+signInternalWith :: (ParamSet a, ByteArrayAccess message)+ => Randomness -> PrivateKey a -> message -> Signature a+signInternalWith (Randomness rnd) sk m' = Internal.sigGen sk m' rnd++-- | Internal procedure @ML-DSA.Verify_internal@. Returns @True@ when the+-- signature is valid for the formatted message.+verifyInternal :: (ParamSet a, ByteArrayAccess message)+ => PublicKey a -> message -> Signature a -> Bool+verifyInternal = Internal.sigVer++-- | Class of hash algorithms that can be used with HashML-DSA.+--+-- The algorithm should be chosen to provide enough collision resistance for the+-- target security level, for example SHA-384 or stronger with ML-DSA-65. This+-- is not enforced by the API.+class HashAlgorithm alg => PreHashAlgorithm alg where+ oid :: proxy alg -> Bytes++instance PreHashAlgorithm SHA256 where+ oid _ = nistHashOID 0x01++instance PreHashAlgorithm SHA384 where+ oid _ = nistHashOID 0x02++instance PreHashAlgorithm SHA512 where+ oid _ = nistHashOID 0x03++instance PreHashAlgorithm SHA224 where+ oid _ = nistHashOID 0x04++instance PreHashAlgorithm SHA512t_224 where+ oid _ = nistHashOID 0x05++instance PreHashAlgorithm SHA512t_256 where+ oid _ = nistHashOID 0x06++instance PreHashAlgorithm SHA3_224 where+ oid _ = nistHashOID 0x07++instance PreHashAlgorithm SHA3_256 where+ oid _ = nistHashOID 0x08++instance PreHashAlgorithm SHA3_384 where+ oid _ = nistHashOID 0x09++instance PreHashAlgorithm SHA3_512 where+ oid _ = nistHashOID 0x0A++instance PreHashAlgorithm (SHAKE128 256) where+ oid _ = nistHashOID 0x0B++instance PreHashAlgorithm (SHAKE256 512) where+ oid _ = nistHashOID 0x0C++nistHashOID :: Word8 -> Bytes+nistHashOID w8 =+ B.pack [ 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, w8 ]
+ src/Equality.hs view
@@ -0,0 +1,25 @@+-- |+-- Module : Equality+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Generate a constraint @'PrimSize' a ~ 'PrimSize' b@+--+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE TypeOperators #-}+module Equality+ ( EqPrimSize, ensureEqPrimSize+ ) where++import Data.Type.Equality++import Base++type EqPrimSize a b = (PrimType a, PrimType b, PrimSize a ~ PrimSize b)++eqPrimSize :: PrimSize a ~ PrimSize b => k a b -> PrimSize a :~: PrimSize b+eqPrimSize _ = Refl++ensureEqPrimSize :: EqPrimSize a b => k a b -> c -> c+ensureEqPrimSize op = case eqPrimSize op of Refl -> id
+ src/Fusion.hs view
@@ -0,0 +1,102 @@+-- |+-- Module : Fusion+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Infrastructure to decrease intermediate allocations and prefer in-place+-- mutation when possible+--+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE TypeFamilyDependencies #-}+module Fusion+ ( Fusion(..), MapF(..)+ , Context, runContext, newContext, thawContext, mapContext, modifyContext+ , foldContext, seqContext+ ) where++import Control.Monad ( forM_, (>=>) )+import Control.Monad.ST++-- class of values that can be mutated in the ST monad+class Fusion a where+ type Mut a s = mut | mut -> a+ newF :: ST s (Mut a s)+ thawF :: a -> ST s (Mut a s)+ unsafeFreezeF :: Mut a s -> ST s a++-- a transformation step in the fusion pipeline, with two implementations+-- provided: one that operates on an existing mutation context, and one that+-- initiates a new context from the input+data MapF a b = MapF+ { mapUpdate :: forall s. Mut a s -> ST s (Mut b s)+ , mapInit :: forall s. a -> ST s (Mut b s)+ }++-- MapF is almost a category except for the 'Fusion' constraint on objects+--+-- idMapF :: Fusion a => MapF a a+-- idMapF = MapF { mapUpdate = pure, mapInit = thawF }++composeMapF :: MapF b c -> MapF a b -> MapF a c+composeMapF m2 m1 = MapF+ { mapUpdate = mapUpdate m1 >=> mapUpdate m2+ , mapInit = mapInit m1 >=> mapUpdate m2+ }++-- fusion context+newtype Context a = Context (forall s. ST s (Mut a s))++newContext :: Fusion a => Context a+newContext = Context newF++thawContext :: Fusion a => a -> Context a+thawContext a = Context $ thawF a+{-# INLINE [0] thawContext #-}++modifyContext :: (forall s. Mut a s -> ST s ()) -> Context a -> Context a+modifyContext f = bindContext $ \ma -> f ma >> return ma++mapContext :: MapF a b -> Context a -> Context b+mapContext m = bindContext (mapUpdate m)+{-# INLINE [0] mapContext #-}++initContext :: MapF a b -> a -> Context b+initContext m a = Context $ mapInit m a++bindContext :: (forall s. Mut a s -> ST s (Mut b s)) -> Context a -> Context b+bindContext f (Context ctx) = Context $ ctx >>= f++foldContext :: Foldable t => (forall s. b -> Mut a s -> ST s ()) -> Context a -> t b -> Context a+foldContext f c bs = modifyContext (\ma -> forM_ bs $ \b -> f b ma) c++runContext :: Fusion a => Context a -> a+runContext (Context ctx) = runST (ctx >>= unsafeFreezeF)+{-# INLINE [0] runContext #-}++seqContext :: a -> Context b -> Context b+seqContext = seq+{-# INLINE [0] seqContext #-}+++-- Fusion rules+--+-- "thawContext/runContext" is the canonical optimization that eliminates an+-- allocation + value copy. Instead, it sequences two transformations on the+-- same mutation context.+--+-- "mapContext/seqContext" moves strictness annotations upstream so that they+-- do not prevent other rules from firing.+--+-- "mapContext/mapContext" is not strictly needed: the function is ultimately+-- inlined to the same code. But we keep it so that simplifications fire early+-- and do not wait for the final phase.+--+-- "mapContext/thawContext" is the rule that invokes mapInit instead of copying+-- the input and calling mapUpdate.++{-# RULES+"thawContext/runContext" [~0] forall c. thawContext (runContext c) = c+"mapContext/seqContext" [~0] forall a m c. mapContext m (seqContext a c) = seqContext a (mapContext m c)+"mapContext/mapContext" [~0] forall m1 m2 c. mapContext m2 (mapContext m1 c) = mapContext (composeMapF m2 m1) c+"mapContext/thawContext" [1] forall m a. mapContext m (thawContext a) = initContext m a+ #-}
+ src/Internal.hs view
@@ -0,0 +1,389 @@+-- |+-- Module : Internal+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- ML-DSA main internal algorithms+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE TypeFamilies #-}+module Internal+ ( ParamSet(..), Params(..), Encode(..), Decode(..)+ , PrivateKey, PublicKey, Signature+ , keyGen, sigGen, sigGenMu, sigVer, sigVerMu, toPublic, checkKeyPair+ ) where++import Control.DeepSeq (NFData(..))+import Control.Monad++import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes)+import qualified Data.ByteArray as B+import qualified Data.Memory.Endian as B++import Data.Bits+import Data.Proxy+import Data.Word++import Auxiliary (Hints, Rq, Tq)+import Base+import BlockN (BlockN)+import Builder (Builder)+import qualified Builder+import qualified Crypto+import Marking (SecurityMarking(..), Leak(..))+import Math+import Vector (Vector)+import qualified Auxiliary as Aux+import qualified Matrix+import qualified Vector++data Params (k :: Nat) (l :: Nat) = Params+ { tau :: {-# UNPACK #-} !Int+ , lambdaDiv4 :: {-# UNPACK #-} !Int+ , gamma1Bits :: {-# UNPACK #-} !Int+ , gamma2 :: {-# UNPACK #-} !Word32+ , wBits :: {-# UNPACK #-} !Int -- bitlen ((𝑞-1)/(2𝛾2)-1)+ , eta :: {-# UNPACK #-} !Int+ , twoEtaBits :: {-# UNPACK #-} !Int+ , beta :: {-# UNPACK #-} !Word+ , omega :: {-# UNPACK #-} !Int+ }++kdim :: KnownNat k => Params k l -> Int+kdim = fromIntegral . natVal . up+ where+ up :: a b (c :: k) -> Proxy b+ up _ = Proxy++ldim :: KnownNat l => Params k l -> Int+ldim = fromIntegral . natVal++-- | The class of ML-DSA parameter sets.+class (KnownNat (K a), KnownNat (L a)) => ParamSet a where+ type K a :: Nat+ type L a :: Nat+ getParams :: proxy a -> Params (K a) (L a)++-- | Utility class to serialize ML-DSA objects to byte arrays.+class Encode obj where+ -- | Serializes an object to a sequence of bytes.+ encode :: (ParamSet a, ByteArray ba) => obj a -> ba++-- | Utility class to deserialize ML-DSA 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-DSA private key.+data PrivateKey a = PrivateKey+ { skPub :: {-# UNPACK #-} !(PublicKey a)+ , skK :: {-# UNPACK #-} !ScrubbedBytes+ , skTr :: {-# UNPACK #-} !Bytes+ , skE1 :: {-# UNPACK #-} !ScrubbedBytes -- serialized s1 & s2+ , skE2 :: {-# UNPACK #-} !Bytes -- serialized t0+ , skS1 :: Vector (L a) (Tq Sec)+ , skS2 :: Vector (K a) (Tq Sec)+ , skT0 :: Vector (K a) (Tq Pub)+ }++skRho :: PrivateKey a -> Bytes+skRho = pkRho . skPub++skT1 :: PrivateKey a -> Vector (K a) (Tq Pub)+skT1 = pkT1 . skPub++skA :: PrivateKey a -> Vector (K a) (Vector (L a) (Tq Pub))+skA = pkA . skPub++-- | An ML-DSA public key.+data PublicKey a = PublicKey+ { pkRho :: {-# UNPACK #-} !Bytes+ , pkE :: {-# UNPACK #-} !Bytes -- serialized t1+ , pkT1 :: Vector (K a) (Tq Pub)+ , pkA :: Vector (K a) (Vector (L a) (Tq Pub))+ }++-- | An ML-DSA signature.+data Signature a = Signature+ { sigCt :: {-# UNPACK #-} !ScrubbedBytes+ , sigZ :: {-# UNPACK #-} !(Vector (L a) (Rq Sec))+ , sigH :: {-# UNPACK #-} !(Vector (K a) Hints)+ }++instance Eq (PrivateKey a) where+ a == b = Crypto.toBool $+ Crypto.constEqW (skRho a) (skRho b) `Crypto.andW`+ Crypto.constEqW (skK a) (skK b) `Crypto.andW`+ Crypto.constEqW (skTr a) (skTr b) `Crypto.andW`+ Crypto.constEqW (skE1 a) (skE1 b) `Crypto.andW`+ Crypto.constEqW (skE2 a) (skE2 b)++instance Eq (PublicKey a) where+ a == b = Crypto.toBool $+ Crypto.constEqW (pkRho a) (pkRho b) `Crypto.andW`+ Crypto.constEqW (pkE a) (pkE b)++instance Eq (Signature a) where+ a == b = Crypto.toBool $+ Crypto.constEqW (sigCt a) (sigCt b) `Crypto.andW`+ Crypto.constEqW (sigZ a) (sigZ b) `Crypto.andW`+ Crypto.constEqW (sigH a) (sigH b)++instance Show (PrivateKey a) where+#ifdef ML_DSA_TESTING+ showsPrec d sk = showParen (d > 10) $+ showString "PrivateKey " . showsPrec 11 (skEncode sk :: Bytes)+#else+ showsPrec _ _ = showString "PrivateKey"+#endif++instance Show (PublicKey a) where+ showsPrec d pk = showParen (d > 10) $+ showString "PublicKey " . showsPrec 11 (pkEncode pk :: Bytes)++instance ParamSet a => Show (Signature a) where+ showsPrec d sig = showParen (d > 10) $+ showString "Signature " . showsPrec 11 (sigEncode sig :: Bytes)++instance NFData (PrivateKey a) where+ rnf sk = rnf (skRho sk) `seq`+ rnf (skK sk) `seq`+ rnf (skTr sk) `seq`+ rnf (skE1 sk) `seq`+ rnf (skE2 sk)+ -- skS1, skS2, skT0, skT1, skA omitted because just for caching++instance NFData (PublicKey a) where+ rnf pk = rnf (pkRho pk) `seq` rnf (pkE pk)+ -- pkT1, pkA omitted because just for caching++instance NFData (Signature a) where+ rnf sig = rnf (sigCt sig) `seq`+ Vector.toNormalForm (sigZ sig) `seq`+ Vector.toNormalForm (sigH sig)++instance Encode PublicKey where+ encode = pkEncode++-- Encodes a public key for ML-DSA into a byte string+pkEncode :: ByteArray ba => PublicKey a -> ba+pkEncode pk = Builder.runRelaxed $+ Builder.bytes (pkRho pk) <> Builder.bytes (pkE pk)++instance Decode PublicKey where+ -- Reverses the procedure pkEncode+ decode p input = do+ guard (B.length input == 32 + 320 * k)+ let rho = B.convert $ B.takeView input 32+ pe = B.convert $ B.dropView input 32+ t1 = Vector.create $ \i -> Aux.simpleBitUnpack10 (view320 i)+ aa = expandA rho+ Just PublicKey { pkRho = rho, pkE = pe, pkT1 = Aux.ntt <$> t1, pkA = aa }+ where+ params = getParams p+ k = kdim params+ view320 (Offset i) = B.view input (32 + 320 * i) 320++instance Encode PrivateKey where+ encode = skEncode++-- Encodes a secret key for ML-DSA into a byte string+skEncode :: ByteArray ba => PrivateKey a -> ba+skEncode sk = Builder.runRelaxed $+ Builder.bytes (skRho sk) <>+ leak (Builder.bytes (skK sk)) <>+ Builder.bytes (skTr sk) <>+ leak (Builder.bytes (skE1 sk)) <>+ Builder.bytes (skE2 sk)++instance Decode PrivateKey where+ -- Reverses the procedure skEncode+ decode p input = do+ guard (B.length input == 128 + teb32 * (l + k) + 416 * k)+ let rho = B.convert $ B.view input 0 32+ kk = B.convert $ B.view input 32 32+ tr' = B.convert $ B.view input 64 64+ e1 = B.convert $ B.view input 128 (teb32 * (l + k))+ e2 = B.convert $ B.dropView input (128 + teb32 * (l + k))+ t0' = Vector.create $ \i -> Aux.bitUnpackSafe 13 (t0Elem i)+ aa = expandA rho+ s1 <- Vector.createMaybe $ \i -> Aux.bitUnpack (fromIntegral eta) twoEtaBits (s1Elem i)+ s2 <- Vector.createMaybe $ \i -> Aux.bitUnpack (fromIntegral eta) twoEtaBits (s2Elem i)+ let ss1 = Aux.ntt <$> s1+ ss2 = Aux.ntt <$> s2+ t = Aux.nttInv <$> Matrix.mmulAdd aa ss1 ss2+ (t1, t0) = Vector.unzipWith Aux.powerTwoRound t+ pk = PublicKey { pkRho = rho, pkE = pe, pkT1 = Aux.ntt <$> t1, pkA = aa }+ tr = Builder.run $ Crypto.h64 (encode pk :: Bytes)+ sk = PrivateKey { skPub = pk, skK = kk, skTr = tr, skE1 = e1, skE2 = e2, skS1 = ss1, skS2 = ss2, skT0 = Aux.ntt <$> t0 }+ pe = Builder.run $ Vector.concatMap Aux.simpleBitPack10 t1+ guard $ Crypto.toBool $+ Crypto.constEqW t0 t0' `Crypto.andW` Crypto.constEqW tr tr'+ return sk+ where+ s1Elem (Offset i) = B.view input (128 + teb32 * i) teb32+ s2Elem (Offset i) = B.view input (128 + teb32 * (l + i)) teb32+ t0Elem (Offset i) = B.view input (128 + teb32 * (l + k) + 416 * i) 416++ teb32 = 32 * twoEtaBits++ params@Params{..} = getParams p+ k = kdim params+ l = ldim params++instance Encode Signature where+ encode = sigEncode++-- Encodes a signature into a byte string+sigEncode :: (ParamSet a, ByteArray ba) => Signature a -> ba+sigEncode sig = Builder.runRelaxed $ leak $+ Builder.bytes (sigCt sig) <>+ Vector.concatMap (Aux.bitPackSafe (1 + gamma1Bits)) (sigZ sig) <>+ Aux.hintBitPack omega (sigH sig)+ where Params{..} = getParams sig++instance Decode Signature where+ -- Reverses the procedure sigEncode+ decode p input = do+ guard (B.length input == lambdaDiv4 + zElemLen * l + omega + k)+ let ct = B.convert $ B.view input 0 lambdaDiv4+ z = Vector.create $ \i -> Aux.bitUnpackSafe (1 + gamma1Bits) (zElem i)+ y = B.view input (lambdaDiv4 + zElemLen * l) (omega + k)+ h <- Aux.hintBitUnpack omega y+ return Signature { sigCt = ct, sigZ = z, sigH = h }+ where+ zElem (Offset i) = B.view input (lambdaDiv4 + zElemLen * i) zElemLen+ zElemLen = 32 * (1 + gamma1Bits)++ params@Params{..} = getParams p+ k = kdim params+ l = ldim params++-- Generates a public-private key pair from a seed+keyGen :: (ParamSet a, ByteArrayAccess seed) => proxy a -> seed -> (PublicKey a, PrivateKey a)+keyGen p xi = (pk, sk)+ where+ (rho, rho', kk) = Crypto.h128 xi (fromIntegral k) (fromIntegral l)+ aa = expandA rho+ (s1, s2) = expandS l eta rho'+ ss1 = Aux.ntt <$> s1+ ss2 = Aux.ntt <$> s2+ t = Aux.nttInv <$> Matrix.mmulAdd aa ss1 ss2+ (t1, t0) = Vector.unzipWith Aux.powerTwoRound t+ pk = PublicKey { pkRho = rho, pkE = pe, pkT1 = Aux.ntt <$> t1, pkA = aa }+ tr = Builder.run $ Crypto.h64 (encode pk :: Bytes)+ sk = PrivateKey { skPub = pk, skK = kk, skTr = tr, skE1 = e1, skE2 = e2, skS1 = ss1, skS2 = ss2, skT0 = Aux.ntt <$> t0 }+ pe = Builder.run $ Vector.concatMap Aux.simpleBitPack10 t1+ e1 = Builder.run $+ Vector.concatMap (Aux.bitPack (fromIntegral eta) twoEtaBits) s1 <>+ Vector.concatMap (Aux.bitPack (fromIntegral eta) twoEtaBits) s2+ e2 = Builder.run $ Vector.concatMap (Aux.bitPackSafe 13) t0++ params@Params{..} = getParams p+ k = kdim params+ l = ldim params++-- Deterministic algorithm to generate a signature for a formatted message 𝑀′+sigGen :: (ParamSet a, ByteArrayAccess m, ByteArrayAccess rnd) => PrivateKey a -> m -> rnd -> Signature a+sigGen sk m' = sigGenMu sk mu+ where+ tr = Builder.bytes (skTr sk)+ mu = Crypto.h64 (Builder.run $ Builder.promote tr <> Builder.secret m')++sigGenMu :: (ParamSet a, ByteArrayAccess rnd) => PrivateKey a -> Builder Sec -> rnd -> Signature a+sigGenMu sk mu rnd = loop 0+ where+ rhos = Crypto.h64 (Builder.run $ Builder.bytes (skK sk) <> Builder.secret rnd <> mu)++ loop kappa+ | Crypto.toBool cn1 && Crypto.toBool cn2 =+ Signature { sigCt = ct, sigZ = z, sigH = h }+ | otherwise = loop (kappa + l)+ where+ y = expandMask rhos kappa+ w = Aux.nttInv <$> Matrix.mmul (skA sk) (Aux.ntt <$> y)+ w1 = Aux.highBits gamma2 <$> w+ ct = Crypto.h lambdaDiv4 (Builder.run $ mu <> w1Encode wBits w1)+ cc = Aux.ntt $ Aux.sampleInBall tau ct+ cs1 = Aux.nttInv . (cc ..*) <$> Vector.seq cc (skS1 sk)+ cs2 = Aux.nttInv . (cc ..*) <$> Vector.seq cc (skS2 sk)+ z = y .+ cs1+ r0 = Aux.lowBits gamma2 <$> (w .- cs2)+ cn1 = Crypto.ltW (normVector z) ((1 `unsafeShiftL` gamma1Bits) - beta) `Crypto.andW`+ Crypto.ltW (normVector r0) (fromIntegral gamma2 - beta)+ ct0 = Aux.nttInv . (..* cc) <$> Vector.seq cc (skT0 sk)+ h = Vector.zipWith (Aux.makeHint gamma2) (neg ct0) (w .- cs2 .+ ct0)+ cn2 = Crypto.ltW (normVector ct0) (fromIntegral gamma2) `Crypto.andW`+ Crypto.lteW (countHints h) (fromIntegral omega)++ params@Params{..} = getParams sk+ l = ldim params++ -- Samples a vector 𝐲 ∈ 𝑅ℓ such that each polynomial 𝐲[𝑟] has+ -- coefficients between -𝛾1 + 1 and 𝛾1+ expandMask :: KnownNat l => Builder Sec -> Int -> Vector l (Rq Sec)+ expandMask rho mu' = Vector.create $ \(Offset r) ->+ let i = fromIntegral (mu' + r) :: Word16+ rho' = Builder.run $ rho <> Builder.storable (B.toLE i)+ v = Crypto.h (32 * c) rho'+ in Aux.bitUnpackSafe c v+ where c = 1 + gamma1Bits++-- Internal function to verify a signature 𝜎 for a formatted message 𝑀′+sigVer :: (ParamSet a, ByteArrayAccess m) => PublicKey a -> m -> Signature a -> Bool+sigVer pk m' = sigVerMu pk mu+ where+ tr = Crypto.h64 (encode pk :: Bytes)+ mu = Crypto.h64 (Builder.run $ Builder.promote tr <> Builder.secret m')++sigVerMu :: ParamSet a => PublicKey a -> Builder Sec -> Signature a -> Bool+sigVerMu pk mu sig =+ normVector (sigZ sig) < (1 `unsafeShiftL` gamma1Bits) - beta &&+ Crypto.toBool (Crypto.constEqW (sigCt sig) ct')+ where+ cc = Aux.ntt $ neg $ Aux.sampleInBall tau (sigCt sig)+ wApprox = Aux.nttInv <$> Matrix.mmulAdd (pkA pk) (Aux.ntt <$> sigZ sig) ((..* cc) <$> Vector.seq cc (pkT1 pk))+ w1 = Vector.zipWith (Aux.useHint gamma2) (sigH sig) wApprox+ ct' = Crypto.h lambdaDiv4 (Builder.run $ mu <> w1Encode wBits w1)++ Params{..} = getParams pk++-- Encodes a polynomial vector 𝐰1 into a byte string+w1Encode :: Int -> Vector n (BlockN Sec 256 Word32) -> Builder Sec+w1Encode wBits = Vector.concatMap (Aux.simpleBitPack wBits)++-- Samples a 𝑘 × ℓ matrix 𝐀 of elements of 𝑇𝑞+expandA :: (KnownNat k, KnownNat l) => Bytes -> Vector k (Vector l (Tq Pub))+expandA !rho = Matrix.create $ \(Offset s) (Offset r) ->+ Aux.rejNttPoly rho (fromIntegral s) (fromIntegral r)++-- Samples vectors 𝐬1 ∈ 𝑅ℓ and 𝐬2 ∈ 𝑅𝑘, each with polynomial coordinates+-- whose coefficients are in the interval [-𝜂, 𝜂]+expandS :: (KnownNat k, KnownNat l) => Int -> Int -> ScrubbedBytes -> (Vector l (Rq Sec), Vector k (Rq Sec))+expandS l eta rho = (s1, s2)+ where+ s1 = Vector.create $ \(Offset r) -> Aux.rejBoundedPoly rho eta (fromIntegral r)+ s2 = Vector.create $ \(Offset r) -> Aux.rejBoundedPoly rho eta (fromIntegral $ r + l)++countHints :: Vector n Hints -> Word+countHints = Vector.foldl' Aux.countFrom 0++normVector :: Vector n (Rq Sec) -> Word+normVector = fromIntegral . Aux.getNorm . Vector.foldl' Aux.normFrom mempty++-- | Returns the public key associated to the given private key.+toPublic :: PrivateKey a -> PublicKey a+toPublic = skPub++-- | Returns @True@ when the public key and private key both match. Note that+-- this does not fully guarantee that the key pair was properly generated.+checkKeyPair :: (PublicKey a, PrivateKey a) -> Bool+checkKeyPair (pk, sk) = Crypto.toBool $+ Crypto.constEqW (pkRho pk) (skRho sk) `Crypto.andW` Crypto.constEqW (pkT1 pk) (skT1 sk)+ -- consistency of skT0 with respect to skRho is verified when decoding+ -- the private key
+ src/Iterate.hs view
@@ -0,0 +1,47 @@+-- |+-- Module : Iterate+-- License : BSD-3-Clause+-- Copyright : (c) 2025 Olivier Chéron+--+-- @'offsetsFrom' x y@ is similar to @[x .. y - 1]@ but does not verify that+-- @x <= y@. This removes an unnecessary branch.+--+{-# LANGUAGE MagicHash #-}+module Iterate+ ( offsets, offsetsFrom+ )+where++import Control.Exception (assert)++import GHC.Exts++{-# INLINE offsets #-}+offsets :: Int -> [Int]+offsets = offsetsFrom 0++{-# INLINE offsetsFrom #-}+offsetsFrom :: Int -> Int -> [Int]+offsetsFrom xx@(I# x) yy@(I# y) = assert (xx <= yy) (indices x y)++-- Implementation is derived from enumFromTo, this acts as good producer for+-- list fusion.++{-# RULES+"indices" [~1] forall x y. indices x y = build (\c n -> indicesFB c n x y)+"indicesList" [1] indicesFB (:) [] = indices+ #-}++indices :: Int# -> Int# -> [Int]+indices x0 y = go x0+ where+ go x = if isTrue# (x <# y) then I# x : go (x +# 1#) else []+{-# NOINLINE [1] indices #-}++indicesFB :: (Int -> r -> r) -> r -> Int# -> Int# -> r+indicesFB c n x0 y = go x0+ where+ go x = if isTrue# (x <# y) then I# x `c` go (x +# 1#) else n+ -- Be very careful not to have more than one "c" so that when indicesFB is+ -- inlined we can inline whatever is bound to "c"+{-# INLINE [0] indicesFB #-}
+ src/Machine.hs view
@@ -0,0 +1,92 @@+-- |+-- Module : Machine+-- License : BSD-3-Clause+-- Copyright : (c) 2025 Olivier Chéron+--+-- Architecture-dependent utilities, to read/write unaligned machine words+-- in little-endian order+--+{-# LANGUAGE CPP #-}+module Machine+ ( WordM, WordLE, assertMultM, fromLE, toLE, wordBits, wordBytes+ ) where++#include "MachDeps.h"++-- Taken from `bytestring`, a list of architectures known to accept+-- unaligned loads and stores+#if defined(i386_HOST_ARCH) || defined(x86_64_HOST_ARCH) \+ || ((defined(arm_HOST_ARCH) || defined(aarch64_HOST_ARCH)) \+ && defined(__ARM_FEATURE_UNALIGNED)) \+ || defined(powerpc_HOST_ARCH) || defined(powerpc64_HOST_ARCH) \+ || defined(powerpc64le_HOST_ARCH)+#define MLDSA_ALLOW_UNALIGNED_OP 1++-- Little-endian conversion in `memory` / `ram` is avoided at compile+-- time only for AMD/Intel, here we will short circuit on ARM too+#if (defined(arm_HOST_ARCH) || defined(aarch64_HOST_ARCH)) \+ && !defined(WORDS_BIGENDIAN)+#define MLDSA_FORCE_LITTLE_ENDIAN_ARCH 1+#endif++#endif++import Control.Exception (assert)++#ifdef MLDSA_ALLOW_UNALIGNED_OP+import qualified Data.Memory.Endian as B+#endif++import Data.Bits+import Data.Word++#ifdef MLDSA_ALLOW_UNALIGNED_OP++-- our preferred word size+#if WORD_SIZE_IN_BITS == 64+type WordM = Word64+#else+type WordM = Word32+#endif++type WordLE = B.LE WordM++fromLE :: WordLE -> WordM+#ifdef MLDSA_FORCE_LITTLE_ENDIAN_ARCH+fromLE = B.unLE -- unwrap constructor with no byte swapping+#else+fromLE = B.fromLE -- byte swap if necessary+#endif++toLE :: WordM -> WordLE+#ifdef MLDSA_FORCE_LITTLE_ENDIAN_ARCH+toLE = B.LE -- wrap constructor with no byte swapping+#else+toLE = B.toLE -- byte swap if necessary+#endif++#else++-- unaligned memory access is not allowed so we fallback to one byte at a time+-- and endianness does not matter++type WordM = Word8+type WordLE = WordM++fromLE :: WordLE -> WordM+fromLE = id++toLE :: WordM -> WordLE+toLE = id++#endif++wordBits :: Int+wordBits = finiteBitSize (0 :: WordM)++wordBytes :: Int+wordBytes = div wordBits 8++assertMultM :: Int -> a -> a+assertMultM n = assert (n .&. mask == 0)+ where mask = wordBytes - 1
+ src/Marking.hs view
@@ -0,0 +1,133 @@+-- |+-- 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(..), index+ , Marking.toNormalForm, unsafeCast+#ifdef ML_DSA_TESTING+ , Marking.toList+#endif+ ) where++import Control.DeepSeq (NFData(..))+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 Base+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++ new :: (PrimType ty, PrimMonad prim) => proxy marking -> CountOf ty -> prim (MutableBlock ty (PrimState prim))+ thaw :: PrimMonad m => SecureBlock marking ty -> m (MutableBlock ty (PrimState m))+ unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty)++#ifdef ML_DSA_TESTING+ eq :: (Eq ty, PrimType ty) => SecureBlock marking ty -> SecureBlock marking ty -> Bool+ showsPrec :: (PrimType ty, Show ty) => Int -> SecureBlock marking ty -> ShowS+ lengthBlock :: PrimType ty => SecureBlock marking ty -> CountOf ty+#endif++ 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++ new _ = Block.new+ thaw = Block.thaw+ unsafeFreeze = Block.unsafeFreeze++#ifdef ML_DSA_TESTING+ eq = (==)+ showsPrec = Prelude.showsPrec+ 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++ new _ = ScrubbedBlock.new+ thaw = ScrubbedBlock.thaw+ unsafeFreeze = ScrubbedBlock.unsafeFreeze++#ifdef ML_DSA_TESTING+ eq = (==)+ showsPrec = Prelude.showsPrec+ lengthBlock = ScrubbedBlock.length+#endif++ type SecureBytes Sec = ScrubbedBytes+ unsafeCreate = ST.unsafeCreate+ {-# INLINE unsafeCreate #-}+ lengthBytes = B.length+ copyByteArrayToPtr = B.copyByteArrayToPtr+++-- for some functions we use the fact that Block and SecureBlock have the same+-- representation and implementation++unwrap :: SecureBlock marking a -> Block a+unwrap = unsafeCoerce++wrap :: Block b -> SecureBlock marking b+wrap = unsafeCoerce++index :: PrimType ty => SecureBlock marking ty -> Offset ty -> ty+index = blockIndex . unwrap++#ifdef ML_DSA_TESTING+toList :: PrimType ty => SecureBlock marking ty -> [ty]+toList = Block.toList . unwrap+#endif++toNormalForm :: SecureBlock marking ty -> ()+toNormalForm = rnf . unwrap++unsafeCast :: SecureBlock marking a -> SecureBlock marking b+unsafeCast = wrap . Block.unsafeCast . unwrap
+ src/Math.hs view
@@ -0,0 +1,55 @@+-- |+-- 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 CPP #-}+{-# LANGUAGE MultiParamTypeClasses #-}+module Math+ ( Add(..), Mul(..), MulAdd(..), BiMul(..), BiMulAdd(..)+ ) where++#if !(MIN_VERSION_base(4,20,0))+import Data.List (foldl')+#endif++infixl 7 .*+infixr 7 ..*+infixl 6 .+, .-++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+ {-# MINIMAL biMulAdd | biMulFold #-}++ -- invariant: biMulAdd a b c == a ..* b .+ c+ biMulAdd :: b -> a -> a -> a+ biMulAdd b a x = biMulFold x [(b, a)]+ {-# INLINE biMulAdd #-}++ -- repeated biMulAdd+ biMulFold :: Foldable t => a -> t (b, a) -> a+ biMulFold = foldl' $ \c (b, a) -> biMulAdd b a c+ {-# INLINE biMulFold #-}
+ src/Matrix.hs view
@@ -0,0 +1,35 @@+-- |+-- Module : Matrix+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- A matrix here is simply a vector of vectors. The module also implements two+-- utility functions 'mmul' and 'mmulAdd' that multiply a matrix and a vector.+--+module Matrix+ ( create, mmul, mmulAdd+ ) where++import Base+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 (Vector n ty) -> Offset ty -> ty+index a i = Vector.index (Vector.index a i)++mmul :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a+mmul a u = Vector.seq u (fmap (`Vector.dot` u) a)+{-# INLINE mmul #-}++mmulAdd :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a -> Vector m a+mmulAdd a u b = Vector.seq u (Vector.mapIx f b)+ where+ f (Offset i) bv =+ let g (Offset j) vu = (index a (Offset i) (Offset j), vu)+ in Vector.biMulFoldIndexWith g bv u+{-# INLINE mmulAdd #-}
+ src/ScrubbedBlock.hs view
@@ -0,0 +1,90 @@+-- |+-- 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.+--+-- A complication here is that we distinguish between mutable and immutable+-- values. And for resiliency against asynchronous exceptions, we need to+-- schedule block scrubbing with a finalizer right at the beginning when the+-- block is still in mutable form. Fortunately, for the perspective of the GC,+-- ByteArray# and MutableByteArray# are really the same heap object in disguise+-- and unsafeFreezeByteArray# is a true no-op. So the finalizer set on the+-- initial MutableByteArray# value gets transferred transparently to the final+-- ByteArray# form.+--+-- See GHC note [primOpEffect of unsafe freezes and thaws]+--+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}+module ScrubbedBlock+ ( ScrubbedBlock, foldZipWith, ScrubbedBlock.length+ , new, thaw, unsafeFreeze+ ) where++import Data.Primitive.PrimArray as Block++import Control.Exception (assert)+import Control.Monad.ST++import Data.Word++import Unsafe.Coerce++import Base+import Block (Block, MutableBlock)+import qualified Block++import GHC.Base (IO(IO))+import GHC.Exts (mkWeak#)++newtype ScrubbedBlock ty = ScrubbedBlock (Block ty)+ deriving (Eq, Show)++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 #-}++length :: PrimType ty => ScrubbedBlock ty -> CountOf ty+length (ScrubbedBlock b) = Block.length b++new :: (PrimType ty, PrimMonad prim) => CountOf ty -> prim (MutableBlock ty (PrimState prim))+new n = Block.newPinned n >>= scrubbed -- always pinned++thaw :: PrimMonad m => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m))+thaw (ScrubbedBlock b) = Block.thawPinned b >>= scrubbed -- always pinned++unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty)+unsafeFreeze mb = checkPinned <$> Block.unsafeFreeze mb+++{- internal -}++assertPinned :: Block ty -> a -> a+assertPinned mb = assert (Block.isPrimArrayPinned mb)++checkPinned :: Block ty -> ScrubbedBlock ty+checkPinned b = assertPinned b (ScrubbedBlock b)++scrubbed :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (MutableBlock ty (PrimState prim))+scrubbed b = unsafePrimFromIO (scheduleBlockScrubbing b >> return b)++wakeUpAfterInception :: MutableBlock ty s -> MutableBlock ty RealWorld+wakeUpAfterInception = unsafeCoerce -- sometimes disappointing++scheduleBlockScrubbing :: MutableBlock ty s -> IO ()+scheduleBlockScrubbing b = addBlockFinalizer b (scrub $ Block.unsafeCastMut b')+ where b' = wakeUpAfterInception b+{-# NOINLINE scheduleBlockScrubbing #-}++scrub :: MutableBlock Word8 RealWorld -> IO ()+scrub b = Block.getMutableLength b >>= \len -> Block.erase len b++addBlockFinalizer :: MutableBlock ty s -> IO () -> IO ()+addBlockFinalizer (Block.MutablePrimArray mbarr) (IO finalizer) = IO $ \s ->+ case mkWeak# mbarr () finalizer s of { (# s1, _ #) -> (# s1, () #) }
+ src/SecureBlock.hs view
@@ -0,0 +1,18 @@+-- |+-- 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, index, new, thaw+ , unsafeCast, unsafeFreeze, toNormalForm+#ifdef ML_DSA_TESTING+ , eq, Marking.showsPrec, toList+#endif+ ) where++import Marking
+ 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,253 @@+-- |+-- Module : Vector+-- License : BSD-3-Clause+-- Copyright : (c) 2025 Olivier Chéron+--+-- A vector of lifted elements with the vector dimension at type level.+-- Backed by type t'SmallArray' from primitive.+--+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ScopedTypeVariables #-}+module Vector+ ( Vector, Vector.concatMap, Vector.dot, Vector.index, Vector.toNormalForm+ , Vector.create, Vector.createMaybe, Vector.createMaybeAccum, Vector.foldl'+ , Vector.fold1ZipWith, Vector.biMulFoldIndexWith, mapIx, Vector.seq+ , Vector.unzipWith, Vector.zipWith+#ifdef ML_DSA_TESTING+ , Vector.replicateM+#endif+ ) where++import Data.Primitive.SmallArray++import Control.DeepSeq (NFData(..))+#ifdef ML_DSA_TESTING+import Control.Monad+#endif+import Control.Monad.ST++#if !(MIN_VERSION_base(4,20,0))+import Data.List as Prelude (foldl')+#endif+import Data.Proxy++import Base+import Iterate+import Math++type Array = SmallArray+type MArray ty s = SmallMutableArray s ty++newtype Vector (n :: Nat) a = Vector { unVector :: Array a }+ deriving (Eq, Show)++instance Functor (Vector n) where+ fmap = mapVector+ {-# INLINE fmap #-}++instance (Add a, KnownNat n) => Add (Vector n a) where+ zero = create (const zero)+ {-# INLINE zero #-}+ (.+) = Vector.zipWith (.+)+ {-# INLINE (.+) #-}+ (.-) = Vector.zipWith (.-)+ {-# INLINE (.-) #-}+ neg = mapVector neg+ {-# INLINE neg #-}++arrayCreate :: forall ty. CountOf ty -> (Offset ty -> ty) -> Array ty+arrayCreate n initializer = runST (arrayNew n >>= iter initializer)+ where+ iter :: PrimMonad prim => (Offset ty -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)+ iter f ma = loop 0+ where+ loop s@(Offset i)+ | s .==# n = unsafeFreezeSmallArray ma+ | otherwise = writeSmallArray ma i (f s) >> loop (s + 1)+ {-# INLINE loop #-}+ {-# INLINE iter #-}++arrayLength :: Array ty -> CountOf ty+arrayLength = CountOf . sizeofSmallArray++arrayMapIx :: (Offset a -> a -> b) -> Array a -> Array b+arrayMapIx f a = arrayCreate (CountOf sz) $ \(Offset i) ->+ let off = Offset i in f off (arrayIndex a off)+ where CountOf sz = arrayLength a++arrayNew :: PrimMonad prim => CountOf ty -> prim (MArray ty (PrimState prim))+arrayNew (CountOf c) = newSmallArray c placeholder+ where placeholder = error "arrayNew: unexpected evaluation"++create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a+create f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i))+ where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+{-# INLINE [2] create #-}++mapIx :: (Offset a -> a -> b) -> Vector n a -> Vector n b+mapIx = mapVectorIx+{-# INLINE [2] mapIx #-}++mapVector :: (a -> b) -> Vector n a -> Vector n b+mapVector f = mapVectorIx $ \_ x -> f x+{-# INLINE [2] mapVector #-}++mapVectorIx :: (Offset a -> a -> b) -> Vector n a -> Vector n b+mapVectorIx f = Vector <$> arrayMapIx f . unVector+{-# INLINE [1] mapVectorIx #-}++arrayIndex :: Array a -> Offset a -> a+#ifdef ML_DSA_TESTING+arrayIndex a off@(Offset i) =+ checkBounds (arrayLength a) off $ indexSmallArray a i++replicateM :: forall n m a. (KnownNat n, Applicative m) => m a -> m (Vector n a)+replicateM f = Vector . smallArrayFromList <$> Control.Monad.replicateM sz f+ where !sz = fromIntegral $ natVal (Proxy :: Proxy n)+#else+arrayIndex a (Offset i) = indexSmallArray a i+#endif++index :: Vector n a -> Offset a -> a+index = arrayIndex . unVector++concatMap :: Monoid b => (a -> b) -> Vector n a -> b+concatMap f = mconcat . mapToList f+{-# INLINE concatMap #-}++mapToList :: (a -> b) -> Vector n a -> [b]+mapToList f (Vector a) = Prelude.map (f . arrayIndex a . Offset) (offsets sa)+ where CountOf sa = arrayLength a++foldl' :: (b -> a -> b) -> b -> Vector n a -> b+foldl' f b (Vector a) = Prelude.foldl' g b (offsets sa)+ where+ g acc i = acc `f` arrayIndex a (Offset i)+ CountOf !sa = arrayLength a+{-# INLINE foldl' #-}++seq :: b -> Vector n a -> Vector n a+seq = Prelude.seq+{-# INLINE [1] seq #-}++zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c+zipWith f a (Vector !b) = mapVectorIx g a+ where g (Offset i) x = f x $ arrayIndex b (Offset i)+{-# INLINE [2] zipWith #-}++fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c+fold1ZipWith f g (Vector a) (Vector !b) =+ Prelude.foldl' ff gg (offsetsFrom 1 sa)+ where+ ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i))+ gg = g (arrayIndex a 0) (arrayIndex b 0)+ CountOf !sa = arrayLength a+{-# INLINE fold1ZipWith #-}++biMulFoldIndexWith :: BiMulAdd b a => (Offset ty -> t -> (b, a)) -> a -> Vector n t -> a+biMulFoldIndexWith f c (Vector a) =+ biMulFold c (map g $ offsets sa)+ where+ g i = f (Offset i) (arrayIndex a (Offset i))+ CountOf !sa = arrayLength a+{-# INLINE biMulFoldIndexWith #-}++dot :: BiMulAdd b a => Vector n b -> Vector n a -> a+dot (Vector b) (Vector a) =+ biMulFold (arrayIndex b 0 ..* arrayIndex a 0) (map g $ offsetsFrom 1 sb)+ where+ g i = (arrayIndex b (Offset i), arrayIndex a (Offset i))+ CountOf !sb = arrayLength b+{-# INLINE dot #-}++toNormalForm :: NFData a => Vector n a -> ()+toNormalForm = Vector.foldl' (\acc x -> acc `Prelude.seq` rnf x) ()++createMaybe :: KnownNat n => (Offset a -> Maybe a) -> Maybe (Vector n a)+createMaybe f = createMaybeAccum (\_ off -> ((), ) <$> f off) (const Just) ()+{-# INLINE createMaybe #-}++createMaybeAccum :: forall n a b r. KnownNat n => (b -> Offset a -> Maybe (b, a)) -> (b -> Vector n a -> Maybe r) -> b -> Maybe r+createMaybeAccum f k b0 = runST $ arrayNew (CountOf n) >>= \ma -> loop ma b0 0+ where+ !n = fromIntegral $ natVal (Proxy :: Proxy n)++ loop :: MArray a s -> b -> Offset a -> ST s (Maybe r)+ loop !ma b s@(Offset i)+ | s .==# CountOf n =+ unsafeFreezeSmallArray ma >>= \a -> return (b `k` Vector a)+ | otherwise = do+ case f b s of+ Nothing -> return Nothing+ Just (b', a) -> do+ writeSmallArray ma i a+ loop ma b' (s + 1)+{-# INLINE createMaybeAccum #-}++unzipWith :: forall n a b c. KnownNat n => (a -> (b, c)) -> Vector n a -> (Vector n b, Vector n c)+unzipWith f (Vector !a) = runST $ do+ mb <- arrayNew (CountOf n)+ mc <- arrayNew (CountOf n)+ loop mb mc 0+ where+ !n = fromIntegral $ natVal (Proxy :: Proxy n)++ loop :: MArray b s -> MArray c s -> Offset a -> ST s (Vector n b, Vector n c)+ loop !mb !mc s@(Offset i)+ | s .==# CountOf n = do+ b <- unsafeFreezeSmallArray mb+ c <- unsafeFreezeSmallArray mc+ return (Vector b, Vector c)+ | otherwise = do+ let (b, c) = f (arrayIndex a s)+ writeSmallArray mb i b+ writeSmallArray mc i c+ loop mb mc (s + 1)+{-# INLINE unzipWith #-}+++-- Rewrite rules+--+-- A first set of rules before Phase 2 performs simplifications between the four+-- main functions: create, mapVector, mapIx, and zipWith.+--+-- During Phase 2, the functions are then inlined. While function create+-- is replaced with its final form as call to arrayCreate, the three other+-- functions mapVector, mapIx, and zipWith all become calls to mapVectorIx.+-- Then, nested calls to mapVectorIx can further be simplified using rule+-- "mapVectorIx/mapVectorIx".+--+-- Finally at Phase 1 the function mapVectorIx is replaced with a call to+-- arrayCreate.+--+-- Both layers of rules have transformations that push calls to Vector.seq+-- outwards. Normal 'Prelude.seq' or the use of bang patterns would prevent+-- rewrite rules from firing.++{-# RULES+"mapVector/mapVector" [~2] forall f g a. mapVector f (mapVector g a) = mapVector (f . g) a+"mapVector/mapIx" [~2] forall f g a. mapVector f (mapIx g a) = mapIx (\i -> f . g i) a+"mapVector/create" [~2] forall f g. mapVector f (create g) = create (\(Offset i) -> f (g (Offset i)))+"mapVector/zipWith" [~2] forall f g a b. mapVector f (Vector.zipWith g a b) = Vector.zipWith (\x -> f . g x) a b+"mapVector/seq" [~2] forall f a b. mapVector f (Vector.seq b a) = Vector.seq b (mapVector f a)++"zipWith/mapVector left" [~2] forall f g a. Vector.zipWith f (mapVector g a) = Vector.zipWith (f . g) a+"zipWith/mapVector right" [~2] forall f g a b. Vector.zipWith f a (mapVector g b) = Vector.zipWith (\aa bb -> f aa (g bb)) a b+"zipWith/create left" [~2] forall f g. Vector.zipWith f (create g) = mapIx (\(Offset i) -> f (g (Offset i)))+"zipWith/create right" [~2] forall f g a. Vector.zipWith f a (create g) = mapIx (\(Offset i) x -> f x (g (Offset i))) a+"zipWith/zipWith right" [~2] forall f g a b c. Vector.zipWith f a (Vector.zipWith g b c) = Vector.zipWith (flip f) (Vector.zipWith g b c) a+"zipWith/seq left" [~2] forall f a b c. Vector.zipWith f (Vector.seq c a) b = Vector.seq c (Vector.zipWith f a b)+"zipWith/seq right" [~2] forall f a b c. Vector.zipWith f a (Vector.seq c b) = Vector.seq c (Vector.zipWith f a b)++"mapIx/mapVector" [~2] forall f g a. mapIx f (mapVector g a) = mapIx (\(Offset i) -> f (Offset i) . g) a+"mapIx/mapIx" [~2] forall f g a. mapIx f (mapIx g a) = mapIx (\(Offset i) -> f (Offset i) . g (Offset i)) a+"mapIx/create" [~2] forall f g. mapIx f (create g) = create (\(Offset i) -> f (Offset i) (g (Offset i)))+"mapIx/seq" [~2] forall f a b. mapIx f (Vector.seq b a) = Vector.seq b (mapIx f a)++"mapVectorIx/mapVectorIx" [~1] forall f g a. mapVectorIx f (mapVectorIx g a) = mapVectorIx (\(Offset i) -> f (Offset i) . g (Offset i)) a+"mapVectorIx/seq" [~1] forall f a b. mapVectorIx f (Vector.seq b a) = Vector.seq b (mapVectorIx f a)+ #-}
+ tests/KeyGen.hs view
@@ -0,0 +1,46 @@+-- |+-- 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+ , seed :: ByteString+ , pk :: ByteString+ , sk :: ByteString+ } deriving Show++instance FromJSON Test where+ parseJSON = withObject "Test" $ \o -> Test+ <$> o .: "tcId"+ <*> o .: "deferred"+ <*> o .:: "seed"+ <*> o .:: "pk"+ <*> o .:: "sk"
+ tests/SigExt.hs view
@@ -0,0 +1,57 @@+-- |+-- Module : SigExt+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Types that are common to signature-generation and verification test vectors+--+{-# LANGUAGE OverloadedStrings #-}+module SigExt+ ( TestExt(..) , PureExt(..), PreHashExt(..), ExternalMuExt(..)+ , InternalExt(..)+ ) where++import Data.Aeson+import Data.Aeson.Types+import Data.ByteString (ByteString)++import Util++class TestExt ext where+ parseExt :: Object -> Parser ext++data PureExt = PureExt+ { msgEncP :: ByteString+ , ctxEncP :: ByteString+ } deriving Show++instance TestExt PureExt where+ parseExt o = PureExt+ <$> o .:: "message"+ <*> o .:: "context"++data PreHashExt = PreHashExt+ { msgEncPH :: ByteString+ , ctxEncPH :: ByteString+ , hashPH :: String+ } deriving Show++instance TestExt PreHashExt where+ parseExt o = PreHashExt+ <$> o .:: "message"+ <*> o .:: "context"+ <*> o .: "hashAlg"++newtype ExternalMuExt = ExternalMuExt+ { muEncEM :: ByteString+ } deriving Show++instance TestExt ExternalMuExt where+ parseExt o = ExternalMuExt <$> o .:: "mu"++newtype InternalExt = InternalExt+ { msgEncIM :: ByteString+ } deriving Show++instance TestExt InternalExt where+ parseExt o = InternalExt <$> o .:: "message"
+ tests/SigGen.hs view
@@ -0,0 +1,81 @@+-- |+-- Module : SigVer+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Types that are specific to signature-generation test vectors+--+{-# LANGUAGE OverloadedStrings #-}+module SigGen+ ( TestGroup(..), Test(..), TestGroupPayload(..)+ ) where++import Data.Aeson+import Data.Aeson.Types+import Data.ByteString (ByteString)++import SigExt+import Util++data TestGroup = TestGroup+ { tgId :: Int+ , testType :: String+ , parameterSet :: String+ , deterministic :: Bool+ , signatureInterface :: String+ , preHash :: String+ , externalMu :: Bool+ , cornerCase :: String+ , payload :: TestGroupPayload+ } deriving Show++data TestGroupPayload+ = ModePure [Test PureExt]+ | ModePreHash [Test PreHashExt]+ | ModeExternalMu [Test ExternalMuExt]+ | ModeInternal [Test InternalExt]+ deriving Show++parsePayload :: Object -> Parser TestGroupPayload+parsePayload o = do+ si <- o .: "signatureInterface"+ ph <- o .: "preHash"+ em <- o .: "externalMu"+ case (si :: String, ph :: String, em) of+ ("external", "pure", False) -> ModePure <$> (o .: "tests")+ ("external", "preHash", False) -> ModePreHash <$> (o .: "tests")+ ("internal", "none", True) -> ModeExternalMu <$> (o .: "tests")+ ("internal", "none", False) -> ModeInternal <$> (o .: "tests")+ unknown -> fail ("parsePayload: unknown mode " ++ show unknown)++instance FromJSON TestGroup where+ parseJSON = withObject "TestGroup" $ \o -> TestGroup+ <$> o .: "tgId"+ <*> o .: "testType"+ <*> o .: "parameterSet"+ <*> o .: "deterministic"+ <*> o .: "signatureInterface"+ <*> o .: "preHash"+ <*> o .: "externalMu"+ <*> o .: "cornerCase"+ <*> parsePayload o++data Test ext = Test+ { tcId :: Int+ , deferred :: Bool+ , pkEnc :: ByteString+ , skEnc :: ByteString+ , sigEnc :: ByteString+ , rndEnc :: Maybe ByteString+ , tcExt :: ext+ } deriving Show++instance TestExt ext => FromJSON (Test ext) where+ parseJSON = withObject "Test" $ \o -> Test+ <$> o .: "tcId"+ <*> o .: "deferred"+ <*> o .:: "pk"+ <*> o .:: "sk"+ <*> o .:: "signature"+ <*> o .::? "rnd"+ <*> parseExt o
+ tests/SigVer.hs view
@@ -0,0 +1,79 @@+-- |+-- Module : SigVer+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- Types that are specific to signature-verification test vectors+--+{-# LANGUAGE OverloadedStrings #-}+module SigVer+ ( TestGroup(..), Test(..), TestGroupPayload(..)+ ) where++import Data.Aeson+import Data.Aeson.Types+import Data.ByteString (ByteString)++import SigExt+import Util++data TestGroup = TestGroup+ { tgId :: Int+ , testType :: String+ , parameterSet :: String+ , signatureInterface :: String+ , preHash :: String+ , externalMu :: Bool+ , payload :: TestGroupPayload+ } deriving Show++data TestGroupPayload+ = ModePure [Test PureExt]+ | ModePreHash [Test PreHashExt]+ | ModeExternalMu [Test ExternalMuExt]+ | ModeInternal [Test InternalExt]+ deriving Show++parsePayload :: Object -> Parser TestGroupPayload+parsePayload o = do+ si <- o .: "signatureInterface"+ ph <- o .: "preHash"+ em <- o .: "externalMu"+ case (si :: String, ph :: String, em) of+ ("external", "pure", False) -> ModePure <$> (o .: "tests")+ ("external", "preHash", False) -> ModePreHash <$> (o .: "tests")+ ("internal", "none", True) -> ModeExternalMu <$> (o .: "tests")+ ("internal", "none", False) -> ModeInternal <$> (o .: "tests")+ unknown -> fail ("parsePayload: unknown mode " ++ show unknown)++instance FromJSON TestGroup where+ parseJSON = withObject "TestGroup" $ \o -> TestGroup+ <$> o .: "tgId"+ <*> o .: "testType"+ <*> o .: "parameterSet"+ <*> o .: "signatureInterface"+ <*> o .: "preHash"+ <*> o .: "externalMu"+ <*> parsePayload o++data Test ext = Test+ { tcId :: Int+ , testPassed :: Bool+ , deferred :: Bool+ , pkEnc :: ByteString+ , skEnc :: ByteString+ , sigEnc :: ByteString+ , reason :: String+ , tcExt :: ext+ } deriving Show++instance TestExt ext => FromJSON (Test ext) where+ parseJSON = withObject "Test" $ \o -> Test+ <$> o .: "tcId"+ <*> o .: "testPassed"+ <*> o .: "deferred"+ <*> o .:: "pk"+ <*> o .:: "sk"+ <*> o .:: "signature"+ <*> o .: "reason"+ <*> parseExt o
+ tests/Tests.hs view
@@ -0,0 +1,739 @@+-- |+-- Module : Main+-- License : BSD-3-Clause+-- Copyright : (c) 2026 Olivier Chéron+--+-- The ML-DSA test suite. Can be instanciated twice, with and without the+-- @ML_DSA_TESTING@ macro to run property testing with assertions enabled in+-- the internal modules.+--+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE RankNTypes #-}+module Main (main) where++import Data.ByteArray (Bytes)+import qualified Data.ByteArray as B++import Crypto.Hash (hashWith)+import Crypto.Hash.Algorithms++import Control.Monad++import Data.List (isPrefixOf)+import Data.Maybe (fromJust)+import Data.Proxy++import GHC.IO.Exception (IOErrorType(..))++import System.Directory (doesFileExist)+import System.FileLock (FileLock, SharedExclusive(..), unlockFile, withFileLock)+import System.IO.Error (catchIOError, mkIOError)+import System.Process (readProcess)++#ifdef ML_DSA_TESTING++import Data.Bits+import Data.Word++import GHC.TypeNats++import Foreign.Ptr (plusPtr)++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_DSA as Lib++import Test.Tasty+import Test.Tasty.HUnit+import Test.Tasty.QuickCheck++import qualified KeyGen+import qualified SigExt+import qualified SigGen+import qualified SigVer+import qualified Vectors++arbitraryBytes :: Int -> Gen Bytes+arbitraryBytes n = B.pack <$> vectorOf n arbitrary++#ifdef ML_DSA_TESTING++truncateRq :: Word32 -> Rq Sec -> Rq Sec+truncateRq m = fromJust . fromCoeffs . map f . toCoeffs+ where f = toZq . (`mod` m) . fromZq++truncateRq' :: Word32 -> Rq Sec -> Rq Sec+truncateRq' m = fromJust . fromCoeffs . map f . toCoeffs+ where f = toZq . (\x -> x `mod` (2 * m - 1) + 8380418 - m) . fromZq++newtype FE = FE { unFE :: Zq } deriving Show++instance Arbitrary FE where+#if (MIN_VERSION_tasty_quickcheck(0,10,2))+ arbitrary = FE . toZq <$> chooseBoundedIntegral (0, 8380416)+#else+ arbitrary = FE . toZq <$> choose (0, 8380416)+#endif++newtype ME = ME { unME :: Mq } deriving Show++instance Arbitrary ME where+ arbitrary = ME . toMontgomery . unFE <$> arbitrary++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++arbitraryHints :: Gen Hints+arbitraryHints = fromJust . fromBools <$> vectorOf 256 h+ where h = elements (True : replicate 9 False) -- probability 1/10++arbitraryHintsVector :: KnownNat k => proxy k -> Gen (Vector k Hints)+arbitraryHintsVector _ = Vector.replicateM arbitraryHints++countHints :: Vector n Hints -> Word+countHints = Vector.foldl' countFrom 0++newtype D = D Int deriving Show++instance Arbitrary D where+ arbitrary = sized $ \n -> D <$> choose (0, min n 22)++data Dim = forall (n :: Nat). KnownNat n => Dim (Proxy n)++instance Show Dim where+ show (Dim n) = show n++instance Arbitrary Dim where+ arbitrary = sized $ \n -> toDim <$> choose (1, min (1 + n) 9)++toDim :: Int -> Dim+toDim n = case someNatVal (fromIntegral n) of SomeNat p -> Dim p++type VElem = Mq -- test with any ring but Tq would also work here++arbitraryVector :: KnownNat n => proxy n -> Gen (Vector n VElem)+arbitraryVector _ = Vector.replicateM (unME <$> arbitrary)++arbitraryMatrix :: (KnownNat m, KnownNat n) => proxy n -> proxy m -> Gen (Vector n (Vector m VElem))+arbitraryMatrix _ m = Vector.replicateM (arbitraryVector m)++simpleBitPackBytes :: Int -> BlockN Sec 256 Word32 -> Bytes+simpleBitPackBytes d = runBytes . simpleBitPack d++runPubBytes :: Builder Pub -> Bytes+runPubBytes = Builder.run++runPubUnaligned :: Builder Pub -> UnalignedBytes+runPubUnaligned = Builder.runRelaxed++runBytes :: Builder Sec -> Bytes+runBytes = runPubBytes . leak++runUnaligned :: Builder Sec -> UnalignedBytes+runUnaligned = runPubUnaligned . leak++newtype UnalignedBytes = UnalignedBytes (B.View Bytes)+ deriving (Eq,Ord,Show)++instance Semigroup UnalignedBytes where+ UnalignedBytes a <> UnalignedBytes b =+ B.allocAndFreeze (na + nb) $ \p -> do+ B.copyByteArrayToPtr a p+ B.copyByteArrayToPtr b (p `plusPtr` na)+ where+ na = B.length a+ nb = B.length b++instance Monoid UnalignedBytes where+ mempty = B.convert (B.empty :: Bytes)++instance B.ByteArrayAccess UnalignedBytes where+ length (UnalignedBytes v) = B.length v+ withByteArray (UnalignedBytes v) = B.withByteArray v++instance B.ByteArray UnalignedBytes where+ allocRet n f = do+ let build ba = UnalignedBytes (B.dropView ba offset)+ (a, ba) <- B.allocRet (n + offset) $ \p -> f (p `plusPtr` offset)+ return (a, build ba)+ where offset = 1++#endif++data P = forall a. (ParamSet a, Show a) => P (Proxy a)++instance Show P where+ show (P p) = show p++instance Arbitrary P where+ arbitrary = elements+ [ P (Proxy :: Proxy ML_DSA_44)+ , P (Proxy :: Proxy ML_DSA_65)+ , P (Proxy :: Proxy ML_DSA_87)+ ]++toP :: String -> P+toP "ML-DSA-44" = P (Proxy :: Proxy ML_DSA_44)+toP "ML-DSA-65" = P (Proxy :: Proxy ML_DSA_65)+toP "ML-DSA-87" = P (Proxy :: Proxy ML_DSA_87)+toP paramSet = error ("unknown parameter set " ++ paramSet)++data PH = forall alg. (PreHashAlgorithm alg, Show alg) => PH (Proxy alg)++instance Show PH where+ show (PH ph) = show ph++instance Arbitrary PH where+ arbitrary = elements+ [ PH (Proxy :: Proxy SHA224)+ , PH (Proxy :: Proxy SHA256)+ , PH (Proxy :: Proxy SHA384)+ , PH (Proxy :: Proxy SHA512)+ , PH (Proxy :: Proxy SHA512t_224)+ , PH (Proxy :: Proxy SHA512t_256)+ , PH (Proxy :: Proxy SHA3_224)+ , PH (Proxy :: Proxy SHA3_256)+ , PH (Proxy :: Proxy SHA3_384)+ , PH (Proxy :: Proxy SHA3_512)+ , PH (Proxy :: Proxy (SHAKE128 256))+ , PH (Proxy :: Proxy (SHAKE256 512))+ ]++toPH :: String -> PH+toPH "SHA2-224" = PH (Proxy :: Proxy SHA224)+toPH "SHA2-256" = PH (Proxy :: Proxy SHA256)+toPH "SHA2-384" = PH (Proxy :: Proxy SHA384)+toPH "SHA2-512" = PH (Proxy :: Proxy SHA512)+toPH "SHA2-512/224" = PH (Proxy :: Proxy SHA512t_224)+toPH "SHA2-512/256" = PH (Proxy :: Proxy SHA512t_256)+toPH "SHA3-224" = PH (Proxy :: Proxy SHA3_224)+toPH "SHA3-256" = PH (Proxy :: Proxy SHA3_256)+toPH "SHA3-384" = PH (Proxy :: Proxy SHA3_384)+toPH "SHA3-512" = PH (Proxy :: Proxy SHA3_512)+toPH "SHAKE-128" = PH (Proxy :: Proxy (SHAKE128 256))+toPH "SHAKE-256" = PH (Proxy :: Proxy (SHAKE256 512))+toPH hashAlg = error ("unknown hash algorithm " ++ hashAlg)++preHash :: Proxy alg -> alg+preHash _ = undefined++withPreHash :: String -> (forall alg. PreHashAlgorithm alg => alg -> r) -> r+withPreHash s f = case toPH s of+ PH ph -> f (preHash ph)++newtype Ctx = Ctx { unCtx :: Context } deriving Show++instance Arbitrary Ctx where+ arbitrary = sized $ \n -> do+ len <- choose (0, min n 255)+ Ctx . fromJust . Lib.context <$> arbitraryBytes len++newtype Msg = Msg { unMsg :: Bytes } deriving Show++instance Arbitrary Msg where+ arbitrary = Msg . B.pack <$> arbitrary++newtype ExternalMu = ExternalMu { unExternalMu :: Mu } deriving Show++instance Arbitrary ExternalMu where+ arbitrary = ExternalMu . fromJust . Lib.externalMu <$> arbitraryBytes 64++withVectors :: (IO () -> TestTree) -> TestTree+withVectors = withResource alloc free+ where+ scriptPath = "tests/get-vectors.sh"+ free _ = return ()+ whenNeeded action = do+ keyGenExists <- doesFileExist "tests/keyGen.json.gz"+ sigGenExists <- doesFileExist "tests/sigGen.json.gz"+ sigVerExists <- doesFileExist "tests/sigVer.json.gz"+ unless (keyGenExists && sigGenExists && sigVerExists) action+ alloc = withTestLock Shared $ \lock -> whenNeeded $ do+ unlockFile lock -- sanity before lock upgrade+ withTestLock Exclusive $ \_ -> whenNeeded $ 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)+ )++withTestLock :: SharedExclusive -> (FileLock -> IO a) -> IO a+withTestLock mode what = do+ -- locking a hidden file in the directory prevents a race condition between+ -- two instances of the test suite trying both to download the test vectors+ -- (otherwise one instance may try to run with files not fully downloaded+ -- yet by the other instance)+ let path = "tests/.lock"+ exists <- doesFileExist path+ unless exists $ writeFile path "DO NOT DELETE"+ withFileLock path mode what++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+ pks = Lib.encode pk+ sks = Lib.encode sk+ (pk, sk) = fromJust $ Lib.generateWith p (KeyGen.seed t)+ assertEqual ("pk mismatch for tcId=" ++ show tcId) (KeyGen.pk t) pks+ assertEqual ("sk mismatch for tcId=" ++ show tcId) (KeyGen.sk t) sks++sigGenVectors :: (String -> IO ()) -> Assertion+sigGenVectors step = do+ step "Reading test vectors ..."+ file <- Vectors.readJson "tests/sigGen.json.gz"+ forM_ (Vectors.testGroups file) $ \group -> do+ let paramSet = SigGen.parameterSet group+ step (paramSet ++ " (" ++ mode group ++ ")")+ case toP paramSet of+ P p -> case SigGen.payload group of+ SigGen.ModePure tests ->+ forM_ tests (testExt doSignPure p)+ SigGen.ModePreHash tests ->+ forM_ tests (testExt doSignPreHash p)+ SigGen.ModeExternalMu tests ->+ forM_ tests (testExt doSignExternalMu p)+ SigGen.ModeInternal tests ->+ forM_ tests (testExt doSignInternal p)+ where+ mode group+ | SigGen.signatureInterface group == "external" = SigGen.preHash group+ | SigGen.externalMu group = "external µ"+ | otherwise = SigGen.signatureInterface group+ testExt signExtWith p test =+ assertEqual ("sig mismatch for tcId=" ++ show tcId) sig' sig+ where+ tcId = SigGen.tcId test+ ext = SigGen.tcExt test+ sk = fromJust $ Lib.decode p (SigGen.skEnc test)+ sig' = fromJust $ Lib.decode p (SigGen.sigEnc test)+ rnd = maybe Lib.deterministic (fromJust . Lib.randomness) (SigGen.rndEnc test)+ sig = signExtWith ext rnd sk+ doSignPure ext rnd sk = Lib.signWith rnd sk m ctx+ where+ m = SigExt.msgEncP ext+ ctx = fromJust $ Lib.context (SigExt.ctxEncP ext)+ doSignPreHash ext rnd sk = withPreHash (SigExt.hashPH ext) $ \alg ->+ let phm = hashWith alg m+ in Lib.signDigestWith rnd sk phm ctx+ where+ m = SigExt.msgEncPH ext+ ctx = fromJust $ Lib.context (SigExt.ctxEncPH ext)+ doSignExternalMu ext rnd sk = Lib.signExternalMuWith rnd sk mu+ where mu = fromJust $ Lib.externalMu (SigExt.muEncEM ext)+ doSignInternal ext rnd sk = Lib.signInternalWith rnd sk m+ where m = SigExt.msgEncIM ext++sigVerVectors :: (String -> IO ()) -> Assertion+sigVerVectors step = do+ step "Reading test vectors ..."+ file <- Vectors.readJson "tests/sigVer.json.gz"+ forM_ (Vectors.testGroups file) $ \group -> do+ let paramSet = SigVer.parameterSet group+ step (paramSet ++ " (" ++ mode group ++ ")")+ case toP paramSet of+ P p -> case SigVer.payload group of+ SigVer.ModePure tests ->+ forM_ tests (testExt doVerifyPure p)+ SigVer.ModePreHash tests ->+ forM_ tests (testExt doVerifyPreHash p)+ SigVer.ModeExternalMu tests ->+ forM_ tests (testExt doVerifyExternalMu p)+ SigVer.ModeInternal tests ->+ forM_ tests (testExt doVerifyInternal p)+ where+ mode group+ | SigVer.signatureInterface group == "external" = SigVer.preHash group+ | SigVer.externalMu group = "external µ"+ | otherwise = SigVer.signatureInterface group+ testExt doVerify p test =+ case Lib.decode p (SigVer.sigEnc test) of+ Just sig -> do+ assertBool ("opposite outcome for tcId=" ++ show tcId)+ (SigVer.testPassed test == doVerify ext pk sig)+ Nothing ->+ assertBool ("could not decode signature for tcId=" ++ show tcId)+ ("modified signature - " `isPrefixOf` SigVer.reason test)+ where+ tcId = SigVer.tcId test+ ext = SigVer.tcExt test+ pk = fromJust $ Lib.decode p (SigVer.pkEnc test)+ doVerifyPure ext pk sig = Lib.verify pk m sig ctx+ where+ m = SigExt.msgEncP ext+ ctx = fromJust $ Lib.context (SigExt.ctxEncP ext)+ doVerifyPreHash ext pk sig = withPreHash (SigExt.hashPH ext) $ \alg ->+ let phm = hashWith alg m+ in Lib.verifyDigest pk phm sig ctx+ where+ m = SigExt.msgEncPH ext+ ctx = fromJust $ Lib.context (SigExt.ctxEncPH ext)+ doVerifyExternalMu ext pk = Lib.verifyExternalMu pk mu+ where mu = fromJust $ Lib.externalMu (SigExt.muEncEM ext)+ doVerifyInternal ext pk = Lib.verifyInternal pk m+ where m = SigExt.msgEncIM ext++main :: IO ()+main = defaultMain $ testGroup "mldsa"+ [ withVectors $ \_ -> testGroup "vectors"+ [ testCaseSteps "keyGen" keyGenVectors+ , testCaseSteps "sigGen" sigGenVectors+ , testCaseSteps "sigVer" sigVerVectors+ ]+ , testGroup "properties"+ [ testGroup "ML-DSA"+ [ testProperty "sign/verify (pure)" $ \(P p) (Msg m) (Ctx ctx) -> ioProperty $ do+ (pk, sk) <- Lib.generate p+ sig <- Lib.sign sk m ctx+ return (Lib.verify pk m sig ctx)+ , testProperty "sign/verify (preHash)" $ \(P p) (PH ph) (Msg m) (Ctx ctx) -> ioProperty $ do+ (pk, sk) <- Lib.generate p+ let phm = hashWith (preHash ph) m+ sig <- Lib.signDigest sk phm ctx+ return (Lib.verifyDigest pk phm sig ctx)+ , testProperty "sign/verify (external µ)" $ \(P p) (ExternalMu mu) -> ioProperty $ do+ (pk, sk) <- Lib.generate p+ sig <- Lib.signExternalMu sk mu+ return (Lib.verifyExternalMu pk mu sig)+ , testProperty "sign/verify (internal)" $ \(P p) (Msg m) -> ioProperty $ do+ (pk, sk) <- Lib.generate p+ sig <- Lib.signInternal sk m+ return (Lib.verifyInternal pk m sig)+ , testProperty "encode/decode keys" $ \(P p) -> ioProperty $ do+ (pk, sk) <- Lib.generate p+ return $ conjoin+ [ Just pk === Lib.decode p (Lib.encode pk :: Bytes)+ , Just sk === Lib.decode p (Lib.encode sk :: Bytes)+ ]+ , testProperty "encode/decode signatures" $ \(P p) (Msg m) (Ctx ctx) -> ioProperty $ do+ (_, sk) <- Lib.generate p+ sig <- Lib.sign sk m ctx+ return $ Just sig === Lib.decode p (Lib.encode sig :: Bytes)+ , testProperty "toPublic" $ \(P p) -> ioProperty $ do+ (ek, sk) <- Lib.generate p+ return (ek === toPublic sk)+ , testProperty "checkKeyPair" $ \(P p) -> ioProperty $+ checkKeyPair <$> Lib.generate p+ ]+#ifdef ML_DSA_TESTING+ , testGroup "bitRev8"+ [ testCase "powers of two" $+ let powers = [1, 2, 4, 8, 16, 32, 64, 128]+ in reverse powers @=? map bitRev8 powers+ , testProperty "or" $ \a b ->+ bitRev8 (a .|. b) === bitRev8 a .|. bitRev8 b+ , testProperty "not" $ \a ->+ let comp = xor 255+ in bitRev8 (comp a) === comp (bitRev8 a)+ , testProperty "involutive" $ \a ->+ a === bitRev8 (bitRev8 a)+ , testProperty "preserves bit count" $ \a ->+ popCount a === popCount (bitRev8 a)+ ]+ , testGroup "compression"+ [ testProperty "powerTwoRound" $ \(FE r) ->+ let (r1, r0) = powerTwoRoundZq r+ in fromZq r === (fromZq r1 + fromZq r0) `mod` 8380417 .&&.+ (fromZq r1 `mod` 8192 == 0) .&&.+ (fromZq r0 <= 4096 .||. fromZq r0 > 8380417 - 4096)+ , testProperty "decompose" $ \(FE r) -> do+ gamma2 <- elements [ 95232, 261888 ]+ let (r1, r0) = decomposeZq gamma2 r+ return $ r === toZq (2 * gamma2 * r1) .+ r0 .&&.+ (fromZq r0 <= gamma2 .||. fromZq r0 > 8380417 - gamma2)+ ]+ , testGroup "hints"+ [ testProperty "fromBools . toBools == id " $ do+ a <- arbitraryHints+ return $ Just a === fromBools (toBools a)+ , testProperty "toBools . fromBools == id " $ do+ a <- vector 256+ return $ Just a === (toBools <$> fromBools a)+ , testProperty "useHint . makeHint" $ \(Poly z') (Poly r) -> do+ gamma2 <- elements [ 95232, 261888 ]+ let z = truncateRq' gamma2 z' -- to ensure |z| < gamma2+ return $ useHint gamma2 (makeHint gamma2 z r) r === highBits gamma2 (r .+ z)+ ]+ , testGroup "conversions"+ [ testProperty "simpleBitPack . simpleBitUnpack == id" $ \(D d) -> do+ b <- arbitraryBytes (32 * d)+ return (b === runBytes (simpleBitPack d (simpleBitUnpack d b)))+ , testProperty "simpleBitPack . simpleBitUnpack == id (unaligned)" $ \(D d) -> do+ b <- arbitraryBytes (32 * d)+ return (B.convert b === runUnaligned (simpleBitPack d (simpleBitUnpack d b)))+ , testProperty "simpleBitPack 8" $ \x ->+ B.replicate 256 x === simpleBitPackBytes 8 (BlockN.replicate $ fromIntegral x)+ , testCase "simpleBitPack 1 (zeros)" $+ B.replicate 32 0 @=? simpleBitPackBytes 1 (BlockN.replicate 0)+ , testCase "simpleBitPack 1 (ones)" $+ B.replicate 32 255 @=? simpleBitPackBytes 1 (BlockN.replicate 1)+ , testProperty "simpleBitPack10 . simpleBitUnpack10 == id" $ do+ b <- arbitraryBytes 320+ return (b === runPubBytes (simpleBitPack10 (simpleBitUnpack10 b)))+ , testProperty "simpleBitPack10 . simpleBitUnpack10 == id (unaligned)" $ do+ b <- arbitraryBytes 320+ return (B.convert b === runPubUnaligned (simpleBitPack10 (simpleBitUnpack10 b)))+ , testProperty "bitPackSafe . bitUnpackSafe == id" $ \(D d) -> do+ b <- arbitraryBytes (32 * d)+ return (b === runBytes (bitPackSafe d (bitUnpackSafe d b)))+ , testProperty "bitPackSafe . bitUnpackSafe == id (unaligned)" $ \(D d) -> do+ b <- arbitraryBytes (32 * d)+ return (B.convert b === runUnaligned (bitPackSafe d (bitUnpackSafe d b)))+ , testProperty "bitUnpack . bitPack == id" $ \(FE t) (Poly w') ->+ let w = truncateRq (fromZq t) w'+ m = getNorm (norm w)+ d = 33 - countLeadingZeros m+ in Just w === bitUnpack m d (runBytes $ bitPack m d w)+ , testProperty "bitUnpack . bitPack == id (unaligned)" $ \(FE t) (Poly w') ->+ let w = truncateRq (fromZq t) w'+ m = getNorm (norm w)+ d = 33 - countLeadingZeros m+ in Just w === bitUnpack m d (runUnaligned $ bitPack m d w)+ , testProperty "hintBitUnpack . hintBitPack == id" $ \(Dim k) extra -> do+ h <- arbitraryHintsVector k+ let omega = fromIntegral (extra + countHints h)+ return $ omega < 256 ==>+ Just h === hintBitUnpack omega (runBytes $ hintBitPack omega h)+ , testProperty "hintBitUnpack . hintBitPack == id (unaligned)" $ \(Dim k) extra -> do+ h <- arbitraryHintsVector k+ let omega = fromIntegral (extra + countHints h)+ return $ omega < 256 ==>+ Just h === hintBitUnpack omega (runUnaligned $ hintBitPack omega h)+ ]+ , testGroup "Zq"+ [ testProperty "toZq . fromZq == id " $ \(FE a) ->+ a === toZq (fromZq a)+ , testProperty "fromZq . toZq == id " $ \a ->+ mod a 8380417 === fromZq (toZq a)+ , testCase "field order" $ zero @=? toZq 8380417+ , 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+ ]+ ]+ , testGroup "Mq"+ [ testProperty "toMontgomery . fromMontgomery == id " $ \(ME a) ->+ a === toMontgomery (fromMontgomery a)+ , testProperty "fromMontgomery . toMontgomery == id" $ \(FE a) ->+ a === fromMontgomery (toMontgomery a)+ , testProperty "addition with zero" $ \(ME a) ->+ conjoin [ a === zero .+ a+ , a === a .+ zero+ ]+ , testProperty "addition associative" $ \(ME a) (ME b) (ME c) ->+ a .+ (b .+ c) === (a .+ b) .+ c+ , testProperty "addition commutative" $ \(ME a) (ME b) ->+ a .+ b === b .+ a+ , testProperty "substraction with zero" $ \(ME a) ->+ a === a .- zero+ , testProperty "substraction non-associative" $ \(ME a) (ME b) (ME c) ->+ a .- (b .- c) === (a .- b) .+ c+ , testProperty "substraction anti-commutative" $ \(ME a) (ME b) ->+ a .- b === neg (b .- a)+ , testProperty "negation" $ \(ME a) ->+ neg a === zero .- a+ , testProperty "double negation" $ \(ME a) ->+ a === neg (neg a)+ , testProperty "multiplication with zero" $ \(ME a) ->+ conjoin [ zero === zero .* a+ , zero === a .* zero+ ]+ , testProperty "multiplication with one" $ \(ME a) ->+ conjoin [ a === one .* a+ , a === a .* one+ ]+ , testProperty "multiplication associative" $ \(ME a) (ME b) (ME c) ->+ a .* (b .* c) === (a .* b) .* c+ , testProperty "multiplication commutative" $ \(ME a) (ME b) ->+ a .* b === b .* a+ , testProperty "multiplication distributive" $ \(ME a) (ME b) (ME c) ->+ conjoin [ (a .* b) .+ (a .* c) === a .* (b .+ c)+ , (b .* a) .+ (c .* a) === (b .+ c) .* a+ ]+ ]+ , 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)+ , testCase "norm zero" $ norm (zero :: Rq Sec) @=? 0+ , testProperty "norm positiveness" $ \(Poly a) ->+ (norm a == 0) === (a == zero)+ , testProperty "norm even" $ \(Poly a) ->+ norm (neg a) === norm a+ , testProperty "norm sub-additivity" $ \(Poly a) (Poly b) ->+ norm (a .+ b) <= norm a + norm b+ ]+ , 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))+ , testProperty "dot product commutative" $ \(Dim n) -> do+ (u, v) <- (,) <$> arbitraryVector n <*> arbitraryVector n+ return (u `dot` v === v `dot` u)+ , testProperty "dot product distributive" $ \(Dim n) -> do+ (u, v, w) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n+ return $ conjoin+ [ u `dot` (v .+ w) === (u `dot` v) .+ (u `dot` w)+ , (u .+ v) `dot` w === (u `dot` w) .+ (v `dot` w)+ ]+ ]+ , testGroup "Matrix"+ [ testProperty "mmul distributive left" $ \(Dim n) (Dim m) -> do+ (a, b, u) <- (,,) <$> arbitraryMatrix n m <*> arbitraryMatrix n m <*> arbitraryVector m+ return ((a .+ b) `mmul` u === (a `mmul` u) .+ (b `mmul` u))+ , testProperty "mmul distributive right" $ \(Dim n) (Dim m) -> do+ (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector m+ return (a `mmul` (u .+ v) === (a `mmul` u) .+ (a `mmul` v))+ , testProperty "mmulAdd definition" $ \(Dim n) (Dim m) -> do+ (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector n+ return (mmulAdd a u v == (a `mmul` u) .+ v)+ , testProperty "mmulAdd distributive left" $ \(Dim n) (Dim m) -> do+ (a, b, u, v) <- (,,,) <$> arbitraryMatrix n m <*> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector n+ return (mmulAdd (a .+ b) u v === mmulAdd a u (mmulAdd b u zero) .+ v)+ , testProperty "mmulAdd distributive right" $ \(Dim n) (Dim m) -> do+ (a, u, v, w) <- (,,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector m <*> arbitraryVector n+ return (mmulAdd a (u .+ v) w === mmulAdd a u (mmulAdd a v w))+ ]+#endif+ ]+ ]
+ tests/Util.hs view
@@ -0,0 +1,47 @@+-- |+-- 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++(.::?) :: Object -> Key -> Parser (Maybe ByteString)+o .::? name = (o .:? name) >>= opt fromBase16++opt :: Monad m => (a -> m b) -> Maybe a -> m (Maybe b)+opt f = maybe (return Nothing) (fmap Just . f)++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-DSA 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/79e78ba49d1605baaf9acdbf475304af6ae36a59+CURL=curl++for KEY in keyGen sigGen sigVer; do+ FILENAME="$DESTDIR"/$KEY.json.gz+ URL=https://codeberg.org/ocheron/hs-mldsa/raw/$REF/tests/$KEY.json.gz++ "$CURL" --silent --fail -o "$FILENAME" "$URL" || exit $?+done