{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE LambdaCase #-}

module Test.Util.HardFork.Future (
    EraSize (..)
  , Future (..)
  , futureEpochInFirstEra
  , futureFirstEpochSize
  , futureFirstSlotLength
  , futureSlotLengths
  , futureSlotToEpoch
  , futureSlotToTime
  , futureTimeToSlot
  , singleEraFuture
  ) where

import qualified Data.Fixed
import           Data.Time (NominalDiffTime)
import           Data.Word (Word64)
import           GHC.Generics (Generic)
import           Ouroboros.Consensus.Block
import           Ouroboros.Consensus.BlockchainTime (SlotLength (..))
import           Ouroboros.Consensus.Util (nTimes)
import           Quiet (Quiet (..))
import           Test.Util.Slots (NumSlots (..))
import           Test.Util.Stream (Stream (..))

{-------------------------------------------------------------------------------
  Careful counts
-------------------------------------------------------------------------------}

-- | Number of epochs
newtype EraSize = EraSize {EraSize -> Word64
unEraSize :: Word64}
  deriving (EraSize -> EraSize -> Bool
(EraSize -> EraSize -> Bool)
-> (EraSize -> EraSize -> Bool) -> Eq EraSize
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EraSize -> EraSize -> Bool
== :: EraSize -> EraSize -> Bool
$c/= :: EraSize -> EraSize -> Bool
/= :: EraSize -> EraSize -> Bool
Eq, (forall x. EraSize -> Rep EraSize x)
-> (forall x. Rep EraSize x -> EraSize) -> Generic EraSize
forall x. Rep EraSize x -> EraSize
forall x. EraSize -> Rep EraSize x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. EraSize -> Rep EraSize x
from :: forall x. EraSize -> Rep EraSize x
$cto :: forall x. Rep EraSize x -> EraSize
to :: forall x. Rep EraSize x -> EraSize
Generic)
  deriving (Int -> EraSize -> ShowS
[EraSize] -> ShowS
EraSize -> String
(Int -> EraSize -> ShowS)
-> (EraSize -> String) -> ([EraSize] -> ShowS) -> Show EraSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EraSize -> ShowS
showsPrec :: Int -> EraSize -> ShowS
$cshow :: EraSize -> String
show :: EraSize -> String
$cshowList :: [EraSize] -> ShowS
showList :: [EraSize] -> ShowS
Show) via (Quiet EraSize)

{-------------------------------------------------------------------------------
  A test's whole timeline
-------------------------------------------------------------------------------}

-- | Every era in the test
--
-- INVARIANT: every number is @> 0@
data Future =
      EraFinal SlotLength EpochSize
    | EraCons  SlotLength EpochSize EraSize Future
  deriving (Future -> Future -> Bool
(Future -> Future -> Bool)
-> (Future -> Future -> Bool) -> Eq Future
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Future -> Future -> Bool
== :: Future -> Future -> Bool
$c/= :: Future -> Future -> Bool
/= :: Future -> Future -> Bool
Eq, Int -> Future -> ShowS
[Future] -> ShowS
Future -> String
(Int -> Future -> ShowS)
-> (Future -> String) -> ([Future] -> ShowS) -> Show Future
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Future -> ShowS
showsPrec :: Int -> Future -> ShowS
$cshow :: Future -> String
show :: Future -> String
$cshowList :: [Future] -> ShowS
showList :: [Future] -> ShowS
Show)

-- | 'Future' with only one era
singleEraFuture :: SlotLength -> EpochSize -> Future
singleEraFuture :: SlotLength -> EpochSize -> Future
singleEraFuture = SlotLength -> EpochSize -> Future
EraFinal

-- | 'SlotLength' of the first era
futureFirstSlotLength :: Future -> SlotLength
futureFirstSlotLength :: Future -> SlotLength
futureFirstSlotLength Future
future = case Future
future of
    EraCons  SlotLength
slotLength EpochSize
_epochSize EraSize
_eraSize Future
_future -> SlotLength
slotLength
    EraFinal SlotLength
slotLength EpochSize
_epochSize                  -> SlotLength
slotLength

-- | 'EpochSize' of the first era
futureFirstEpochSize :: Future -> EpochSize
futureFirstEpochSize :: Future -> EpochSize
futureFirstEpochSize Future
future = case Future
future of
    EraCons  SlotLength
_slotLength EpochSize
epochSize EraSize
_eraSize Future
_future -> EpochSize
epochSize
    EraFinal SlotLength
_slotLength EpochSize
epochSize                  -> EpochSize
epochSize

-- | Length of each slot in the whole 'Future'
futureSlotLengths :: Future -> Stream SlotLength
futureSlotLengths :: Future -> Stream SlotLength
futureSlotLengths = \case
    EraFinal SlotLength
slotLength EpochSize
_epochSize ->
        let x :: Stream SlotLength
x = SlotLength
slotLength SlotLength -> Stream SlotLength -> Stream SlotLength
forall a. a -> Stream a -> Stream a
:< Stream SlotLength
x in Stream SlotLength
x
    EraCons SlotLength
slotLength EpochSize
epochSize EraSize
eraSize Future
future ->
        (Stream SlotLength -> Stream SlotLength)
-> Word64 -> Stream SlotLength -> Stream SlotLength
forall a. (a -> a) -> Word64 -> a -> a
nTimes (SlotLength
slotLength SlotLength -> Stream SlotLength -> Stream SlotLength
forall a. a -> Stream a -> Stream a
:<) Word64
eraSlots (Stream SlotLength -> Stream SlotLength)
-> Stream SlotLength -> Stream SlotLength
forall a b. (a -> b) -> a -> b
$
        Future -> Stream SlotLength
futureSlotLengths Future
future
      where
        NumSlots Word64
eraSlots = EpochSize -> EraSize -> NumSlots
calcEraSlots EpochSize
epochSize EraSize
eraSize

-- | @(slot, time left in slot, length of slot)@
futureTimeToSlot :: Future
                 -> NominalDiffTime
                 -> (SlotNo, NominalDiffTime, SlotLength)
futureTimeToSlot :: Future -> NominalDiffTime -> (SlotNo, NominalDiffTime, SlotLength)
futureTimeToSlot = \Future
future NominalDiffTime
d -> Word64
-> NominalDiffTime
-> Future
-> (SlotNo, NominalDiffTime, SlotLength)
go Word64
0 NominalDiffTime
d Future
future
  where
    done :: Word64
-> NominalDiffTime
-> SlotLength
-> (SlotNo, NominalDiffTime, SlotLength)
done Word64
acc NominalDiffTime
d SlotLength
slotLength =
        (Word64 -> SlotNo
SlotNo (Word64 -> SlotNo) -> Word64 -> SlotNo
forall a b. (a -> b) -> a -> b
$ Word64
acc Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
n, SlotLength -> NominalDiffTime
getSlotLength SlotLength
slotLength NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
- NominalDiffTime
timeInSlot, SlotLength
slotLength)
      where
        n :: Word64
n          = NominalDiffTime -> SlotLength -> Word64
divide NominalDiffTime
d SlotLength
slotLength
        timeInSlot :: NominalDiffTime
timeInSlot = NominalDiffTime
d NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
- Word64 -> SlotLength -> NominalDiffTime
multiply Word64
n SlotLength
slotLength

    go :: Word64
-> NominalDiffTime
-> Future
-> (SlotNo, NominalDiffTime, SlotLength)
go Word64
acc NominalDiffTime
d (EraFinal SlotLength
slotLength EpochSize
_epochSize) =
        Word64
-> NominalDiffTime
-> SlotLength
-> (SlotNo, NominalDiffTime, SlotLength)
done Word64
acc NominalDiffTime
d SlotLength
slotLength
    go Word64
acc NominalDiffTime
d (EraCons SlotLength
slotLength EpochSize
epochSize EraSize
eraSize Future
future) =
        case NominalDiffTime
d NominalDiffTime -> NominalDiffTime -> Maybe NominalDiffTime
forall a. (Num a, Ord a) => a -> a -> Maybe a
`safeSub` NominalDiffTime
eraLength of
          Maybe NominalDiffTime
Nothing -> Word64
-> NominalDiffTime
-> SlotLength
-> (SlotNo, NominalDiffTime, SlotLength)
done Word64
acc NominalDiffTime
d SlotLength
slotLength
          Just NominalDiffTime
d' -> Word64
-> NominalDiffTime
-> Future
-> (SlotNo, NominalDiffTime, SlotLength)
go (Word64
acc Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
eraSlots) NominalDiffTime
d' Future
future
      where
        NumSlots Word64
eraSlots = EpochSize -> EraSize -> NumSlots
calcEraSlots EpochSize
epochSize EraSize
eraSize
        eraLength :: NominalDiffTime
eraLength         = Word64 -> SlotLength -> NominalDiffTime
multiply Word64
eraSlots SlotLength
slotLength

-- | Which epoch the slot is in
futureSlotToEpoch :: Future
                  -> SlotNo
                  -> EpochNo
futureSlotToEpoch :: Future -> SlotNo -> EpochNo
futureSlotToEpoch = \Future
future (SlotNo Word64
s) -> Word64 -> EpochNo
EpochNo (Word64 -> EpochNo) -> Word64 -> EpochNo
forall a b. (a -> b) -> a -> b
$ Word64 -> Word64 -> Future -> Word64
go Word64
0 Word64
s Future
future
  where
    go :: Word64 -> Word64 -> Future -> Word64
go Word64
acc Word64
s = \case
      EraFinal SlotLength
_slotLength (EpochSize Word64
epSz)                ->
          Word64
acc Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
s Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`div` Word64
epSz
      EraCons  SlotLength
slotLength  EpochSize
epochSize        EraSize
eraSize Future
future ->
          case Word64
s Word64 -> Word64 -> Maybe Word64
forall a. (Num a, Ord a) => a -> a -> Maybe a
`safeSub` Word64
eraSlots of
            Maybe Word64
Nothing -> Word64 -> Word64 -> Future -> Word64
go Word64
acc Word64
s (SlotLength -> EpochSize -> Future
EraFinal SlotLength
slotLength EpochSize
epochSize)
            Just Word64
s' -> Word64 -> Word64 -> Future -> Word64
go (Word64
acc Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
n) Word64
s' Future
future
        where
          EraSize Word64
n = EraSize
eraSize
          NumSlots Word64
eraSlots = EpochSize -> EraSize -> NumSlots
calcEraSlots EpochSize
epochSize EraSize
eraSize

-- | When the slot begins
futureSlotToTime :: Future
                 -> SlotNo
                 -> NominalDiffTime
futureSlotToTime :: Future -> SlotNo -> NominalDiffTime
futureSlotToTime = \Future
future (SlotNo Word64
s) -> NominalDiffTime -> Word64 -> Future -> NominalDiffTime
go NominalDiffTime
0 Word64
s Future
future
  where
    done :: NominalDiffTime -> Word64 -> SlotLength -> NominalDiffTime
done NominalDiffTime
acc Word64
s SlotLength
slotLength =
        NominalDiffTime
acc NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
+ Word64 -> SlotLength -> NominalDiffTime
multiply Word64
s SlotLength
slotLength

    go :: NominalDiffTime -> Word64 -> Future -> NominalDiffTime
go NominalDiffTime
acc Word64
s = \case
      EraFinal SlotLength
slotLength EpochSize
_epochSize                ->
          NominalDiffTime -> Word64 -> SlotLength -> NominalDiffTime
done NominalDiffTime
acc Word64
s SlotLength
slotLength
      EraCons  SlotLength
slotLength EpochSize
epochSize  EraSize
eraSize Future
future ->
          case Word64
s Word64 -> Word64 -> Maybe Word64
forall a. (Num a, Ord a) => a -> a -> Maybe a
`safeSub` Word64
eraSlots of
            Maybe Word64
Nothing -> NominalDiffTime -> Word64 -> SlotLength -> NominalDiffTime
done NominalDiffTime
acc Word64
s SlotLength
slotLength
            Just Word64
s' -> NominalDiffTime -> Word64 -> Future -> NominalDiffTime
go (NominalDiffTime
acc NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
+ NominalDiffTime
eraLength) Word64
s' Future
future
        where
          NumSlots Word64
eraSlots = EpochSize -> EraSize -> NumSlots
calcEraSlots EpochSize
epochSize EraSize
eraSize
          eraLength :: NominalDiffTime
eraLength         = Word64 -> SlotLength -> NominalDiffTime
multiply Word64
eraSlots SlotLength
slotLength

-- | Whether the epoch is in the first era
futureEpochInFirstEra :: Future -> EpochNo -> Bool
futureEpochInFirstEra :: Future -> EpochNo -> Bool
futureEpochInFirstEra = \case
    EraCons SlotLength
_slotLength EpochSize
_epochSize (EraSize Word64
n) Future
_future ->
        \(EpochNo Word64
e) -> Word64
e Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
n
    EraFinal{} -> Bool -> EpochNo -> Bool
forall a b. a -> b -> a
const Bool
True

{-------------------------------------------------------------------------------
  Miscellany
-------------------------------------------------------------------------------}

divide :: NominalDiffTime -> SlotLength -> Word64
divide :: NominalDiffTime -> SlotLength -> Word64
divide NominalDiffTime
d SlotLength
slotLength = NominalDiffTime -> NominalDiffTime -> Word64
forall a b. (Real a, Integral b) => a -> a -> b
Data.Fixed.div' NominalDiffTime
d (SlotLength -> NominalDiffTime
getSlotLength SlotLength
slotLength)

multiply :: Word64 -> SlotLength -> NominalDiffTime
multiply :: Word64 -> SlotLength -> NominalDiffTime
multiply Word64
m SlotLength
sl = Word64 -> NominalDiffTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
m NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* SlotLength -> NominalDiffTime
getSlotLength SlotLength
sl

-- | Find the non-negative difference
safeSub :: (Num a, Ord a) => a -> a -> Maybe a
safeSub :: forall a. (Num a, Ord a) => a -> a -> Maybe a
safeSub a
x a
y = if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
y then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
y)

calcEraSlots :: EpochSize -> EraSize -> NumSlots
calcEraSlots :: EpochSize -> EraSize -> NumSlots
calcEraSlots (EpochSize Word64
slotPerEpoch) (EraSize Word64
epochPerEra) =
    Word64 -> NumSlots
NumSlots (Word64
slotPerEpoch Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
epochPerEra)