-- | Utility functions for the WFALS model tests.
module Test.Consensus.Committee.WFALS.Model.Utils
  ( -- * Generators
    genStake
  , genUniqueVoterIds
  , genStakeDistr

    -- * Property helpers
  , forAllValidStakeDistrAndNumSeats
  , forAllPossiblyInvalidStakeDistrAndNumSeats
  , tabulateTargetNumSeats
  , tabulateStakeDistrSize
  , tabulatePersistentToNonPersistentRatio
  ) where

import qualified Data.Map.Strict as Map
import Data.Ratio ((%))
import qualified Data.Set as Set
import Data.Word (Word64)
import Test.Consensus.Committee.Utils (mkBucket)
import Test.Consensus.Committee.WFALS.Model
  ( IsStake (..)
  , NumSeats
  , Stake (..)
  , StakeDistr
  , StakeRole (..)
  , StakeType (..)
  , VoterId
  )
import Test.QuickCheck
  ( Arbitrary (..)
  , Gen
  , Positive (..)
  , Property
  , choose
  , forAll
  , sized
  , tabulate
  , vectorOf
  )

-- * Generators

-- | Generate a random stake as a Rational number
genStake :: Gen (Stake Ledger Global)
genStake :: Gen (Stake Ledger Global)
genStake = do
  Positive num <- Gen (Positive Integer)
forall a. Arbitrary a => Gen a
arbitrary
  Positive den <- arbitrary
  pure (StakeLedgerGlobal (num % den))

-- | Generate a non-empty list of unique voter IDs of a given size.
genUniqueVoterIds :: Int -> Gen [VoterId]
genUniqueVoterIds :: Int -> Gen [VoterId]
genUniqueVoterIds Int
size =
  (Word64 -> VoterId) -> [Word64] -> [VoterId]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word64 -> VoterId
forall a. Show a => a -> VoterId
show ([Word64] -> [VoterId]) -> Gen [Word64] -> Gen [VoterId]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Set Word64 -> Gen [Word64]
forall {t}. (Eq t, Num t) => t -> Set Word64 -> Gen [Word64]
go Int
size Set Word64
forall a. Set a
Set.empty
 where
  go :: t -> Set Word64 -> Gen [Word64]
go t
0 Set Word64
acc =
    [Word64] -> Gen [Word64]
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set Word64 -> [Word64]
forall a. Set a -> [a]
Set.toList Set Word64
acc)
  go t
k Set Word64
acc = do
    voterId <- (Word64, Word64) -> Gen Word64
forall a. Random a => (a, a) -> Gen a
choose (Word64
0, Word64
forall a. Bounded a => a
maxBound :: Word64)
    if voterId `Set.member` acc
      then
        go k acc
      else do
        rest <- go (k - 1) (Set.insert voterId acc)
        pure (voterId : rest)

-- | Generate a non-empty random stake distribution of a given size
--
-- NOTE: the size is shifted up by one to ensure non-emptiness.
genStakeDistr :: Int -> Gen (StakeDistr Ledger Global)
genStakeDistr :: Int -> Gen (StakeDistr Ledger Global)
genStakeDistr Int
size = do
  ids <- Int -> Gen [VoterId]
genUniqueVoterIds (Int -> Int
forall a. Enum a => a -> a
succ Int
size)
  stakes <- vectorOf (succ size) genStake
  pure (Map.fromList (zip ids stakes))

-- * Property helpers

-- | Helper to generate stake distributions along with a number of seats that
-- lies within the acceptable range [1, #{nodes with positive stake}]
forAllValidStakeDistrAndNumSeats ::
  (StakeDistr Ledger Global -> NumSeats Global -> Property) ->
  Property
forAllValidStakeDistrAndNumSeats :: (StakeDistr Ledger Global -> NumSeats Global -> Property)
-> Property
forAllValidStakeDistrAndNumSeats StakeDistr Ledger Global -> NumSeats Global -> Property
p =
  Gen (StakeDistr Ledger Global)
-> (StakeDistr Ledger Global -> Property) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll ((Int -> Gen (StakeDistr Ledger Global))
-> Gen (StakeDistr Ledger Global)
forall a. (Int -> Gen a) -> Gen a
sized Int -> Gen (StakeDistr Ledger Global)
genStakeDistr) ((StakeDistr Ledger Global -> Property) -> Property)
-> (StakeDistr Ledger Global -> Property) -> Property
forall a b. (a -> b) -> a -> b
$ \StakeDistr Ledger Global
stakeDistr -> do
    let numPositiveStakeNodes :: Integer
numPositiveStakeNodes =
          Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
            (Int -> Integer)
-> (StakeDistr Ledger Global -> Int)
-> StakeDistr Ledger Global
-> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Stake Ledger Global] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
            ([Stake Ledger Global] -> Int)
-> (StakeDistr Ledger Global -> [Stake Ledger Global])
-> StakeDistr Ledger Global
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stake Ledger Global -> Bool)
-> [Stake Ledger Global] -> [Stake Ledger Global]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
> Rational
0) (Rational -> Bool)
-> (Stake Ledger Global -> Rational) -> Stake Ledger Global -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stake Ledger Global -> Rational
forall stake. IsStake stake => stake -> Rational
stakeToRational)
            ([Stake Ledger Global] -> [Stake Ledger Global])
-> (StakeDistr Ledger Global -> [Stake Ledger Global])
-> StakeDistr Ledger Global
-> [Stake Ledger Global]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StakeDistr Ledger Global -> [Stake Ledger Global]
forall k a. Map k a -> [a]
Map.elems
            (StakeDistr Ledger Global -> Integer)
-> StakeDistr Ledger Global -> Integer
forall a b. (a -> b) -> a -> b
$ StakeDistr Ledger Global
stakeDistr
    let genNumSeats :: Gen Natural
genNumSeats =
          Integer -> Natural
forall a. Num a => Integer -> a
fromInteger (Integer -> Natural) -> Gen Integer -> Gen Natural
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer, Integer) -> Gen Integer
forall a. Random a => (a, a) -> Gen a
choose (Integer
1, Integer
numPositiveStakeNodes)
    Gen Natural -> (Natural -> Property) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll Gen Natural
genNumSeats ((Natural -> Property) -> Property)
-> (Natural -> Property) -> Property
forall a b. (a -> b) -> a -> b
$ \Natural
numSeats ->
      StakeDistr Ledger Global -> NumSeats Global -> Property
p StakeDistr Ledger Global
stakeDistr Natural
NumSeats Global
numSeats

-- | Helper to generate stake distributions along with a number of seats that
-- could possibly be invalid in different ways, e.g., too many expected seats
-- or an empty stake distribution.
forAllPossiblyInvalidStakeDistrAndNumSeats ::
  (StakeDistr Ledger Global -> NumSeats Global -> Property) ->
  Property
forAllPossiblyInvalidStakeDistrAndNumSeats :: (StakeDistr Ledger Global -> NumSeats Global -> Property)
-> Property
forAllPossiblyInvalidStakeDistrAndNumSeats StakeDistr Ledger Global -> NumSeats Global -> Property
p =
  Gen (StakeDistr Ledger Global)
-> (StakeDistr Ledger Global -> Property) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll ((Int -> Gen (StakeDistr Ledger Global))
-> Gen (StakeDistr Ledger Global)
forall a. (Int -> Gen a) -> Gen a
sized Int -> Gen (StakeDistr Ledger Global)
genStakeDistr) ((StakeDistr Ledger Global -> Property) -> Property)
-> (StakeDistr Ledger Global -> Property) -> Property
forall a b. (a -> b) -> a -> b
$ \StakeDistr Ledger Global
stakeDistr -> do
    let numNodes :: Integer
numNodes =
          Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (StakeDistr Ledger Global -> Int
forall k a. Map k a -> Int
Map.size StakeDistr Ledger Global
stakeDistr)
    let genNumSeats :: Gen Natural
genNumSeats =
          Integer -> Natural
forall a. Num a => Integer -> a
fromInteger (Integer -> Natural) -> Gen Integer -> Gen Natural
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer, Integer) -> Gen Integer
forall a. Random a => (a, a) -> Gen a
choose (Integer
1, Integer
numNodes Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
    Gen Natural -> (Natural -> Property) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll Gen Natural
genNumSeats ((Natural -> Property) -> Property)
-> (Natural -> Property) -> Property
forall a b. (a -> b) -> a -> b
$ \Natural
numSeats ->
      StakeDistr Ledger Global -> NumSeats Global -> Property
p StakeDistr Ledger Global
stakeDistr Natural
NumSeats Global
numSeats

-- | Tabulate the target number of seats
tabulateTargetNumSeats :: NumSeats Global -> Property -> Property
tabulateTargetNumSeats :: NumSeats Global -> Property -> Property
tabulateTargetNumSeats NumSeats Global
numSeats =
  VoterId -> [VoterId] -> Property -> Property
forall prop.
Testable prop =>
VoterId -> [VoterId] -> prop -> Property
tabulate
    VoterId
"Target number of seats"
    [Integer -> Integer -> VoterId
mkBucket Integer
10 (Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
NumSeats Global
numSeats)]

-- | Tabulate the size of a stake distribution
tabulateStakeDistrSize :: StakeDistr Ledger Global -> Property -> Property
tabulateStakeDistrSize :: StakeDistr Ledger Global -> Property -> Property
tabulateStakeDistrSize StakeDistr Ledger Global
stakeDistr =
  VoterId -> [VoterId] -> Property -> Property
forall prop.
Testable prop =>
VoterId -> [VoterId] -> prop -> Property
tabulate
    VoterId
"Stake distribution size"
    [Integer -> Integer -> VoterId
mkBucket Integer
10 (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (StakeDistr Ledger Global -> Int
forall k a. Map k a -> Int
Map.size StakeDistr Ledger Global
stakeDistr))]

-- | Tabulate the ratio of persistent to non-persistent seats as the percentage
-- of persistent seats out of the total number of seats.
tabulatePersistentToNonPersistentRatio ::
  NumSeats Global ->
  NumSeats Global ->
  Property ->
  Property
tabulatePersistentToNonPersistentRatio :: NumSeats Global -> NumSeats Global -> Property -> Property
tabulatePersistentToNonPersistentRatio
  NumSeats Global
numPersistentSeats
  NumSeats Global
numNonPersistentSeats =
    VoterId -> [VoterId] -> Property -> Property
forall prop.
Testable prop =>
VoterId -> [VoterId] -> prop -> Property
tabulate
      VoterId
"Persistent to non-persistent seat ratio (%)"
      [Integer -> Integer -> VoterId
mkBucket Integer
10 Integer
ratio]
   where
    ratio :: Integer
ratio = Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
NumSeats Global
numPersistentSeats Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
100 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
totalSeats
    totalSeats :: Natural
totalSeats = Natural
NumSeats Global
numPersistentSeats Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
NumSeats Global
numNonPersistentSeats