{-# LANGUAGE TupleSections #-}

module Test.Cardano.Chain.Slotting.Gen (
  genEpochNumber,
  genSlotNumber,
  genEpochSlots,
  genWithEpochSlots,
  genSlotCount,
  genEpochAndSlotCount,
  genConsistentEpochAndSlotCountEpochSlots,
  feedPMEpochSlots,
) where

import Cardano.Chain.Slotting (
  EpochAndSlotCount (..),
  EpochNumber (..),
  EpochSlots (..),
  SlotCount (..),
  SlotNumber (..),
  WithEpochSlots (WithEpochSlots),
 )
import Cardano.Crypto (ProtocolMagicId)
import Cardano.Prelude
import Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.Cardano.Crypto.Gen (genProtocolMagicId)

genEpochNumber :: Gen EpochNumber
genEpochNumber :: Gen EpochNumber
genEpochNumber = Word64 -> EpochNumber
EpochNumber (Word64 -> EpochNumber) -> GenT Identity Word64 -> Gen EpochNumber
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word64 -> GenT Identity Word64
forall (m :: * -> *). MonadGen m => Range Word64 -> m Word64
Gen.word64 Range Word64
forall a. (Bounded a, Num a) => Range a
Range.constantBounded

genSlotNumber :: Gen SlotNumber
genSlotNumber :: Gen SlotNumber
genSlotNumber = Word64 -> SlotNumber
SlotNumber (Word64 -> SlotNumber) -> GenT Identity Word64 -> Gen SlotNumber
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word64 -> GenT Identity Word64
forall (m :: * -> *). MonadGen m => Range Word64 -> m Word64
Gen.word64 Range Word64
forall a. (Bounded a, Num a) => Range a
Range.constantBounded

-- | Generator for slots-per-epoch. This will generate a positive number of
-- slots per-epoch, and it will have an upper bound of @maxBound :: Word16 =
-- 2^16@. The reason for this upper bound is that when converting a slot number
-- (which is an absolute value) to a pair of epoch and slot-count, this
-- slot-count, which represents a local index of a slot within the epoch and is
-- represented using a 'Word16', is calculated taking the reminder of dividing
-- the slot number by the number of slots-per-epoch ('EpochSlots'). So if the
-- generated epoch would be greater than @2^16@ we couldn't guarantee that the
-- local-index would fit inside its representation.
genEpochSlots :: Gen EpochSlots
genEpochSlots :: Gen EpochSlots
genEpochSlots =
  Word64 -> EpochSlots
EpochSlots (Word64 -> EpochSlots)
-> (Word16 -> Word64) -> Word16 -> EpochSlots
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word16 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16 -> EpochSlots) -> GenT Identity Word16 -> Gen EpochSlots
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word16 -> GenT Identity Word16
forall (m :: * -> *). MonadGen m => Range Word16 -> m Word16
Gen.word16 (Word16 -> Word16 -> Range Word16
forall a. a -> a -> Range a
Range.constant Word16
1 Word16
forall a. Bounded a => a
maxBound)

-- | Generate a value wrapped on a 'WithEpochSlots' context, using the given
-- generator functions, and its arguments
genWithEpochSlots ::
  (ProtocolMagicId -> EpochSlots -> Gen a) ->
  ProtocolMagicId ->
  EpochSlots ->
  Gen (WithEpochSlots a)
genWithEpochSlots :: forall a.
(ProtocolMagicId -> EpochSlots -> Gen a)
-> ProtocolMagicId -> EpochSlots -> Gen (WithEpochSlots a)
genWithEpochSlots ProtocolMagicId -> EpochSlots -> Gen a
gen ProtocolMagicId
pm EpochSlots
es = EpochSlots -> a -> WithEpochSlots a
forall a. EpochSlots -> a -> WithEpochSlots a
WithEpochSlots EpochSlots
es (a -> WithEpochSlots a)
-> Gen a -> GenT Identity (WithEpochSlots a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProtocolMagicId -> EpochSlots -> Gen a
gen ProtocolMagicId
pm EpochSlots
es

genSlotCount :: Gen SlotCount
genSlotCount :: Gen SlotCount
genSlotCount = Word64 -> SlotCount
SlotCount (Word64 -> SlotCount) -> GenT Identity Word64 -> Gen SlotCount
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word64 -> GenT Identity Word64
forall (m :: * -> *). MonadGen m => Range Word64 -> m Word64
Gen.word64 Range Word64
forall a. (Bounded a, Num a) => Range a
Range.constantBounded

genEpochAndSlotCount :: EpochSlots -> Gen EpochAndSlotCount
genEpochAndSlotCount :: EpochSlots -> Gen EpochAndSlotCount
genEpochAndSlotCount EpochSlots
epochSlots =
  EpochNumber -> SlotCount -> EpochAndSlotCount
EpochAndSlotCount (EpochNumber -> SlotCount -> EpochAndSlotCount)
-> Gen EpochNumber
-> GenT Identity (SlotCount -> EpochAndSlotCount)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen EpochNumber
genEpochNumber GenT Identity (SlotCount -> EpochAndSlotCount)
-> Gen SlotCount -> Gen EpochAndSlotCount
forall a b.
GenT Identity (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> EpochSlots -> Gen SlotCount
genEpochSlotCount EpochSlots
epochSlots

-- | Generate a 'SlotCount' constrained by the number of 'EpochSlots'
genEpochSlotCount :: EpochSlots -> Gen SlotCount
genEpochSlotCount :: EpochSlots -> Gen SlotCount
genEpochSlotCount EpochSlots
epochSlots =
  Word64 -> SlotCount
SlotCount (Word64 -> SlotCount) -> GenT Identity Word64 -> Gen SlotCount
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word64 -> GenT Identity Word64
forall (m :: * -> *). MonadGen m => Range Word64 -> m Word64
Gen.word64 (Word64 -> Word64 -> Range Word64
forall a. Integral a => a -> a -> Range a
Range.linear Word64
0 (EpochSlots -> Word64
unEpochSlots EpochSlots
epochSlots Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1))

-- Generates a `EpochAndSlotCount` and a `EpochSlots` that does not exceed
-- the `Word64` maximum boundary of `flattenEpochAndSlotCount` when flattened.
genConsistentEpochAndSlotCountEpochSlots :: Gen (EpochAndSlotCount, EpochSlots)
genConsistentEpochAndSlotCountEpochSlots :: Gen (EpochAndSlotCount, EpochSlots)
genConsistentEpochAndSlotCountEpochSlots = do
  EpochSlots
epochSlots <- Gen EpochSlots
genEpochSlots
  (EpochAndSlotCount -> (EpochAndSlotCount, EpochSlots))
-> Gen EpochAndSlotCount -> Gen (EpochAndSlotCount, EpochSlots)
forall a b. (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,EpochSlots
epochSlots)
    (Gen EpochAndSlotCount -> Gen (EpochAndSlotCount, EpochSlots))
-> Gen EpochAndSlotCount -> Gen (EpochAndSlotCount, EpochSlots)
forall a b. (a -> b) -> a -> b
$ EpochNumber -> SlotCount -> EpochAndSlotCount
EpochAndSlotCount
    (EpochNumber -> SlotCount -> EpochAndSlotCount)
-> Gen EpochNumber
-> GenT Identity (SlotCount -> EpochAndSlotCount)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64 -> Gen EpochNumber
genRestrictedEpochNumber (Word64
forall a. Bounded a => a
maxBound Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` EpochSlots -> Word64
unEpochSlots EpochSlots
epochSlots)
    GenT Identity (SlotCount -> EpochAndSlotCount)
-> Gen SlotCount -> Gen EpochAndSlotCount
forall a b.
GenT Identity (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> EpochSlots -> Gen SlotCount
genEpochSlotCount EpochSlots
epochSlots
  where
    genRestrictedEpochNumber :: Word64 -> Gen EpochNumber
    genRestrictedEpochNumber :: Word64 -> Gen EpochNumber
genRestrictedEpochNumber Word64
bound =
      Word64 -> EpochNumber
EpochNumber (Word64 -> EpochNumber) -> GenT Identity Word64 -> Gen EpochNumber
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word64 -> GenT Identity Word64
forall (m :: * -> *). MonadGen m => Range Word64 -> m Word64
Gen.word64 (Word64 -> Word64 -> Range Word64
forall a. Integral a => a -> a -> Range a
Range.linear Word64
0 Word64
bound)

feedPMEpochSlots :: (ProtocolMagicId -> EpochSlots -> Gen a) -> Gen a
feedPMEpochSlots :: forall a. (ProtocolMagicId -> EpochSlots -> Gen a) -> Gen a
feedPMEpochSlots ProtocolMagicId -> EpochSlots -> Gen a
genA = do
  ProtocolMagicId
pm <- Gen ProtocolMagicId
genProtocolMagicId
  EpochSlots
epochSlots <- Gen EpochSlots
genEpochSlots
  ProtocolMagicId -> EpochSlots -> Gen a
genA ProtocolMagicId
pm EpochSlots
epochSlots