{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}

-- | Utility functions for the voting committee tests.
module Test.Consensus.Committee.Utils
  ( -- * General utilities
    mkPoolId
  , unfairWFATiebreaker

    -- * QuickCheck generators
  , genEpochNonce
  , genPositiveStake
  , genPools

    -- * Property helpers
  , eqWithShowCmp
  , onError
  , mkBucket

    -- * Tabulation helpers
  , tabulateNumPools
  , tabulatePoolStake
  ) where

import qualified Cardano.Crypto.DSIGN.Class as SL
import qualified Cardano.Crypto.Seed as SL
import Cardano.Ledger.BaseTypes (Nonce (..), mkNonceFromNumber)
import qualified Cardano.Ledger.Core as SL
import qualified Cardano.Ledger.Keys as SL
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.String (IsString (..))
import Ouroboros.Consensus.Committee.Types (LedgerStake (..), PoolId (..))
import Ouroboros.Consensus.Committee.WFA (WFATiebreaker (..))
import Test.QuickCheck
  ( Arbitrary (..)
  , Gen
  , Property
  , choose
  , counterexample
  , elements
  , frequency
  , tabulate
  , vectorOf
  )
import Test.Util.QuickCheck (geometric)

-- * General utilities

-- | Create a pool ID from an arbitrary string of any length.
--
-- NOTE: we are assuming that this function preserves uniqueness.
mkPoolId :: String -> PoolId
mkPoolId :: [Char] -> PoolId
mkPoolId [Char]
str =
  KeyHash StakePool -> PoolId
PoolId
    (KeyHash StakePool -> PoolId)
-> ([Char] -> KeyHash StakePool) -> [Char] -> PoolId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VKey StakePool -> KeyHash StakePool
forall (kd :: KeyRole). VKey kd -> KeyHash kd
SL.hashKey
    (VKey StakePool -> KeyHash StakePool)
-> ([Char] -> VKey StakePool) -> [Char] -> KeyHash StakePool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VerKeyDSIGN DSIGN -> VKey StakePool
forall (kd :: KeyRole). VerKeyDSIGN DSIGN -> VKey kd
SL.VKey
    (VerKeyDSIGN DSIGN -> VKey StakePool)
-> ([Char] -> VerKeyDSIGN DSIGN) -> [Char] -> VKey StakePool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SignKeyDSIGN DSIGN -> VerKeyDSIGN DSIGN
forall v. DSIGNAlgorithm v => SignKeyDSIGN v -> VerKeyDSIGN v
SL.deriveVerKeyDSIGN
    (SignKeyDSIGN DSIGN -> VerKeyDSIGN DSIGN)
-> ([Char] -> SignKeyDSIGN DSIGN) -> [Char] -> VerKeyDSIGN DSIGN
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Seed -> SignKeyDSIGN DSIGN
forall v. DSIGNAlgorithm v => Seed -> SignKeyDSIGN v
SL.genKeyDSIGN
    (Seed -> SignKeyDSIGN DSIGN)
-> ([Char] -> Seed) -> [Char] -> SignKeyDSIGN DSIGN
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Seed
SL.mkSeedFromBytes
    (ByteString -> Seed) -> ([Char] -> ByteString) -> [Char] -> Seed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ByteString
forall a. IsString a => [Char] -> a
fromString
    ([Char] -> PoolId) -> [Char] -> PoolId
forall a b. (a -> b) -> a -> b
$ [Char]
paddedStr
 where
  paddedStr :: [Char]
paddedStr
    | [Char] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Char]
str Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
neededBytes = Int -> [Char] -> [Char]
forall a. Int -> [a] -> [a]
take Int
neededBytes [Char]
str
    | Bool
otherwise = [Char]
str [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> Char -> [Char]
forall a. Int -> a -> [a]
replicate (Int
neededBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Char] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Char]
str) Char
'0'

  neededBytes :: Int
neededBytes = Int
32

-- | An unfair tie-breaker that compares pool IDs lexicographically.
unfairWFATiebreaker :: WFATiebreaker
unfairWFATiebreaker :: WFATiebreaker
unfairWFATiebreaker =
  (PoolId -> PoolId -> Ordering) -> WFATiebreaker
WFATiebreaker PoolId -> PoolId -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

-- * QuickCheck generators

-- | Generate a random nonce for testing purposes.
genEpochNonce :: Gen Nonce
genEpochNonce :: Gen Nonce
genEpochNonce =
  [(Int, Gen Nonce)] -> Gen Nonce
forall a. HasCallStack => [(Int, Gen a)] -> Gen a
frequency
    [ (Int
1, Nonce -> Gen Nonce
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Nonce
NeutralNonce)
    , (Int
9, Word64 -> Nonce
mkNonceFromNumber (Word64 -> Nonce) -> Gen Word64 -> Gen Nonce
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
forall a. Arbitrary a => Gen a
arbitrary)
    ]

-- | Generate a positive stake value using a geometric distribution.
genPositiveStake :: Gen LedgerStake
genPositiveStake :: Gen LedgerStake
genPositiveStake =
  Ratio Integer -> LedgerStake
LedgerStake
    (Ratio Integer -> LedgerStake)
-> (Int -> Ratio Integer) -> Int -> LedgerStake
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Ratio Integer
forall a. Real a => a -> Ratio Integer
toRational
    (Int -> Ratio Integer) -> (Int -> Int) -> Int -> Ratio Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    (Int -> LedgerStake) -> Gen Int -> Gen LedgerStake
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Gen Int
geometric Double
0.25

-- | Generate a non-empty map of pools with crypto keys and varying stakes.
--
-- NOTE: the generator ensures at least one pool has positive stake.
genPools ::
  -- | Maximum number of pools to generate
  Int ->
  -- | Keypair generator
  Gen (privateKey, publicKey) ->
  Gen (Map PoolId (privateKey, publicKey, LedgerStake))
genPools :: forall privateKey publicKey.
Int
-> Gen (privateKey, publicKey)
-> Gen (Map PoolId (privateKey, publicKey, LedgerStake))
genPools Int
maxPools Gen (privateKey, publicKey)
genKeyPair = do
  numPools <-
    (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int
maxPools)
  numPoolsWithZeroStake <-
    choose (0, numPools - 1)
  poolsWithZeroStake <-
    vectorOf
      numPoolsWithZeroStake
      (genOnePool (pure (LedgerStake 0)))
  poolsWithPositiveStake <-
    vectorOf
      (numPools - numPoolsWithZeroStake)
      (genOnePool genPositiveStake)
  pure $
    Map.fromList (poolsWithZeroStake <> poolsWithPositiveStake)
 where
  genOnePool :: Gen c -> Gen (PoolId, (privateKey, publicKey, c))
genOnePool Gen c
genStake = do
    poolId <- Gen [Char]
alphaNumString
    (privateKey, publicKey) <- genKeyPair
    stake <- genStake
    pure (mkPoolId poolId, (privateKey, publicKey, stake))

  alphaNumString :: Gen [Char]
alphaNumString =
    Int -> Gen Char -> Gen [Char]
forall a. Int -> Gen a -> Gen [a]
vectorOf Int
8 (Gen Char -> Gen [Char]) -> Gen Char -> Gen [Char]
forall a b. (a -> b) -> a -> b
$
      [Char] -> Gen Char
forall a. HasCallStack => [a] -> Gen a
elements ([Char] -> Gen Char) -> [Char] -> Gen Char
forall a b. (a -> b) -> a -> b
$
        [Char
'a' .. Char
'z']
          [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char
'A' .. Char
'Z']
          [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char
'0' .. Char
'9']

-- * Property helpers

-- | Check equality using a custom comparison and show function.
eqWithShowCmp ::
  -- | Custom show function
  (a -> String) ->
  -- | Custom equality function
  (a -> a -> Bool) ->
  -- | First value
  a ->
  -- | Second value
  a ->
  Property
eqWithShowCmp :: forall a. (a -> [Char]) -> (a -> a -> Bool) -> a -> a -> Property
eqWithShowCmp a -> [Char]
showValue a -> a -> Bool
eqValue a
x a
y =
  [Char] -> Bool -> Property
forall prop. Testable prop => [Char] -> prop -> Property
counterexample (a -> [Char]
showValue a
x [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Bool -> [Char]
interpret Bool
res [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> a -> [Char]
showValue a
y) Bool
res
 where
  res :: Bool
res = a -> a -> Bool
eqValue a
x a
y
  interpret :: Bool -> [Char]
interpret Bool
True = [Char]
" == "
  interpret Bool
False = [Char]
" /= "

-- | Handle 'Either' errors by converting them to values.
onError :: Either err a -> (err -> a) -> a
onError :: forall err a. Either err a -> (err -> a) -> a
onError Either err a
action err -> a
onLeft =
  case Either err a
action of
    Left err
err -> err -> a
onLeft err
err
    Right a
val -> a
val

-- | Create a bucketized label for tabulation.
mkBucket ::
  -- | Bucket size
  Integer ->
  -- | Value to bucket
  Integer ->
  String
mkBucket :: Integer -> Integer -> [Char]
mkBucket Integer
size Integer
val
  | Integer
val Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 =
      [Char]
"<= 0"
  | Bool
otherwise =
      [Char]
"[ " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
lo [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
", " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
hi [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" )"
 where
  lo :: Integer
lo = (Integer
val Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
size) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
size
  hi :: Integer
hi = Integer
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
size

-- * Tabulation helpers

-- | Tabulate the number of pools in a test run.
tabulateNumPools ::
  Map
    PoolId
    ( privateKey
    , publicKey
    , LedgerStake
    ) ->
  Property ->
  Property
tabulateNumPools :: forall privateKey publicKey.
Map PoolId (privateKey, publicKey, LedgerStake)
-> Property -> Property
tabulateNumPools Map PoolId (privateKey, publicKey, LedgerStake)
pools =
  [Char] -> [[Char]] -> Property -> Property
forall prop.
Testable prop =>
[Char] -> [[Char]] -> prop -> Property
tabulate
    [Char]
"Number of pools"
    [Integer -> Integer -> [Char]
mkBucket Integer
100 (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Map PoolId (privateKey, publicKey, LedgerStake) -> Int
forall k a. Map k a -> Int
Map.size Map PoolId (privateKey, publicKey, LedgerStake)
pools))]

-- | Tabulate whether a pool has positive or zero stake.
tabulatePoolStake ::
  LedgerStake ->
  Property ->
  Property
tabulatePoolStake :: LedgerStake -> Property -> Property
tabulatePoolStake (LedgerStake Ratio Integer
stake) =
  [Char] -> [[Char]] -> Property -> Property
forall prop.
Testable prop =>
[Char] -> [[Char]] -> prop -> Property
tabulate
    [Char]
"Pool stake"
    [ if Ratio Integer
stake Ratio Integer -> Ratio Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Ratio Integer
0
        then [Char]
"> 0"
        else [Char]
"== 0"
    ]