{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | These functions iteratively produce all race windows in a slot vector.
--
-- The first window is produced by 'init', which unconditionally starts the window at the first slot.
-- This window can then be passed to 'next', which starts the new window after the first active slot.
--
-- @
-- ---X--X--X--X-- ...
-- ^ start of window 1 from 'init'
--     ^ start of window 2 from 'next'
--        ^ start of window 3 from 'next'
-- @
--
-- Valid windows must have @k+1@ active slots.
-- If the vector doesn't have sufficient slots to meet this condition, 'init' and 'next' return 'Nothing' and we fall back
-- to 'initConservative' and 'nextConservative', which return windows truncated at the end of time.
module Test.Ouroboros.Consensus.ChainGenerator.RaceIterator
  ( Race (Race, UnsafeRace)
  , RaceLbl
  , init
  , initConservative
  , next
  , nextConservative
  ) where

import Control.Monad (when)
import Data.Proxy (Proxy (Proxy))
import qualified Test.Ouroboros.Consensus.ChainGenerator.BitVector as BV
import qualified Test.Ouroboros.Consensus.ChainGenerator.Counting as C
import Test.Ouroboros.Consensus.ChainGenerator.Params (Kcp (Kcp))
import Test.Ouroboros.Consensus.ChainGenerator.Slot
  ( E (ActiveSlotE, SlotE)
  , S
  )
import Prelude hiding (init)

-----

-- | See 'Race'
data RaceLbl

-- | A window whose last slot contains the @k+1@st active slot in it
newtype Race base = UnsafeRace (C.SomeWindow RaceLbl base SlotE)
  deriving (Race base -> Race base -> Bool
(Race base -> Race base -> Bool)
-> (Race base -> Race base -> Bool) -> Eq (Race base)
forall base. Race base -> Race base -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall base. Race base -> Race base -> Bool
== :: Race base -> Race base -> Bool
$c/= :: forall base. Race base -> Race base -> Bool
/= :: Race base -> Race base -> Bool
Eq, ReadPrec [Race base]
ReadPrec (Race base)
Int -> ReadS (Race base)
ReadS [Race base]
(Int -> ReadS (Race base))
-> ReadS [Race base]
-> ReadPrec (Race base)
-> ReadPrec [Race base]
-> Read (Race base)
forall base. ReadPrec [Race base]
forall base. ReadPrec (Race base)
forall base. Int -> ReadS (Race base)
forall base. ReadS [Race base]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: forall base. Int -> ReadS (Race base)
readsPrec :: Int -> ReadS (Race base)
$creadList :: forall base. ReadS [Race base]
readList :: ReadS [Race base]
$creadPrec :: forall base. ReadPrec (Race base)
readPrec :: ReadPrec (Race base)
$creadListPrec :: forall base. ReadPrec [Race base]
readListPrec :: ReadPrec [Race base]
Read, Int -> Race base -> ShowS
[Race base] -> ShowS
Race base -> String
(Int -> Race base -> ShowS)
-> (Race base -> String)
-> ([Race base] -> ShowS)
-> Show (Race base)
forall base. Int -> Race base -> ShowS
forall base. [Race base] -> ShowS
forall base. Race base -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall base. Int -> Race base -> ShowS
showsPrec :: Int -> Race base -> ShowS
$cshow :: forall base. Race base -> String
show :: Race base -> String
$cshowList :: forall base. [Race base] -> ShowS
showList :: [Race base] -> ShowS
Show)

pattern Race :: C.SomeWindow RaceLbl base SlotE -> Race base
pattern $mRace :: forall {r} {base}.
Race base
-> (SomeWindow RaceLbl base 'SlotE -> r) -> ((# #) -> r) -> r
Race x <- UnsafeRace x

{-# COMPLETE Race #-}

-----

-- | Find the nth active slot /in/ the given race window and return the slot number in the context of the entire chain.
--
-- Race windows are anchored in an active slot, and so could start with an empty or active slot.
nthActiveSlotIndex ::
  forall base adv.
  C.Index adv ActiveSlotE ->
  C.Vector base SlotE S ->
  C.Contains SlotE base adv ->
  Maybe (C.Index base SlotE)
nthActiveSlotIndex :: forall base adv.
Index adv 'ActiveSlotE
-> Vector base 'SlotE S
-> Contains 'SlotE base adv
-> Maybe (Index base 'SlotE)
nthActiveSlotIndex Index adv 'ActiveSlotE
n Vector base 'SlotE S
v Contains 'SlotE base adv
raceWin =
  -- the given race window has at least k+1 blocks in it and 0<=k, so this pattern can't fail
  -- TODO by invariant during construction of the honest chain?
  case Vector adv 'SlotE S -> Index adv 'ActiveSlotE -> MaybeFound adv
forall base.
Vector base 'SlotE S -> Index base 'ActiveSlotE -> MaybeFound base
BV.findIthActiveInV (Contains 'SlotE base adv
-> Vector base 'SlotE S -> Vector adv 'SlotE S
forall {k} a (elem :: k) outer inner.
Unbox a =>
Contains elem outer inner
-> Vector outer elem a -> Vector inner elem a
C.sliceV Contains 'SlotE base adv
raceWin Vector base 'SlotE S
v) Index adv 'ActiveSlotE
n of
    MaybeFound adv
BV.NothingFound -> Maybe (Index base 'SlotE)
forall a. Maybe a
Nothing -- would be impossible if we never called next after *Conservative
    BV.JustFound Index adv 'SlotE
slot -> Index base 'SlotE -> Maybe (Index base 'SlotE)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Index base 'SlotE -> Maybe (Index base 'SlotE))
-> Index base 'SlotE -> Maybe (Index base 'SlotE)
forall a b. (a -> b) -> a -> b
$! Contains 'SlotE base adv -> Index adv 'SlotE -> Index base 'SlotE
forall {kelem} (elem :: kelem) outer inner.
Contains elem outer inner -> Index inner elem -> Index outer elem
C.fromWindow Contains 'SlotE base adv
raceWin Index adv 'SlotE
slot

-- | Yields the race window starting at position 0 of the given
-- vector if the @k+1@ active slot exists.
init :: Kcp -> C.Vector base SlotE S -> Maybe (Race base)
init :: forall base. Kcp -> Vector base 'SlotE S -> Maybe (Race base)
init (Kcp Int
k) Vector base 'SlotE S
v =
  -- find the @k+1@st active slot in the given race window
  case Vector base 'SlotE S -> Index base 'ActiveSlotE -> MaybeFound base
forall base.
Vector base 'SlotE S -> Index base 'ActiveSlotE -> MaybeFound base
BV.findIthActiveInV Vector base 'SlotE S
v (Int -> Index base 'ActiveSlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
k) of
    MaybeFound base
BV.NothingFound -> Maybe (Race base)
forall a. Maybe a
Nothing
    BV.JustFound Index base 'SlotE
kPlus1st ->
      Race base -> Maybe (Race base)
forall a. a -> Maybe a
Just (Race base -> Maybe (Race base)) -> Race base -> Maybe (Race base)
forall a b. (a -> b) -> a -> b
$!
        SomeWindow RaceLbl base 'SlotE -> Race base
forall base. SomeWindow RaceLbl base 'SlotE -> Race base
UnsafeRace (SomeWindow RaceLbl base 'SlotE -> Race base)
-> SomeWindow RaceLbl base 'SlotE -> Race base
forall a b. (a -> b) -> a -> b
$
          Size base 'SlotE
-> Lbl RaceLbl
-> Index base 'SlotE
-> Index base 'SlotE
-> SomeWindow RaceLbl base 'SlotE
forall {kelem} {klbl} outer (elem :: kelem) (lbl :: klbl).
Size outer elem
-> Lbl lbl
-> Index outer elem
-> Index outer elem
-> SomeWindow lbl outer elem
C.withWindowBetween
            (Vector base 'SlotE S -> Size base 'SlotE
forall {kelem} a base (elem :: kelem).
Unbox a =>
Vector base elem a -> Size base elem
C.lengthV Vector base 'SlotE S
v)
            (forall lbl. Lbl lbl
forall {k} (lbl :: k). Lbl lbl
C.Lbl @RaceLbl)
            (Int -> Index base 'SlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
0)
            Index base 'SlotE
kPlus1st

-- | @initConservative@ creates a window for the whole vector.
initConservative :: C.Vector base SlotE S -> Race base
initConservative :: forall base. Vector base 'SlotE S -> Race base
initConservative Vector base 'SlotE S
v =
  let sz :: Size base 'SlotE
sz = Vector base 'SlotE S -> Size base 'SlotE
forall {kelem} a base (elem :: kelem).
Unbox a =>
Vector base elem a -> Size base elem
C.lengthV Vector base 'SlotE S
v
   in SomeWindow RaceLbl base 'SlotE -> Race base
forall base. SomeWindow RaceLbl base 'SlotE -> Race base
UnsafeRace (SomeWindow RaceLbl base 'SlotE -> Race base)
-> SomeWindow RaceLbl base 'SlotE -> Race base
forall a b. (a -> b) -> a -> b
$
        Size base 'SlotE
-> Lbl RaceLbl
-> Index base 'SlotE
-> Size (ZonkAny 0) 'SlotE
-> SomeWindow RaceLbl base 'SlotE
forall {kelem} {klbl} outer (elem :: kelem) (lbl :: klbl) x.
Size outer elem
-> Lbl lbl
-> Index outer elem
-> Size x elem
-> SomeWindow lbl outer elem
C.withWindow
          Size base 'SlotE
sz
          (forall lbl. Lbl lbl
forall {k} (lbl :: k). Lbl lbl
C.Lbl @RaceLbl)
          (Int -> Index base 'SlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
0)
          (Int -> Size (ZonkAny 0) 'SlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count (Int -> Size (ZonkAny 0) 'SlotE) -> Int -> Size (ZonkAny 0) 'SlotE
forall a b. (a -> b) -> a -> b
$ Size base 'SlotE -> Int
forall {kelem} {kwhich} base (elem :: kelem) (which :: kwhich).
Count base elem which -> Int
C.getCount Size base 'SlotE
sz)

data RaceStepLbl

-- | @next v r@ yields the race window anchored at the first
-- active slot of @r@ if there is an active slot after @r@.
next ::
  forall base.
  C.Vector base SlotE S ->
  Race base ->
  Maybe (Race base)
next :: forall base. Vector base 'SlotE S -> Race base -> Maybe (Race base)
next Vector base 'SlotE S
v (UnsafeRace (C.SomeWindow Proxy skolem
Proxy Contains 'SlotE base (Win RaceLbl skolem)
raceWin)) = do
  next0 <- Index (Win RaceLbl skolem) 'ActiveSlotE
-> Vector base 'SlotE S
-> Contains 'SlotE base (Win RaceLbl skolem)
-> Maybe (Index base 'SlotE)
forall base adv.
Index adv 'ActiveSlotE
-> Vector base 'SlotE S
-> Contains 'SlotE base adv
-> Maybe (Index base 'SlotE)
nthActiveSlotIndex (Int -> Index (Win RaceLbl skolem) 'ActiveSlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
0) Vector base 'SlotE S
v Contains 'SlotE base (Win RaceLbl skolem)
raceWin

  -- find the first active slot /after/ the given race window
  --
  -- Race windows end in an active slot.
  nextK <- do
    C.SomeWindow Proxy searchWin <-
      pure $
        C.withWindowBetween
          sz
          (C.Lbl @RaceStepLbl)
          (C.windowLast raceWin)
          (C.lastIndex sz)
    nthActiveSlotIndex (C.Count 1) v searchWin

  pure $!
    UnsafeRace $
      C.withWindowBetween
        sz
        (C.Lbl @RaceLbl)
        (next0 C.+ 1)
        nextK
 where
  sz :: Size base 'SlotE
sz = Vector base 'SlotE S -> Size base 'SlotE
forall {kelem} a base (elem :: kelem).
Unbox a =>
Vector base elem a -> Size base elem
C.lengthV Vector base 'SlotE S
v

-- | @nextConservative v r@ yields a window anchored at the first
-- active slot of @r@ if it exists, and extending until the end of @v@.
nextConservative ::
  forall base.
  C.Vector base SlotE S ->
  Race base ->
  Maybe (Race base)
nextConservative :: forall base. Vector base 'SlotE S -> Race base -> Maybe (Race base)
nextConservative Vector base 'SlotE S
v (UnsafeRace (C.SomeWindow Proxy skolem
Proxy Contains 'SlotE base (Win RaceLbl skolem)
raceWin)) = do
  next0 <- Index (Win RaceLbl skolem) 'ActiveSlotE
-> Vector base 'SlotE S
-> Contains 'SlotE base (Win RaceLbl skolem)
-> Maybe (Index base 'SlotE)
forall base adv.
Index adv 'ActiveSlotE
-> Vector base 'SlotE S
-> Contains 'SlotE base adv
-> Maybe (Index base 'SlotE)
nthActiveSlotIndex (Int -> Index (Win RaceLbl skolem) 'ActiveSlotE
forall kelem kwhich base (elem :: kelem) (which :: kwhich).
Int -> Count base elem which
C.Count Int
0) Vector base 'SlotE S
v Contains 'SlotE base (Win RaceLbl skolem)
raceWin

  -- do not return a Race Window that starts after 'Len'
  when (next0 == C.lastIndex sz) Nothing

  pure $!
    UnsafeRace $
      C.withWindowBetween
        sz
        (C.Lbl @RaceLbl)
        (next0 C.+ 1)
        (C.lastIndex sz)
 where
  sz :: Size base 'SlotE
sz = Vector base 'SlotE S -> Size base 'SlotE
forall {kelem} a base (elem :: kelem).
Unbox a =>
Vector base elem a -> Size base elem
C.lengthV Vector base 'SlotE S
v