mldsa-0.1.0.0: benchs/Bench.hs

{-# 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
    ]