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