packages feed

mlkem 0.2.1.0 → 0.2.2.0

raw patch · 11 files changed

+149/−73 lines, 11 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

CHANGELOG.md view
@@ -1,5 +1,17 @@ # Changelog for `mlkem` +## 0.2.2.0 - 2026-06-28++* Clearing buffers containing secrets is more resilient and takes into account+  the possibility of asynchronous exceptions++* Improve optimizations a bit more.  Sometimes vector rewrite rules were not+  triggered early enough to fuse polynomial allocations at element level.++* Modular reduction and constant-time selection now use an arithmetic shift+  instead of a logical shift.  This makes no difference with LLVM but saves one+  instruction with the NCG.+ ## 0.2.1.0 - 2026-05-28  * Function `generateOpen` is added to return not only the key pair but also the
mlkem.cabal view
@@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack  name:           mlkem-version:        0.2.1.0+version:        0.2.2.0 synopsis:       Module-Lattice-based Key-Encapsulation Mechanism description:    Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM) implemented in                 Haskell.
src/Auxiliary.hs view
@@ -38,6 +38,7 @@ import Control.Monad.ST  import Data.Bits+import Data.Int import Data.Proxy import Data.Word @@ -90,12 +91,16 @@     (b `unsafeShiftR` 1 .&. 1) `unsafeShiftL` 5 .|.     (b .&. 1) `unsafeShiftL` 6 +unsafeShiftIR :: Word16 -> Int -> Word16+unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int16) `unsafeShiftR` s)+{-# INLINE unsafeShiftIR #-}+ -- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞 reduceSimple :: Word16 -> Word16 reduceSimple x = (mask .&. x) .|. (complement mask .&. subtracted)   where     subtracted = x - q16-    mask = negate (subtracted `unsafeShiftR` 15)+    mask = subtracted `unsafeShiftIR` 15 {-# INLINE reduceSimple #-}  -- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞² + 𝑞@@ -403,12 +408,12 @@  -- Compress a polynomial with 𝑑 < 12 rcompress :: Classified marking => Int -> Rq marking -> BlockN marking N Word16-rcompress !d (Rq a) = BlockN.mapEqPrimSize (compress d) a+rcompress d (Rq a) = BlockN.seq d $ BlockN.mapEqPrimSize (compress d) a {-# INLINE rcompress #-}  -- Decompress a polynomial with 𝑑 < 12 rdecompress :: Classified marking => Int -> BlockN marking N Word16 -> Rq marking-rdecompress !d = Rq . BlockN.mapEqPrimSize (decompress d)+rdecompress d = Rq . BlockN.seq d . BlockN.mapEqPrimSize (decompress d) {-# INLINE rdecompress #-}  -- Generates a pseudorandom element of T𝑞 from a seed and two indices
src/Block.hs view
@@ -11,7 +11,7 @@ {-# LANGUAGE CPP #-} module Block     ( Block, MutableBlock, blockIndex, blockRead, blockWrite-    , foldZipWith, Block.length, mutableContents+    , foldZipWith, Block.length, mutableContents, getMutableLength     , Block.new, Block.newPinned, Block.thaw, thawPinned     , Block.unsafeCast, unsafeCastMut, Block.unsafeFreeze, Block.unsafeThaw #ifdef ML_KEM_TESTING@@ -67,6 +67,9 @@  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
src/BlockN.hs view
@@ -13,7 +13,7 @@ {-# LANGUAGE ScopedTypeVariables #-} module BlockN     ( BlockN, MutableBlockN, index, iterModify, mapEqPrimSize-    , BlockN.read, runNew, runThaw, runFold, BlockN.unsafeCast+    , BlockN.read, runNew, runThaw, runFold, BlockN.seq, BlockN.unsafeCast     , BlockN.write, BlockN.zipWith #ifdef ML_KEM_TESTING     , create, BlockN.fromList, BlockN.replicate, BlockN.toList@@ -188,6 +188,10 @@ 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 @@ -216,7 +220,7 @@ iterMapIxContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (Int -> a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b) iterMapIxContext f = mapContext m   where m = MapF { mapUpdate = unsafeMapIx f-                 , mapInit = \x -> newF >>= \mb -> iterSet (g x) mb >> return mb+                 , 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 #-}
src/Crypto.hs view
@@ -20,7 +20,7 @@ import Crypto.Hash.Algorithms import Crypto.Hash.IO -import Control.Exception (assert)+import Control.Exception (assert, mask_) import Control.Monad import Control.Monad.ST @@ -47,6 +47,10 @@ 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_KEM_TESTING@@ -74,7 +78,7 @@ eqW a b = isZeroW (a `xor` b)   where     isZeroW x = BoolW $ msbW (complement x .&. (x - 1))-    msbW x = negate (x `unsafeShiftR` (bitsW - 1))+    msbW x = x `unsafeShiftIR` (bitsW - 1)  assertMultW :: Int -> a -> a assertMultW n = assert (n .&. mask == 0)@@ -222,9 +226,9 @@ hashToBlock = BlockDigest . Builder.runToBlock . hashWith (undefined :: a)  hashWith :: forall marking a ba. (HashAlgorithm a, ByteArrayAccess ba) => a -> ba -> Builder marking-hashWith a ba = Builder.unsafeCreate (hashDigestSize a) $ \dig -> do-    ctx <- hashMutableInit-    hashMutableUpdate (ctx :: MutableContext a) ba-    B.withByteArray ctx $ \pctx -> do-        hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig-        fillBytes pctx 0 (B.length ctx)+hashWith a ba = Builder.unsafeCreate (hashDigestSize a) $ \dig ->+    hashMutableInit >>= \ctx -> mask_ $ do+        hashMutableUpdate (ctx :: MutableContext a) ba+        B.withByteArray ctx $ \pctx -> do+            hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig+            fillBytes pctx 0 (B.length ctx)
src/K_PKE.hs view
@@ -100,11 +100,12 @@     view384 (Offset i) = B.view input (384 * i) 384  createMatrix :: KnownNat k => Bytes -> Vector k (Vector k (Tq Pub))-createMatrix !rho = Matrix.create $ \(Offset i) (Offset j) ->+createMatrix rho = Vector.seq rho $ Matrix.create $ \(Offset i) (Offset j) ->     Aux.sampleNTT rho (fromIntegral j) (fromIntegral i)  createVector :: (KnownNat k, ByteArrayAccess s) => Word -> s -> Int -> Vector k (Rq Sec)-createVector !eta !s !j = Vector.create $ \(Offset i) -> sample eta s (i + j)+createVector eta s j = Vector.seq eta $ Vector.seq s $ Vector.seq j $+    Vector.create $ \(Offset i) -> sample eta s (i + j) {-# INLINE createVector #-}  sample :: ByteArrayAccess s => Word -> s -> Int -> Rq Sec
src/Matrix.hs view
@@ -6,7 +6,6 @@ -- A matrix here is simply a vector of vectors.  The module also implements -- two utility functions 'mulw' and 'muly' that multiply a matrix and a vector. ---{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} module Matrix     ( create, mulw, muly@@ -27,13 +26,16 @@ index :: Vector m (Vector n ty) -> Offset ty -> Offset (Vector n ty) -> ty index a i j = Vector.index (Vector.index a j) i -mulw :: (KnownNat n, BiMulAdd b a) => Vector m (Vector n b) -> Vector m a -> Vector n a -> Vector n a-mulw a !u !b = Vector.create $ \(Offset i) ->-    Vector.biMulFoldIndexWith (\(Offset j) vu -> (index a (Offset i) (Offset j), vu)) (Vector.index b (Offset i)) u+mulw :: BiMulAdd b a => Vector m (Vector n b) -> Vector m a -> Vector n a -> Vector n a+mulw 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 mulw #-}  muly :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a-muly a !u = fmap (`Vector.dot` u) a+muly a u = Vector.seq u (fmap (`Vector.dot` u) a) {-# INLINE muly #-}  #ifdef ML_KEM_TESTING
src/ScrubbedBlock.hs view
@@ -7,6 +7,17 @@ -- 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@@ -21,6 +32,8 @@  import Data.Word +import Unsafe.Coerce+ import Base import Block (Block, MutableBlock) import qualified Block@@ -41,13 +54,13 @@ length (ScrubbedBlock b) = Block.length b  new :: (PrimType ty, PrimMonad prim) => CountOf ty -> prim (MutableBlock ty (PrimState prim))-new = Block.newPinned  -- always pinned+new n = Block.newPinned n >>= scrubbed  -- always pinned  thaw :: PrimMonad m => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m))-thaw (ScrubbedBlock b) = Block.thawPinned b  -- always pinned+thaw (ScrubbedBlock b) = Block.thawPinned b >>= scrubbed  -- always pinned  unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty)-unsafeFreeze mb = Block.unsafeFreeze mb >>= scrubbed+unsafeFreeze mb = checkPinned <$> Block.unsafeFreeze mb   {- internal -}@@ -55,21 +68,26 @@ assertPinned :: Block ty -> a -> a assertPinned mb = assert (Block.isPrimArrayPinned mb) -scrubbed :: PrimMonad prim => Block ty -> prim (ScrubbedBlock ty)-scrubbed b = assertPinned b $ unsafePrimFromIO $-    scheduleBlockScrubbing b >> return (ScrubbedBlock b)+checkPinned :: Block ty -> ScrubbedBlock ty+checkPinned b = assertPinned b (ScrubbedBlock b) -scheduleBlockScrubbing :: Block ty -> IO ()-scheduleBlockScrubbing b = addBlockFinalizer b (scrub $ Block.unsafeCast 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 :: Block Word8 -> IO ()-scrub b = Block.unsafeThaw b >>= erase len-  where CountOf len = Block.length b+scrub :: MutableBlock Word8 RealWorld -> IO ()+scrub b = Block.getMutableLength b >>= \(CountOf len) -> erase len b -addBlockFinalizer :: Block ty -> IO () -> IO ()-addBlockFinalizer (Block.PrimArray barr) (IO finalizer) = IO $ \s ->-   case mkWeak# barr () finalizer s of { (# s1, _ #) -> (# s1, () #) }+addBlockFinalizer :: MutableBlock ty s -> IO () -> IO ()+addBlockFinalizer (Block.MutablePrimArray mbarr) (IO finalizer) = IO $ \s ->+   case mkWeak# mbarr () finalizer s of { (# s1, _ #) -> (# s1, () #) }  erase :: Int -> MutableBlock Word8 RealWorld -> IO () erase (I# len) (Block.MutablePrimArray mbarr) = IO $ \s1 ->
src/Vector.hs view
@@ -14,7 +14,7 @@ module Vector     ( Vector, Vector.concatMap, Vector.dot     , Vector.fold1ZipWith, Vector.biMulFoldIndexWith, Vector.toNormalForm-    , Vector.create, Vector.index+    , Vector.create, Vector.index, mapIx, Vector.seq #ifdef ML_KEM_TESTING     , Vector.replicateM #endif@@ -72,8 +72,9 @@ arrayLength :: Array ty -> CountOf ty arrayLength = CountOf . sizeofSmallArray -arrayMap :: (a -> b) -> Array a -> Array b-arrayMap f a = arrayCreate (CountOf sz) $ \(Offset i) -> f $ arrayIndex a (Offset i)+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))@@ -81,26 +82,22 @@   where placeholder = error "arrayNew: unexpected evaluation"  create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a-create = genericCreate-{-# INLINE create #-}--genericCreate :: forall n a a'. KnownNat n => (Offset a' -> a) -> Vector n a-genericCreate f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i))+create f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i))   where !sz = fromIntegral $ natVal (Proxy :: Proxy n)-{-# INLINE [0] genericCreate #-}--genericCreateZipLeft :: KnownNat n => (a -> b -> c) -> (Offset a' -> a) -> Vector n b -> Vector n c-genericCreateZipLeft f g a = genericCreate $ \off@(Offset i) -> f (g off) (index a (Offset i))-{-# INLINE [0] genericCreateZipLeft #-}+{-# INLINE [2] create #-} -genericCreateZipRight :: KnownNat n => (a -> b -> c) -> (Offset b' -> b) -> Vector n a -> Vector n c-genericCreateZipRight f g a = genericCreate $ \off@(Offset i) -> f (index a (Offset i)) (g off)-{-# INLINE [0] genericCreateZipRight #-}+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 = Vector <$> arrayMap f . unVector-{-# INLINE [0] mapVector #-}+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_KEM_TESTING arrayIndex a off@(Offset i) =@@ -124,13 +121,14 @@ mapToList f (Vector a) = Prelude.map (f . arrayIndex a . Offset) (offsets sa)   where CountOf sa = arrayLength a +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 (Vector a) (Vector !b) = Vector $-    arrayCreate (CountOf sa) $ \(Offset i) ->-        f (arrayIndex a (Offset i)) (arrayIndex b (Offset i))-  where-    CountOf sa = arrayLength a-{-# INLINE [0] zipWith #-}+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) =@@ -158,18 +156,47 @@ {-# INLINE dot #-}  toNormalForm :: NFData a => Vector n a -> ()-toNormalForm = foldl' (\acc x -> acc `seq` rnf x) () . unVector+toNormalForm = foldl' (\acc x -> acc `Prelude.seq` rnf x) () . unVector ++-- 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" [~0] forall f g a. mapVector f (mapVector g a) = mapVector (f . g) a-"mapVector/genericCreate" [~0] forall f g. mapVector f (genericCreate g) = genericCreate (f . g)-"zipWith/genericCreate left" [~0] forall f g b. Vector.zipWith f (genericCreate g) b = genericCreateZipLeft f g b-"zipWith/genericCreate right" [~0] forall f g a. Vector.zipWith f a (genericCreate g) = genericCreateZipRight f g a-"genericCreateZipLeft/mapVector" [~0] forall f g h b. genericCreateZipLeft f g (mapVector h b) = genericCreateZipLeft (\aa bb -> f aa (h bb)) g b-"genericCreateZipRight/mapVector" [~0] forall f g h a. genericCreateZipRight f g (mapVector h a) = genericCreateZipRight (f . h) g a-"genericCreateZipLeft/genericCreate" [~0] forall f g h. genericCreateZipLeft f g (genericCreate h) = genericCreate $ \(Offset i) -> f (g (Offset i)) (h (Offset i))-"genericCreateZipRight/genericCreate" [~0] forall f g h. genericCreateZipRight f g (genericCreate h) = genericCreate $ \(Offset i) -> f (h (Offset i)) (g (Offset i))-"zipWith/mapVector left" [~0] forall f g a. Vector.zipWith f (mapVector g a) = Vector.zipWith (f . g) a-"zipWith/mapVector right" [~0] forall f g a b. Vector.zipWith f a (mapVector g b) = Vector.zipWith (\aa bb -> f aa (g bb)) a b-"mapVector/zipWith" [~0] forall f g a b. mapVector f (Vector.zipWith g a b) = Vector.zipWith (\aa bb -> f (g aa bb)) a b+"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/get-vectors.sh view
@@ -9,5 +9,5 @@     FILENAME="$DESTDIR"/$KEY.json.gz     URL=https://codeberg.org/ocheron/hs-mlkem/raw/$REF/tests/$KEY.json.gz -    "$CURL" --silent -o "$FILENAME" "$URL" || exit $?+    "$CURL" --silent --fail -o "$FILENAME" "$URL" || exit $? done