{-# 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 :: String -> PoolId
mkPoolId String
str =
  KeyHash StakePool -> PoolId
PoolId
    (KeyHash StakePool -> PoolId)
-> (String -> KeyHash StakePool) -> String -> 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)
-> (String -> VKey StakePool) -> String -> 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)
-> (String -> VerKeyDSIGN DSIGN) -> String -> 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)
-> (String -> SignKeyDSIGN DSIGN) -> String -> 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)
-> (String -> Seed) -> String -> SignKeyDSIGN DSIGN
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Seed
SL.mkSeedFromBytes
    (ByteString -> Seed) -> (String -> ByteString) -> String -> Seed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
forall a. IsString a => String -> a
fromString
    (String -> PoolId) -> String -> PoolId
forall a b. (a -> b) -> a -> b
$ String
paddedStr
 where
  paddedStr :: String
paddedStr
    | String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
str Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
neededBytes = Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
neededBytes String
str
    | Bool
otherwise = String
str String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
neededBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
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 =
  Rational -> LedgerStake
LedgerStake
    (Rational -> LedgerStake)
-> (Int -> Rational) -> Int -> LedgerStake
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Rational
forall a. Real a => a -> Rational
toRational
    (Int -> Rational) -> (Int -> Int) -> Int -> Rational
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 String
alphaNumString
    (privateKey, publicKey) <- genKeyPair
    stake <- genStake
    pure (mkPoolId poolId, (privateKey, publicKey, stake))

  alphaNumString :: Gen String
alphaNumString =
    Int -> Gen Char -> Gen String
forall a. Int -> Gen a -> Gen [a]
vectorOf Int
8 (Gen Char -> Gen String) -> Gen Char -> Gen String
forall a b. (a -> b) -> a -> b
$
      String -> Gen Char
forall a. HasCallStack => [a] -> Gen a
elements (String -> Gen Char) -> String -> Gen Char
forall a b. (a -> b) -> a -> b
$
        [Char
'a' .. Char
'z']
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [Char
'A' .. Char
'Z']
          String -> String -> String
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 -> String) -> (a -> a -> Bool) -> a -> a -> Property
eqWithShowCmp a -> String
showValue a -> a -> Bool
eqValue a
x a
y =
  String -> Bool -> Property
forall prop. Testable prop => String -> prop -> Property
counterexample (a -> String
showValue a
x String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Bool -> String
interpret Bool
res String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
showValue a
y) Bool
res
 where
  res :: Bool
res = a -> a -> Bool
eqValue a
x a
y
  interpret :: Bool -> String
interpret Bool
True = String
" == "
  interpret Bool
False = String
" /= "

-- | 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 -> String
mkBucket Integer
size Integer
val
  | Integer
val Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 =
      String
"<= 0"
  | Bool
otherwise =
      String
"[ " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Integer -> String
forall a. Show a => a -> String
show Integer
lo String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Integer -> String
forall a. Show a => a -> String
show Integer
hi String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" )"
 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 =
  String -> [String] -> Property -> Property
forall prop.
Testable prop =>
String -> [String] -> prop -> Property
tabulate
    String
"Number of pools"
    [Integer -> Integer -> String
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 Rational
stake) =
  String -> [String] -> Property -> Property
forall prop.
Testable prop =>
String -> [String] -> prop -> Property
tabulate
    String
"Pool stake"
    [ if Rational
stake Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
> Rational
0
        then String
"> 0"
        else String
"== 0"
    ]