{-# 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)