{-# LANGUAGE BangPatterns    #-}
{-# LANGUAGE NamedFieldPuns  #-}
{-# LANGUAGE RecordWildCards #-}

module Ouroboros.Network.PeerSelection.LocalRootPeers
  ( -- * Types
    LocalRootPeers (..)
    -- Export constructors for defining tests.
  , invariant
    -- * Basic operations
  , empty
  , null
  , size
  , member
  , target
  , fromGroups
  , toGroups
  , toGroupSets
  , toMap
  , keysSet
    -- * Special operations
  , clampToLimit
  ) where

import           Prelude hiding (null)

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Set (Set)
import qualified Data.Set as Set

import           Ouroboros.Network.PeerSelection.Types


---------------------------------------
-- Local root peer set representation
--

data LocalRootPeers peeraddr =
     LocalRootPeers
       -- We use two partial & overlapping representations:

       -- The collection of all the peers, with the associated PeerAdvertise
       (Map peeraddr PeerAdvertise)

       -- The groups, but without the associated PeerAdvertise
       [(Int, Set peeraddr)]
  deriving LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool
(LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool)
-> (LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool)
-> Eq (LocalRootPeers peeraddr)
forall peeraddr.
Eq peeraddr =>
LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool
$c/= :: forall peeraddr.
Eq peeraddr =>
LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool
== :: LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool
$c== :: forall peeraddr.
Eq peeraddr =>
LocalRootPeers peeraddr -> LocalRootPeers peeraddr -> Bool
Eq

-- It is an abstract type, so the derived Show is unhelpful, e.g. for replaying
-- test cases.
--
instance (Show peeraddr, Ord peeraddr) => Show (LocalRootPeers peeraddr) where
  show :: LocalRootPeers peeraddr -> String
show LocalRootPeers peeraddr
lrps = String
"fromGroups " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [(Int, Map peeraddr PeerAdvertise)] -> String
forall a. Show a => a -> String
show (LocalRootPeers peeraddr -> [(Int, Map peeraddr PeerAdvertise)]
forall peeraddr.
Ord peeraddr =>
LocalRootPeers peeraddr -> [(Int, Map peeraddr PeerAdvertise)]
toGroups LocalRootPeers peeraddr
lrps)

invariant :: Ord peeraddr => LocalRootPeers peeraddr -> Bool
invariant :: LocalRootPeers peeraddr -> Bool
invariant (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
gs) =

    -- The overlapping representations must be consistent
    [Set peeraddr] -> Set peeraddr
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [ Set peeraddr
g | (Int
_, Set peeraddr
g) <- [(Int, Set peeraddr)]
gs ] Set peeraddr -> Set peeraddr -> Bool
forall a. Eq a => a -> a -> Bool
== Map peeraddr PeerAdvertise -> Set peeraddr
forall k a. Map k a -> Set k
Map.keysSet Map peeraddr PeerAdvertise
m

    -- The localRootPeers groups must not overlap with each other
 Bool -> Bool -> Bool
&& Map peeraddr PeerAdvertise -> Int
forall k a. Map k a -> Int
Map.size Map peeraddr PeerAdvertise
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Set peeraddr -> Int
forall a. Set a -> Int
Set.size Set peeraddr
g | (Int
_, Set peeraddr
g) <- [(Int, Set peeraddr)]
gs ]

    -- Individual group targets must be greater than zero and achievable given
    -- the group sizes.
 Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
t Bool -> Bool -> Bool
&& Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Set peeraddr -> Int
forall a. Set a -> Int
Set.size Set peeraddr
g | (Int
t, Set peeraddr
g) <- [(Int, Set peeraddr)]
gs ]


empty :: LocalRootPeers peeraddr
empty :: LocalRootPeers peeraddr
empty = Map peeraddr PeerAdvertise
-> [(Int, Set peeraddr)] -> LocalRootPeers peeraddr
forall peeraddr.
Map peeraddr PeerAdvertise
-> [(Int, Set peeraddr)] -> LocalRootPeers peeraddr
LocalRootPeers Map peeraddr PeerAdvertise
forall k a. Map k a
Map.empty []

null :: LocalRootPeers peeraddr -> Bool
null :: LocalRootPeers peeraddr -> Bool
null (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
_) = Map peeraddr PeerAdvertise -> Bool
forall k a. Map k a -> Bool
Map.null Map peeraddr PeerAdvertise
m

size :: LocalRootPeers peeraddr -> Int
size :: LocalRootPeers peeraddr -> Int
size (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
_) = Map peeraddr PeerAdvertise -> Int
forall k a. Map k a -> Int
Map.size Map peeraddr PeerAdvertise
m

member :: Ord peeraddr => peeraddr -> LocalRootPeers peeraddr -> Bool
member :: peeraddr -> LocalRootPeers peeraddr -> Bool
member peeraddr
p (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
_) = peeraddr -> Map peeraddr PeerAdvertise -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member peeraddr
p Map peeraddr PeerAdvertise
m

target :: LocalRootPeers peeraddr -> Int
target :: LocalRootPeers peeraddr -> Int
target (LocalRootPeers Map peeraddr PeerAdvertise
_ [(Int, Set peeraddr)]
gs) = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Int
t | (Int
t, Set peeraddr
_) <- [(Int, Set peeraddr)]
gs ]

toMap :: LocalRootPeers peeraddr -> Map peeraddr PeerAdvertise
toMap :: LocalRootPeers peeraddr -> Map peeraddr PeerAdvertise
toMap (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
_) = Map peeraddr PeerAdvertise
m

keysSet :: LocalRootPeers peeraddr -> Set peeraddr
keysSet :: LocalRootPeers peeraddr -> Set peeraddr
keysSet (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
_) = Map peeraddr PeerAdvertise -> Set peeraddr
forall k a. Map k a -> Set k
Map.keysSet Map peeraddr PeerAdvertise
m

toGroupSets :: LocalRootPeers peeraddr -> [(Int, Set peeraddr)]
toGroupSets :: LocalRootPeers peeraddr -> [(Int, Set peeraddr)]
toGroupSets (LocalRootPeers Map peeraddr PeerAdvertise
_ [(Int, Set peeraddr)]
gs) = [(Int, Set peeraddr)]
gs


-- | The local root peers info has some invariants that are not directly
-- enforced in the types, and the config comes from an external source. Of
-- course it's good to validate that at source, but here we need to not fail
-- if we're given imperfect data.
--
-- So what we do is bash it until it is valid. We don't need to be too careful
-- about how we do it, it's ok to be brutal. We should however make sure we
-- trace a warning about dodgy config.
--
fromGroups :: Ord peeraddr
           => [(Int, Map peeraddr PeerAdvertise)]
           -> LocalRootPeers peeraddr
fromGroups :: [(Int, Map peeraddr PeerAdvertise)] -> LocalRootPeers peeraddr
fromGroups =
    (\[(Int, Map peeraddr PeerAdvertise)]
gs -> let m' :: Map peeraddr PeerAdvertise
m'  = [Map peeraddr PeerAdvertise] -> Map peeraddr PeerAdvertise
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
Map.unions [ Map peeraddr PeerAdvertise
g | (Int
_, Map peeraddr PeerAdvertise
g) <- [(Int, Map peeraddr PeerAdvertise)]
gs ]
                gs' :: [(Int, Set peeraddr)]
gs' = [ (Int
t, Map peeraddr PeerAdvertise -> Set peeraddr
forall k a. Map k a -> Set k
Map.keysSet Map peeraddr PeerAdvertise
g) | (Int
t, Map peeraddr PeerAdvertise
g) <- [(Int, Map peeraddr PeerAdvertise)]
gs ]
             in Map peeraddr PeerAdvertise
-> [(Int, Set peeraddr)] -> LocalRootPeers peeraddr
forall peeraddr.
Map peeraddr PeerAdvertise
-> [(Int, Set peeraddr)] -> LocalRootPeers peeraddr
LocalRootPeers Map peeraddr PeerAdvertise
m' [(Int, Set peeraddr)]
gs')
  ([(Int, Map peeraddr PeerAdvertise)] -> LocalRootPeers peeraddr)
-> ([(Int, Map peeraddr PeerAdvertise)]
    -> [(Int, Map peeraddr PeerAdvertise)])
-> [(Int, Map peeraddr PeerAdvertise)]
-> LocalRootPeers peeraddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set peeraddr
-> [(Int, Map peeraddr PeerAdvertise)]
-> [(Int, Map peeraddr PeerAdvertise)]
forall k a. Ord k => Set k -> [(Int, Map k a)] -> [(Int, Map k a)]
establishStructureInvariant Set peeraddr
forall a. Set a
Set.empty
  where
    -- The groups must not overlap; have achievable targets; and be non-empty.
    establishStructureInvariant :: Set k -> [(Int, Map k a)] -> [(Int, Map k a)]
establishStructureInvariant !Set k
_ [] = []
    establishStructureInvariant !Set k
acc ((Int
t, Map k a
g): [(Int, Map k a)]
gs)
      | Int
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0    = (Int
t', Map k a
g') (Int, Map k a) -> [(Int, Map k a)] -> [(Int, Map k a)]
forall a. a -> [a] -> [a]
: Set k -> [(Int, Map k a)] -> [(Int, Map k a)]
establishStructureInvariant Set k
acc' [(Int, Map k a)]
gs
      | Bool
otherwise =            Set k -> [(Int, Map k a)] -> [(Int, Map k a)]
establishStructureInvariant Set k
acc' [(Int, Map k a)]
gs
      where
        !g' :: Map k a
g'   = Map k a
g Map k a -> Set k -> Map k a
forall k a. Ord k => Map k a -> Set k -> Map k a
`Map.withoutKeys` Set k
acc
        !t' :: Int
t'   = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
t (Map k a -> Int
forall k a. Map k a -> Int
Map.size Map k a
g')
        !acc' :: Set k
acc' = Set k
acc Set k -> Set k -> Set k
forall a. Semigroup a => a -> a -> a
<> Map k a -> Set k
forall k a. Map k a -> Set k
Map.keysSet Map k a
g

-- | Inverse of 'fromGroups', for the subset of inputs to 'fromGroups' that
-- satisfy the invariant.
--
toGroups :: Ord peeraddr
         => LocalRootPeers peeraddr
         -> [(Int, Map peeraddr PeerAdvertise)]
toGroups :: LocalRootPeers peeraddr -> [(Int, Map peeraddr PeerAdvertise)]
toGroups (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
gs) =
    [ (Int
t, (peeraddr -> PeerAdvertise)
-> Set peeraddr -> Map peeraddr PeerAdvertise
forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet (Map peeraddr PeerAdvertise
m Map peeraddr PeerAdvertise -> peeraddr -> PeerAdvertise
forall k a. Ord k => Map k a -> k -> a
Map.!) Set peeraddr
g)
    | (Int
t, Set peeraddr
g) <- [(Int, Set peeraddr)]
gs ]


-- | Limit the size of the root peers collection to fit within given bounds.
--
-- The governor needs to be able to do this to enforce its invariant that:
--
-- > LocalRootPeers.size localRootPeers <= targetNumberOfKnownPeers
--
-- It needs to be able to /establish/ that invariant given arbitrary
-- configuration for local root peers. It makes sense to do it this way rather
-- than just enforce that local root peers config fits the invariant because
-- the invariant depends on both the targets and the local root peers config
-- and these can both vary dynamically and independently.
--
-- It is unlikely in practice that there are so many local root peers
-- configured that it goes over this targets, so it's ok to resolve it pretty
-- arbitrarily. We just take the local roots in left to right order up to the
-- limit. So we have the property that
--
-- > LocalRootPeers.size (LocalRootPeers.clampToLimit sz lrps)
-- >  == min sz (LocalRootPeers.size lrps)
--
clampToLimit :: Ord peeraddr
             => Int -- ^ The limit on the total number of local peers
             -> LocalRootPeers peeraddr
             -> LocalRootPeers peeraddr
clampToLimit :: Int -> LocalRootPeers peeraddr -> LocalRootPeers peeraddr
clampToLimit Int
totalLimit (LocalRootPeers Map peeraddr PeerAdvertise
m [(Int, Set peeraddr)]
gs0) =
    let gs' :: [(Int, Set peeraddr)]
gs' = Int -> [(Int, Set peeraddr)] -> [(Int, Set peeraddr)]
forall a. Int -> [(Int, Set a)] -> [(Int, Set a)]
limitTotalSize Int
0 [(Int, Set peeraddr)]
gs0
        m' :: Map peeraddr PeerAdvertise
m'  = Map peeraddr PeerAdvertise
m Map peeraddr PeerAdvertise
-> Set peeraddr -> Map peeraddr PeerAdvertise
forall k a. Ord k => Map k a -> Set k -> Map k a
`Map.restrictKeys` [Set peeraddr] -> Set peeraddr
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [ Set peeraddr
g | (Int
_, Set peeraddr
g) <- [(Int, Set peeraddr)]
gs' ]
     in Map peeraddr PeerAdvertise
-> [(Int, Set peeraddr)] -> LocalRootPeers peeraddr
forall peeraddr.
Map peeraddr PeerAdvertise
-> [(Int, Set peeraddr)] -> LocalRootPeers peeraddr
LocalRootPeers Map peeraddr PeerAdvertise
m' [(Int, Set peeraddr)]
gs'

  where
    limitTotalSize :: Int -> [(Int, Set a)] -> [(Int, Set a)]
limitTotalSize !Int
_ [] = []
    limitTotalSize !Int
n ((Int
t, Set a
g) : [(Int, Set a)]
gs)

        -- No space at all!
      | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
totalLimit
      = []

        -- It fits entirely!
      | let n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
Set.size Set a
g
      , Int
n' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
totalLimit
      = (Int
t, Set a
g) (Int, Set a) -> [(Int, Set a)] -> [(Int, Set a)]
forall a. a -> [a] -> [a]
: Int -> [(Int, Set a)] -> [(Int, Set a)]
limitTotalSize Int
n' [(Int, Set a)]
gs

        -- We can fit a bit more if we chop it up!
      | Bool
otherwise
      , let !g' :: Set a
g' = Int -> Set a -> Set a
forall a. Int -> Set a -> Set a
Set.take (Int
totalLimit Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Set a
g
            !t' :: Int
t' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
t (Set a -> Int
forall a. Set a -> Int
Set.size Set a
g')
      = (Int
t', Set a
g') (Int, Set a) -> [(Int, Set a)] -> [(Int, Set a)]
forall a. a -> [a] -> [a]
: []