{-# LANGUAGE DeriveFunctor  #-}
{-# LANGUAGE NamedFieldPuns #-}

{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Ouroboros.Network.PeerSelection.EstablishedPeers
  ( EstablishedPeers
  , empty
  , toMap
  , toSet
  , readyPeers
  , size
  , sizeReady
  , member
  , insert
  , delete
  , deletePeers
  , setCurrentTime
  , minActivateTime
  , setActivateTime
  , invariant
  ) where

import           Prelude

import           Data.Foldable (foldl')
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.OrdPSQ (OrdPSQ)
import qualified Data.OrdPSQ as PSQ
import           Data.Set (Set)
import qualified Data.Set as Set

import           Control.Exception (assert)
import           Control.Monad.Class.MonadTime


data EstablishedPeers peeraddr peerconn = EstablishedPeers {
    -- | Peers which are either ready to become active or are active.
    --
    EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers          :: !(Map peeraddr peerconn),

    -- | Peers which are not ready to become active.
    EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes :: !(OrdPSQ peeraddr Time ())
  }
  deriving (Int -> EstablishedPeers peeraddr peerconn -> ShowS
[EstablishedPeers peeraddr peerconn] -> ShowS
EstablishedPeers peeraddr peerconn -> String
(Int -> EstablishedPeers peeraddr peerconn -> ShowS)
-> (EstablishedPeers peeraddr peerconn -> String)
-> ([EstablishedPeers peeraddr peerconn] -> ShowS)
-> Show (EstablishedPeers peeraddr peerconn)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall peeraddr peerconn.
(Show peeraddr, Show peerconn) =>
Int -> EstablishedPeers peeraddr peerconn -> ShowS
forall peeraddr peerconn.
(Show peeraddr, Show peerconn) =>
[EstablishedPeers peeraddr peerconn] -> ShowS
forall peeraddr peerconn.
(Show peeraddr, Show peerconn) =>
EstablishedPeers peeraddr peerconn -> String
showList :: [EstablishedPeers peeraddr peerconn] -> ShowS
$cshowList :: forall peeraddr peerconn.
(Show peeraddr, Show peerconn) =>
[EstablishedPeers peeraddr peerconn] -> ShowS
show :: EstablishedPeers peeraddr peerconn -> String
$cshow :: forall peeraddr peerconn.
(Show peeraddr, Show peerconn) =>
EstablishedPeers peeraddr peerconn -> String
showsPrec :: Int -> EstablishedPeers peeraddr peerconn -> ShowS
$cshowsPrec :: forall peeraddr peerconn.
(Show peeraddr, Show peerconn) =>
Int -> EstablishedPeers peeraddr peerconn -> ShowS
Show, a -> EstablishedPeers peeraddr b -> EstablishedPeers peeraddr a
(a -> b)
-> EstablishedPeers peeraddr a -> EstablishedPeers peeraddr b
(forall a b.
 (a -> b)
 -> EstablishedPeers peeraddr a -> EstablishedPeers peeraddr b)
-> (forall a b.
    a -> EstablishedPeers peeraddr b -> EstablishedPeers peeraddr a)
-> Functor (EstablishedPeers peeraddr)
forall a b.
a -> EstablishedPeers peeraddr b -> EstablishedPeers peeraddr a
forall a b.
(a -> b)
-> EstablishedPeers peeraddr a -> EstablishedPeers peeraddr b
forall peeraddr a b.
a -> EstablishedPeers peeraddr b -> EstablishedPeers peeraddr a
forall peeraddr a b.
(a -> b)
-> EstablishedPeers peeraddr a -> EstablishedPeers peeraddr b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> EstablishedPeers peeraddr b -> EstablishedPeers peeraddr a
$c<$ :: forall peeraddr a b.
a -> EstablishedPeers peeraddr b -> EstablishedPeers peeraddr a
fmap :: (a -> b)
-> EstablishedPeers peeraddr a -> EstablishedPeers peeraddr b
$cfmap :: forall peeraddr a b.
(a -> b)
-> EstablishedPeers peeraddr a -> EstablishedPeers peeraddr b
Functor)


empty :: EstablishedPeers peeraddr perconn
empty :: EstablishedPeers peeraddr perconn
empty = Map peeraddr perconn
-> OrdPSQ peeraddr Time () -> EstablishedPeers peeraddr perconn
forall peeraddr peerconn.
Map peeraddr peerconn
-> OrdPSQ peeraddr Time () -> EstablishedPeers peeraddr peerconn
EstablishedPeers Map peeraddr perconn
forall k a. Map k a
Map.empty OrdPSQ peeraddr Time ()
forall k p v. OrdPSQ k p v
PSQ.empty


invariant :: Ord peeraddr
          => EstablishedPeers peeraddr peerconn
          -> Bool
invariant :: EstablishedPeers peeraddr peerconn -> Bool
invariant EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers, OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
     -- nextActivateTimes is a subset of allPeers
     [peeraddr] -> Set peeraddr
forall a. Ord a => [a] -> Set a
Set.fromList (OrdPSQ peeraddr Time () -> [peeraddr]
forall k p v. OrdPSQ k p v -> [k]
PSQ.keys OrdPSQ peeraddr Time ()
nextActivateTimes)
     Set peeraddr -> Set peeraddr -> Bool
forall a. Ord a => Set a -> Set a -> Bool
`Set.isSubsetOf`
     Map peeraddr peerconn -> Set peeraddr
forall k a. Map k a -> Set k
Map.keysSet Map peeraddr peerconn
allPeers


-- | /O(1)/
toMap :: EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
toMap :: EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
toMap = EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers

-- | /O(n)/
toSet :: EstablishedPeers peeraddr peerconn -> Set peeraddr
toSet :: EstablishedPeers peeraddr peerconn -> Set peeraddr
toSet = Map peeraddr peerconn -> Set peeraddr
forall k a. Map k a -> Set k
Map.keysSet (Map peeraddr peerconn -> Set peeraddr)
-> (EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn)
-> EstablishedPeers peeraddr peerconn
-> Set peeraddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers


-- | Map of established peers that are either active or ready to be promoted
-- to active.
--
-- /O(n log m), for n not-ready peers, and m established peers/
--
readyPeers :: Ord peeraddr
           => EstablishedPeers peeraddr peerconn
           -> Set peeraddr
readyPeers :: EstablishedPeers peeraddr peerconn -> Set peeraddr
readyPeers EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers, OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
    (peeraddr -> Time -> () -> Set peeraddr -> Set peeraddr)
-> Set peeraddr -> OrdPSQ peeraddr Time () -> Set peeraddr
forall k p v a. (k -> p -> v -> a -> a) -> a -> OrdPSQ k p v -> a
PSQ.fold'
      (\peeraddr
peeraddr Time
_ ()
_ -> peeraddr -> Set peeraddr -> Set peeraddr
forall a. Ord a => a -> Set a -> Set a
Set.delete peeraddr
peeraddr)
      (Map peeraddr peerconn -> Set peeraddr
forall k a. Map k a -> Set k
Map.keysSet Map peeraddr peerconn
allPeers)
      OrdPSQ peeraddr Time ()
nextActivateTimes


-- | The number of established peers. The size of 'allPeers'
--
-- /O(1)/
--
size :: EstablishedPeers peeraddr peerconn -> Int
size :: EstablishedPeers peeraddr peerconn -> Int
size EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers } = Map peeraddr peerconn -> Int
forall k a. Map k a -> Int
Map.size Map peeraddr peerconn
allPeers


-- | The number of ready peers. The size of 'readyPeers'
--
-- /O(1)/
--
sizeReady :: EstablishedPeers peeraddr peerconn -> Int
sizeReady :: EstablishedPeers peeraddr peerconn -> Int
sizeReady EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers, OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
    Map peeraddr peerconn -> Int
forall k a. Map k a -> Int
Map.size Map peeraddr peerconn
allPeers Int -> Int -> Int
forall a. Num a => a -> a -> a
- OrdPSQ peeraddr Time () -> Int
forall k p v. OrdPSQ k p v -> Int
PSQ.size OrdPSQ peeraddr Time ()
nextActivateTimes


member :: Ord peeraddr => peeraddr -> EstablishedPeers peeraddr peerconn -> Bool
member :: peeraddr -> EstablishedPeers peeraddr peerconn -> Bool
member peeraddr
peeraddr = peeraddr -> Map peeraddr peerconn -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member peeraddr
peeraddr (Map peeraddr peerconn -> Bool)
-> (EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn)
-> EstablishedPeers peeraddr peerconn
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers


-- | Insert a peer into 'EstablishedPeers'.
--
insert :: Ord peeraddr
       => peeraddr
       -> peerconn
       -> EstablishedPeers peeraddr peerconn
       -> EstablishedPeers peeraddr peerconn
insert :: peeraddr
-> peerconn
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
insert peeraddr
peeraddr peerconn
peerconn ep :: EstablishedPeers peeraddr peerconn
ep@EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers } =
  EstablishedPeers peeraddr peerconn
ep { allPeers :: Map peeraddr peerconn
allPeers          = peeraddr
-> peerconn -> Map peeraddr peerconn -> Map peeraddr peerconn
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert peeraddr
peeraddr peerconn
peerconn Map peeraddr peerconn
allPeers }

delete :: Ord peeraddr
       => peeraddr
       -> EstablishedPeers peeraddr peerconn
       -> EstablishedPeers peeraddr peerconn
delete :: peeraddr
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
delete peeraddr
peeraddr es :: EstablishedPeers peeraddr peerconn
es@EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers, OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
    EstablishedPeers peeraddr peerconn
es { allPeers :: Map peeraddr peerconn
allPeers          = peeraddr -> Map peeraddr peerconn -> Map peeraddr peerconn
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete peeraddr
peeraddr Map peeraddr peerconn
allPeers,
         nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes = peeraddr -> OrdPSQ peeraddr Time () -> OrdPSQ peeraddr Time ()
forall k p v. (Ord k, Ord p) => k -> OrdPSQ k p v -> OrdPSQ k p v
PSQ.delete peeraddr
peeraddr OrdPSQ peeraddr Time ()
nextActivateTimes }



-- | Bulk delete of peers from 'EstablishedPeers.
--
deletePeers :: Ord peeraddr
            => Set peeraddr
            -> EstablishedPeers peeraddr peerconn
            -> EstablishedPeers peeraddr peerconn
deletePeers :: Set peeraddr
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
deletePeers Set peeraddr
peeraddrs es :: EstablishedPeers peeraddr peerconn
es@EstablishedPeers { Map peeraddr peerconn
allPeers :: Map peeraddr peerconn
allPeers :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> Map peeraddr peerconn
allPeers, OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
    EstablishedPeers peeraddr peerconn
es { allPeers :: Map peeraddr peerconn
allPeers          = (Map peeraddr peerconn -> peeraddr -> Map peeraddr peerconn)
-> Map peeraddr peerconn -> Set peeraddr -> Map peeraddr peerconn
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((peeraddr -> Map peeraddr peerconn -> Map peeraddr peerconn)
-> Map peeraddr peerconn -> peeraddr -> Map peeraddr peerconn
forall a b c. (a -> b -> c) -> b -> a -> c
flip peeraddr -> Map peeraddr peerconn -> Map peeraddr peerconn
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete) Map peeraddr peerconn
allPeers  Set peeraddr
peeraddrs,
         nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes = (OrdPSQ peeraddr Time () -> peeraddr -> OrdPSQ peeraddr Time ())
-> OrdPSQ peeraddr Time ()
-> Set peeraddr
-> OrdPSQ peeraddr Time ()
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((peeraddr -> OrdPSQ peeraddr Time () -> OrdPSQ peeraddr Time ())
-> OrdPSQ peeraddr Time () -> peeraddr -> OrdPSQ peeraddr Time ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip peeraddr -> OrdPSQ peeraddr Time () -> OrdPSQ peeraddr Time ()
forall k p v. (Ord k, Ord p) => k -> OrdPSQ k p v -> OrdPSQ k p v
PSQ.delete) OrdPSQ peeraddr Time ()
nextActivateTimes Set peeraddr
peeraddrs }


--
-- Time managment
--

setCurrentTime :: Ord peeraddr
               => Time
               -> EstablishedPeers peeraddr peerconn
               -> EstablishedPeers peeraddr peerconn
setCurrentTime :: Time
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
setCurrentTime Time
now ep :: EstablishedPeers peeraddr peerconn
ep@EstablishedPeers { OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
    let ep' :: EstablishedPeers peeraddr peerconn
ep' = EstablishedPeers peeraddr peerconn
ep { nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes = OrdPSQ peeraddr Time ()
nextActivateTimes' }
    in Bool
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (EstablishedPeers peeraddr peerconn -> Bool
forall peeraddr peerconn.
Ord peeraddr =>
EstablishedPeers peeraddr peerconn -> Bool
invariant EstablishedPeers peeraddr peerconn
ep') EstablishedPeers peeraddr peerconn
ep'
  where
    ([(peeraddr, Time, ())]
_, OrdPSQ peeraddr Time ()
nextActivateTimes') = Time
-> OrdPSQ peeraddr Time ()
-> ([(peeraddr, Time, ())], OrdPSQ peeraddr Time ())
forall k p v.
(Ord k, Ord p) =>
p -> OrdPSQ k p v -> ([(k, p, v)], OrdPSQ k p v)
PSQ.atMostView Time
now OrdPSQ peeraddr Time ()
nextActivateTimes


minActivateTime :: Ord peeraddr
                => EstablishedPeers peeraddr peerconn
                -> Maybe Time
minActivateTime :: EstablishedPeers peeraddr peerconn -> Maybe Time
minActivateTime EstablishedPeers { OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes }
  | Just (peeraddr
_k, Time
t, ()
_, OrdPSQ peeraddr Time ()
_psq) <- OrdPSQ peeraddr Time ()
-> Maybe (peeraddr, Time, (), OrdPSQ peeraddr Time ())
forall k p v.
(Ord k, Ord p) =>
OrdPSQ k p v -> Maybe (k, p, v, OrdPSQ k p v)
PSQ.minView OrdPSQ peeraddr Time ()
nextActivateTimes
  = Time -> Maybe Time
forall a. a -> Maybe a
Just Time
t

  | Bool
otherwise
  = Maybe Time
forall a. Maybe a
Nothing


setActivateTime :: Ord peeraddr
                => Set peeraddr
                -> Time
                -> EstablishedPeers peeraddr peerconn
                -> EstablishedPeers peeraddr peerconn
setActivateTime :: Set peeraddr
-> Time
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
setActivateTime Set peeraddr
peeraddrs Time
_time EstablishedPeers peeraddr peerconn
ep | Set peeraddr -> Bool
forall a. Set a -> Bool
Set.null Set peeraddr
peeraddrs = EstablishedPeers peeraddr peerconn
ep
setActivateTime Set peeraddr
peeraddrs Time
time  ep :: EstablishedPeers peeraddr peerconn
ep@EstablishedPeers { OrdPSQ peeraddr Time ()
nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes :: forall peeraddr peerconn.
EstablishedPeers peeraddr peerconn -> OrdPSQ peeraddr Time ()
nextActivateTimes } =
    let ep' :: EstablishedPeers peeraddr peerconn
ep' = EstablishedPeers peeraddr peerconn
ep { nextActivateTimes :: OrdPSQ peeraddr Time ()
nextActivateTimes = (OrdPSQ peeraddr Time () -> peeraddr -> OrdPSQ peeraddr Time ())
-> OrdPSQ peeraddr Time ()
-> Set peeraddr
-> OrdPSQ peeraddr Time ()
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\OrdPSQ peeraddr Time ()
psq peeraddr
peeraddr -> peeraddr
-> Time -> () -> OrdPSQ peeraddr Time () -> OrdPSQ peeraddr Time ()
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
PSQ.insert peeraddr
peeraddr Time
time () OrdPSQ peeraddr Time ()
psq)
                                              OrdPSQ peeraddr Time ()
nextActivateTimes
                                              Set peeraddr
peeraddrs
                 }
    in   Bool
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
forall a. (?callStack::CallStack) => Bool -> a -> a
assert ((peeraddr -> Bool) -> Set peeraddr -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool
not (Bool -> Bool) -> (peeraddr -> Bool) -> peeraddr -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (peeraddr -> Set peeraddr -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` EstablishedPeers peeraddr peerconn -> Set peeraddr
forall peeraddr peerconn.
Ord peeraddr =>
EstablishedPeers peeraddr peerconn -> Set peeraddr
readyPeers EstablishedPeers peeraddr peerconn
ep')) Set peeraddr
peeraddrs)
       (EstablishedPeers peeraddr peerconn
 -> EstablishedPeers peeraddr peerconn)
-> (EstablishedPeers peeraddr peerconn
    -> EstablishedPeers peeraddr peerconn)
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (EstablishedPeers peeraddr peerconn -> Bool
forall peeraddr peerconn.
Ord peeraddr =>
EstablishedPeers peeraddr peerconn -> Bool
invariant EstablishedPeers peeraddr peerconn
ep')
       (EstablishedPeers peeraddr peerconn
 -> EstablishedPeers peeraddr peerconn)
-> EstablishedPeers peeraddr peerconn
-> EstablishedPeers peeraddr peerconn
forall a b. (a -> b) -> a -> b
$ EstablishedPeers peeraddr peerconn
ep'