{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

module Test.Cardano.Network.PeerSelection.Utils where

import Control.Monad.Class.MonadTime.SI
import Data.Set (Set)
import Data.Set qualified as Set
import System.Random (mkStdGen)

import Ouroboros.Network.PeerSelection.Governor (AssociationMode (..),
           PeerSelectionTargets (..), TracePeerSelection (..))
import Ouroboros.Network.PeerSelection.Governor qualified as Governor
import Ouroboros.Network.PeerSelection.PublicRootPeers qualified as PublicRootPeers

import Test.Ouroboros.Network.Data.Signal (Events, Signal)
import Test.Ouroboros.Network.Data.Signal qualified as Signal
import Test.Ouroboros.Network.PeerSelection.Instances (PeerAddr (..))

import Test.Cardano.Network.PeerSelection.MockEnvironment hiding (targets,
           tests)


takeFirstNHours :: DiffTime -> [(Time, a)] -> [(Time, a)]
takeFirstNHours :: forall a. DiffTime -> [(Time, a)] -> [(Time, a)]
takeFirstNHours DiffTime
h = ((Time, a) -> Bool) -> [(Time, a)] -> [(Time, a)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\(Time
t,a
_) -> Time
t Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
< DiffTime -> Time
Time (DiffTime
60DiffTime -> DiffTime -> DiffTime
forall a. Num a => a -> a -> a
*DiffTime
60DiffTime -> DiffTime -> DiffTime
forall a. Num a => a -> a -> a
*DiffTime
h))


selectEnvEvents :: Events (TestTraceEvent extraState extraFlags extraPeers)
                -> Events TraceMockEnv
selectEnvEvents :: forall extraState extraFlags extraPeers.
Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events TraceMockEnv
selectEnvEvents = (TestTraceEvent extraState extraFlags extraPeers
 -> Maybe TraceMockEnv)
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events TraceMockEnv
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                    (\case MockEnvEvent TraceMockEnv
e -> TraceMockEnv -> Maybe TraceMockEnv
forall a. a -> Maybe a
Just (TraceMockEnv -> Maybe TraceMockEnv)
-> TraceMockEnv -> Maybe TraceMockEnv
forall a b. (a -> b) -> a -> b
$! TraceMockEnv
e
                           TestTraceEvent extraState extraFlags extraPeers
_              -> Maybe TraceMockEnv
forall a. Maybe a
Nothing)

selectGovEvents :: Events (TestTraceEvent extraState extraFlags extraPeers)
                -> Events (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
selectGovEvents :: forall extraState extraFlags extraPeers.
Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
selectGovEvents = (TestTraceEvent extraState extraFlags extraPeers
 -> Maybe
      (TracePeerSelection extraState extraFlags extraPeers PeerAddr))
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                    (\case GovernorEvent TracePeerSelection extraState extraFlags extraPeers PeerAddr
e -> TracePeerSelection extraState extraFlags extraPeers PeerAddr
-> Maybe
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a. a -> Maybe a
Just (TracePeerSelection extraState extraFlags extraPeers PeerAddr
 -> Maybe
      (TracePeerSelection extraState extraFlags extraPeers PeerAddr))
-> TracePeerSelection extraState extraFlags extraPeers PeerAddr
-> Maybe
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a b. (a -> b) -> a -> b
$! TracePeerSelection extraState extraFlags extraPeers PeerAddr
e
                           TestTraceEvent extraState extraFlags extraPeers
_               -> Maybe
  (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a. Maybe a
Nothing)

selectGovCounters :: Events (TestTraceEvent extraState extraFlags extraPeers)
                  -> Events (Governor.PeerSelectionCounters (Governor.ViewExtraPeers extraPeers))
selectGovCounters :: forall extraState extraFlags extraPeers.
Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events (PeerSelectionCounters (ViewExtraPeers extraPeers))
selectGovCounters = (TestTraceEvent extraState extraFlags extraPeers
 -> Maybe (PeerSelectionCounters (ViewExtraPeers extraPeers)))
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events (PeerSelectionCounters (ViewExtraPeers extraPeers))
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                      (\case GovernorCounters PeerSelectionCounters (ViewExtraPeers extraPeers)
e -> PeerSelectionCounters (ViewExtraPeers extraPeers)
-> Maybe (PeerSelectionCounters (ViewExtraPeers extraPeers))
forall a. a -> Maybe a
Just (PeerSelectionCounters (ViewExtraPeers extraPeers)
 -> Maybe (PeerSelectionCounters (ViewExtraPeers extraPeers)))
-> PeerSelectionCounters (ViewExtraPeers extraPeers)
-> Maybe (PeerSelectionCounters (ViewExtraPeers extraPeers))
forall a b. (a -> b) -> a -> b
$! PeerSelectionCounters (ViewExtraPeers extraPeers)
e
                             TestTraceEvent extraState extraFlags extraPeers
_                  -> Maybe (PeerSelectionCounters (ViewExtraPeers extraPeers))
forall a. Maybe a
Nothing)

selectGovAssociationMode :: Events (TestTraceEvent extraState extraFlags extraPeers)
                         -> Events AssociationMode
selectGovAssociationMode :: forall extraState extraFlags extraPeers.
Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events AssociationMode
selectGovAssociationMode = (TestTraceEvent extraState extraFlags extraPeers
 -> Maybe AssociationMode)
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events AssociationMode
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                             (\case GovernorAssociationMode AssociationMode
e -> AssociationMode -> Maybe AssociationMode
forall a. a -> Maybe a
Just (AssociationMode -> Maybe AssociationMode)
-> AssociationMode -> Maybe AssociationMode
forall a b. (a -> b) -> a -> b
$! AssociationMode
e
                                    TestTraceEvent extraState extraFlags extraPeers
_                         -> Maybe AssociationMode
forall a. Maybe a
Nothing)

selectGovState :: Eq a
               => (forall peerconn.
                       Governor.PeerSelectionState extraState extraFlags extraPeers
                                                   PeerAddr peerconn
                    -> a
                  )
               -> extraState
               -> extraPeers
               -> Events (TestTraceEvent extraState extraFlags extraPeers)
               -> Signal a
selectGovState :: forall a extraState extraFlags extraPeers.
Eq a =>
(forall peerconn.
 PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> a)
-> extraState
-> extraPeers
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal a
selectGovState forall peerconn.
PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
f extraState
es extraPeers
ep =
    Signal a -> Signal a
forall a. Eq a => Signal a -> Signal a
Signal.nub
  -- TODO: #3182 Rng seed should come from quickcheck.
  --       and `NumberOfBigLedgerPeers`
  (Signal a -> Signal a)
-> (Events (TestTraceEvent extraState extraFlags extraPeers)
    -> Signal a)
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Events a -> Signal a
forall a. a -> Events a -> Signal a
Signal.fromChangeEvents (PeerSelectionState
  extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
-> a
forall peerconn.
PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
f (PeerSelectionState
   extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
 -> a)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
-> a
forall a b. (a -> b) -> a -> b
$! StdGen
-> extraState
-> extraPeers
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
forall extraState extraPeers extraFlags peeraddr peerconn.
StdGen
-> extraState
-> extraPeers
-> PeerSelectionState
     extraState extraFlags extraPeers peeraddr peerconn
Governor.emptyPeerSelectionState (Int -> StdGen
mkStdGen Int
42) extraState
es extraPeers
ep)
  (Events a -> Signal a)
-> (Events (TestTraceEvent extraState extraFlags extraPeers)
    -> Events a)
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TestTraceEvent extraState extraFlags extraPeers -> Maybe a)
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events a
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
      (\case GovernorDebug (Governor.TraceGovernorState Time
_ Maybe DiffTime
_ PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st) -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
forall peerconn.
PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
f PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st
             TestTraceEvent extraState extraFlags extraPeers
_                                                  -> Maybe a
forall a. Maybe a
Nothing)

selectGovTargets :: Eq a
                 => (PeerSelectionTargets -> a)
                 -> extraState
                 -> extraPeers
                 -> Events (TestTraceEvent extraState extraFlags extraPeers)
                 -> Signal a
selectGovTargets :: forall a extraState extraPeers extraFlags.
Eq a =>
(PeerSelectionTargets -> a)
-> extraState
-> extraPeers
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal a
selectGovTargets PeerSelectionTargets -> a
f = (forall peerconn.
 PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> a)
-> extraState
-> extraPeers
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal a
forall a extraState extraFlags extraPeers.
Eq a =>
(forall peerconn.
 PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> a)
-> extraState
-> extraPeers
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal a
selectGovState (PeerSelectionTargets -> a
f (PeerSelectionTargets -> a)
-> (PeerSelectionState
      extraState extraFlags extraPeers PeerAddr peerconn
    -> PeerSelectionTargets)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> PeerSelectionTargets
forall extraState extraFlags extraPeers peeraddr peerconn.
PeerSelectionState
  extraState extraFlags extraPeers peeraddr peerconn
-> PeerSelectionTargets
Governor.targets)

selectForgottenPeers :: Events (TestTraceEvent extraState extraFlags extraPeers)
                     -> Signal (Set PeerAddr)
selectForgottenPeers :: forall extraState extraFlags extraPeers.
Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal (Set PeerAddr)
selectForgottenPeers =
    Set PeerAddr -> Events (Set PeerAddr) -> Signal (Set PeerAddr)
forall a. a -> Events a -> Signal a
Signal.fromChangeEvents Set PeerAddr
forall a. Set a
Set.empty
  (Events (Set PeerAddr) -> Signal (Set PeerAddr))
-> (Events (TestTraceEvent extraState extraFlags extraPeers)
    -> Events (Set PeerAddr))
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Signal (Set PeerAddr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TestTraceEvent extraState extraFlags extraPeers
 -> Maybe (Set PeerAddr))
-> Events (TestTraceEvent extraState extraFlags extraPeers)
-> Events (Set PeerAddr)
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
      (\case GovernorEvent (TraceForgottenPeers Set PeerAddr
peers) -> Set PeerAddr -> Maybe (Set PeerAddr)
forall a. a -> Maybe a
Just (Set PeerAddr -> Maybe (Set PeerAddr))
-> Set PeerAddr -> Maybe (Set PeerAddr)
forall a b. (a -> b) -> a -> b
$! Set PeerAddr
peers
             TestTraceEvent extraState extraFlags extraPeers
_                                         -> Maybe (Set PeerAddr)
forall a. Maybe a
Nothing
      )

-- | filter big ledger peers
--
takeBigLedgerPeers
    :: (   Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn
        -> Set PeerAddr
       )
    -> Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn
    -> Set PeerAddr
takeBigLedgerPeers :: forall extraState extraFlags extraPeers peerconn.
(PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> Set PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
takeBigLedgerPeers PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f =
  \PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st -> PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st Set PeerAddr -> Set PeerAddr -> Set PeerAddr
forall a. Ord a => Set a -> Set a -> Set a
`Set.intersection` (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr
forall extraPeers peeraddr.
PublicRootPeers extraPeers peeraddr -> Set peeraddr
PublicRootPeers.getBigLedgerPeers (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr)
-> (PeerSelectionState
      extraState extraFlags extraPeers PeerAddr peerconn
    -> PublicRootPeers extraPeers PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> PublicRootPeers extraPeers PeerAddr
forall extraState extraFlags extraPeers peeraddr peerconn.
PeerSelectionState
  extraState extraFlags extraPeers peeraddr peerconn
-> PublicRootPeers extraPeers peeraddr
Governor.publicRootPeers) PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st

-- | filter out big ledger peers
--
dropBigLedgerPeers
    :: (Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> Set PeerAddr)
    ->  Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> Set PeerAddr
dropBigLedgerPeers :: forall extraState extraFlags extraPeers peerconn.
(PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> Set PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
dropBigLedgerPeers PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f =
  \PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st -> PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st Set PeerAddr -> Set PeerAddr -> Set PeerAddr
forall a. Ord a => Set a -> Set a -> Set a
Set.\\ (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr
forall extraPeers peeraddr.
PublicRootPeers extraPeers peeraddr -> Set peeraddr
PublicRootPeers.getBigLedgerPeers (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr)
-> (PeerSelectionState
      extraState extraFlags extraPeers PeerAddr peerconn
    -> PublicRootPeers extraPeers PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> PublicRootPeers extraPeers PeerAddr
forall extraState extraFlags extraPeers peeraddr peerconn.
PeerSelectionState
  extraState extraFlags extraPeers peeraddr peerconn
-> PublicRootPeers extraPeers peeraddr
Governor.publicRootPeers) PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st