module Test.ThreadNet.Util.NodeRestarts (
    NodeRestart (..)
  , NodeRestarts (..)
  , genNodeRestarts
  , noRestarts
  , shrinkNodeRestarts
  ) where

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import           Data.Traversable (forM)
import           Ouroboros.Consensus.Block
import           Ouroboros.Consensus.NodeId
import           Ouroboros.Consensus.Util.Condense
import           Test.QuickCheck
import           Test.ThreadNet.Util.NodeJoinPlan
import           Test.Util.Slots (NumSlots (..))

data NodeRestart
  = NodeRekey
    -- ^ restart the node with a fresh operational key and immediately emit a
    -- delegation certificate transaction
  | NodeRestart
    -- ^ restart the node without rekeying
  deriving (NodeRestart -> NodeRestart -> Bool
(NodeRestart -> NodeRestart -> Bool)
-> (NodeRestart -> NodeRestart -> Bool) -> Eq NodeRestart
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NodeRestart -> NodeRestart -> Bool
== :: NodeRestart -> NodeRestart -> Bool
$c/= :: NodeRestart -> NodeRestart -> Bool
/= :: NodeRestart -> NodeRestart -> Bool
Eq, Eq NodeRestart
Eq NodeRestart =>
(NodeRestart -> NodeRestart -> Ordering)
-> (NodeRestart -> NodeRestart -> Bool)
-> (NodeRestart -> NodeRestart -> Bool)
-> (NodeRestart -> NodeRestart -> Bool)
-> (NodeRestart -> NodeRestart -> Bool)
-> (NodeRestart -> NodeRestart -> NodeRestart)
-> (NodeRestart -> NodeRestart -> NodeRestart)
-> Ord NodeRestart
NodeRestart -> NodeRestart -> Bool
NodeRestart -> NodeRestart -> Ordering
NodeRestart -> NodeRestart -> NodeRestart
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: NodeRestart -> NodeRestart -> Ordering
compare :: NodeRestart -> NodeRestart -> Ordering
$c< :: NodeRestart -> NodeRestart -> Bool
< :: NodeRestart -> NodeRestart -> Bool
$c<= :: NodeRestart -> NodeRestart -> Bool
<= :: NodeRestart -> NodeRestart -> Bool
$c> :: NodeRestart -> NodeRestart -> Bool
> :: NodeRestart -> NodeRestart -> Bool
$c>= :: NodeRestart -> NodeRestart -> Bool
>= :: NodeRestart -> NodeRestart -> Bool
$cmax :: NodeRestart -> NodeRestart -> NodeRestart
max :: NodeRestart -> NodeRestart -> NodeRestart
$cmin :: NodeRestart -> NodeRestart -> NodeRestart
min :: NodeRestart -> NodeRestart -> NodeRestart
Ord, Int -> NodeRestart -> ShowS
[NodeRestart] -> ShowS
NodeRestart -> String
(Int -> NodeRestart -> ShowS)
-> (NodeRestart -> String)
-> ([NodeRestart] -> ShowS)
-> Show NodeRestart
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NodeRestart -> ShowS
showsPrec :: Int -> NodeRestart -> ShowS
$cshow :: NodeRestart -> String
show :: NodeRestart -> String
$cshowList :: [NodeRestart] -> ShowS
showList :: [NodeRestart] -> ShowS
Show)

instance Condense NodeRestart where
  condense :: NodeRestart -> String
condense = NodeRestart -> String
forall a. Show a => a -> String
show

-- | Which nodes are scheduled to restart in each slot
--
-- INVARIANT no element 'Map' is empty
--
newtype NodeRestarts = NodeRestarts (Map SlotNo (Map CoreNodeId NodeRestart))
  deriving (NodeRestarts -> NodeRestarts -> Bool
(NodeRestarts -> NodeRestarts -> Bool)
-> (NodeRestarts -> NodeRestarts -> Bool) -> Eq NodeRestarts
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NodeRestarts -> NodeRestarts -> Bool
== :: NodeRestarts -> NodeRestarts -> Bool
$c/= :: NodeRestarts -> NodeRestarts -> Bool
/= :: NodeRestarts -> NodeRestarts -> Bool
Eq, Int -> NodeRestarts -> ShowS
[NodeRestarts] -> ShowS
NodeRestarts -> String
(Int -> NodeRestarts -> ShowS)
-> (NodeRestarts -> String)
-> ([NodeRestarts] -> ShowS)
-> Show NodeRestarts
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NodeRestarts -> ShowS
showsPrec :: Int -> NodeRestarts -> ShowS
$cshow :: NodeRestarts -> String
show :: NodeRestarts -> String
$cshowList :: [NodeRestarts] -> ShowS
showList :: [NodeRestarts] -> ShowS
Show)

instance Condense NodeRestarts where
  condense :: NodeRestarts -> String
condense (NodeRestarts Map SlotNo (Map CoreNodeId NodeRestart)
m) = [(SlotNo, [(NodeId, NodeRestart)])] -> String
forall a. Condense a => a -> String
condense
      [ (,) SlotNo
slot ([(NodeId, NodeRestart)] -> (SlotNo, [(NodeId, NodeRestart)]))
-> [(NodeId, NodeRestart)] -> (SlotNo, [(NodeId, NodeRestart)])
forall a b. (a -> b) -> a -> b
$
        [ (CoreNodeId -> NodeId
fromCoreNodeId CoreNodeId
cid, NodeRestart
r)
        | (CoreNodeId
cid, NodeRestart
r) <- Map CoreNodeId NodeRestart -> [(CoreNodeId, NodeRestart)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map CoreNodeId NodeRestart
m'
        ]
      | (SlotNo
slot, Map CoreNodeId NodeRestart
m') <- Map SlotNo (Map CoreNodeId NodeRestart)
-> [(SlotNo, Map CoreNodeId NodeRestart)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map SlotNo (Map CoreNodeId NodeRestart)
m
      ]

noRestarts :: NodeRestarts
noRestarts :: NodeRestarts
noRestarts = Map SlotNo (Map CoreNodeId NodeRestart) -> NodeRestarts
NodeRestarts Map SlotNo (Map CoreNodeId NodeRestart)
forall k a. Map k a
Map.empty

-- | Generate a valid 'NodeRestarts'
--
-- POSTCONDITION will not restart a node before it joins
--
-- POSTCONDITION will not restart a node when it's scheduled to lead
-- *according* *to* *round-robin*
--
-- POSTCONDITION will not simultaneously restart all nodes that have previously
-- joined
--
genNodeRestarts :: NodeJoinPlan -> NumSlots -> Gen NodeRestarts
genNodeRestarts :: NodeJoinPlan -> NumSlots -> Gen NodeRestarts
genNodeRestarts (NodeJoinPlan Map CoreNodeId SlotNo
m) (NumSlots Word64
t)
  | Word64
t Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
1     = NodeRestarts -> Gen NodeRestarts
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodeRestarts
noRestarts
  | Bool
otherwise =
  ([(SlotNo, Map CoreNodeId NodeRestart)] -> NodeRestarts)
-> Gen [(SlotNo, Map CoreNodeId NodeRestart)] -> Gen NodeRestarts
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map SlotNo (Map CoreNodeId NodeRestart) -> NodeRestarts
NodeRestarts (Map SlotNo (Map CoreNodeId NodeRestart) -> NodeRestarts)
-> ([(SlotNo, Map CoreNodeId NodeRestart)]
    -> Map SlotNo (Map CoreNodeId NodeRestart))
-> [(SlotNo, Map CoreNodeId NodeRestart)]
-> NodeRestarts
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map CoreNodeId NodeRestart -> Bool)
-> Map SlotNo (Map CoreNodeId NodeRestart)
-> Map SlotNo (Map CoreNodeId NodeRestart)
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Bool -> Bool
not (Bool -> Bool)
-> (Map CoreNodeId NodeRestart -> Bool)
-> Map CoreNodeId NodeRestart
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map CoreNodeId NodeRestart -> Bool
forall k a. Map k a -> Bool
Map.null) (Map SlotNo (Map CoreNodeId NodeRestart)
 -> Map SlotNo (Map CoreNodeId NodeRestart))
-> ([(SlotNo, Map CoreNodeId NodeRestart)]
    -> Map SlotNo (Map CoreNodeId NodeRestart))
-> [(SlotNo, Map CoreNodeId NodeRestart)]
-> Map SlotNo (Map CoreNodeId NodeRestart)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(SlotNo, Map CoreNodeId NodeRestart)]
-> Map SlotNo (Map CoreNodeId NodeRestart)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList) (Gen [(SlotNo, Map CoreNodeId NodeRestart)] -> Gen NodeRestarts)
-> Gen [(SlotNo, Map CoreNodeId NodeRestart)] -> Gen NodeRestarts
forall a b. (a -> b) -> a -> b
$ do
    [SlotNo]
ss <- [SlotNo] -> Gen [SlotNo]
forall a. [a] -> Gen [a]
sublistOf [SlotNo
0 .. Word64 -> SlotNo
SlotNo (Word64
t Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)]
    [SlotNo]
-> (SlotNo -> Gen (SlotNo, Map CoreNodeId NodeRestart))
-> Gen [(SlotNo, Map CoreNodeId NodeRestart)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [SlotNo]
ss ((SlotNo -> Gen (SlotNo, Map CoreNodeId NodeRestart))
 -> Gen [(SlotNo, Map CoreNodeId NodeRestart)])
-> (SlotNo -> Gen (SlotNo, Map CoreNodeId NodeRestart))
-> Gen [(SlotNo, Map CoreNodeId NodeRestart)]
forall a b. (a -> b) -> a -> b
$ \SlotNo
s ->
      (Map CoreNodeId NodeRestart
 -> (SlotNo, Map CoreNodeId NodeRestart))
-> Gen (Map CoreNodeId NodeRestart)
-> Gen (SlotNo, Map CoreNodeId NodeRestart)
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) SlotNo
s) (Gen (Map CoreNodeId NodeRestart)
 -> Gen (SlotNo, Map CoreNodeId NodeRestart))
-> Gen (Map CoreNodeId NodeRestart)
-> Gen (SlotNo, Map CoreNodeId NodeRestart)
forall a b. (a -> b) -> a -> b
$
      let alreadyJoined :: Set CoreNodeId
alreadyJoined = Map CoreNodeId SlotNo -> Set CoreNodeId
forall k a. Map k a -> Set k
Map.keysSet (Map CoreNodeId SlotNo -> Set CoreNodeId)
-> Map CoreNodeId SlotNo -> Set CoreNodeId
forall a b. (a -> b) -> a -> b
$ (SlotNo -> Bool) -> Map CoreNodeId SlotNo -> Map CoreNodeId SlotNo
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (SlotNo -> SlotNo -> Bool
forall a. Ord a => a -> a -> Bool
< SlotNo
s) Map CoreNodeId SlotNo
m
          keepSome :: Gen (Map CoreNodeId a) -> Gen (Map CoreNodeId a)
keepSome
            | Set CoreNodeId -> Bool
forall a. Set a -> Bool
Set.null Set CoreNodeId
alreadyJoined = Gen (Map CoreNodeId a) -> Gen (Map CoreNodeId a)
forall a. a -> a
id
            | Bool
otherwise              =
              (Gen (Map CoreNodeId a)
-> (Map CoreNodeId a -> Bool) -> Gen (Map CoreNodeId a)
forall a. Gen a -> (a -> Bool) -> Gen a
`suchThat` \Map CoreNodeId a
x -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Set CoreNodeId
alreadyJoined Set CoreNodeId -> Set CoreNodeId -> Bool
forall a. Ord a => Set a -> Set a -> Bool
`Set.isSubsetOf` Map CoreNodeId a -> Set CoreNodeId
forall k a. Map k a -> Set k
Map.keysSet Map CoreNodeId a
x)
          candidates :: Map CoreNodeId SlotNo
candidates = (CoreNodeId -> SlotNo -> Bool)
-> Map CoreNodeId SlotNo -> Map CoreNodeId SlotNo
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (SlotNo -> CoreNodeId -> SlotNo -> Bool
canRestartIn SlotNo
s) Map CoreNodeId SlotNo
m
      in
      Gen (Map CoreNodeId NodeRestart)
-> Gen (Map CoreNodeId NodeRestart)
forall {a}. Gen (Map CoreNodeId a) -> Gen (Map CoreNodeId a)
keepSome (Gen (Map CoreNodeId NodeRestart)
 -> Gen (Map CoreNodeId NodeRestart))
-> Gen (Map CoreNodeId NodeRestart)
-> Gen (Map CoreNodeId NodeRestart)
forall a b. (a -> b) -> a -> b
$
      if Map CoreNodeId SlotNo -> Bool
forall k a. Map k a -> Bool
Map.null Map CoreNodeId SlotNo
candidates
      then Map CoreNodeId NodeRestart -> Gen (Map CoreNodeId NodeRestart)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map CoreNodeId NodeRestart
forall k a. Map k a
Map.empty
      else ([CoreNodeId] -> Map CoreNodeId NodeRestart)
-> Gen [CoreNodeId] -> Gen (Map CoreNodeId NodeRestart)
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(CoreNodeId, NodeRestart)] -> Map CoreNodeId NodeRestart
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(CoreNodeId, NodeRestart)] -> Map CoreNodeId NodeRestart)
-> ([CoreNodeId] -> [(CoreNodeId, NodeRestart)])
-> [CoreNodeId]
-> Map CoreNodeId NodeRestart
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CoreNodeId -> (CoreNodeId, NodeRestart))
-> [CoreNodeId] -> [(CoreNodeId, NodeRestart)]
forall a b. (a -> b) -> [a] -> [b]
map ((CoreNodeId -> NodeRestart -> (CoreNodeId, NodeRestart))
-> NodeRestart -> CoreNodeId -> (CoreNodeId, NodeRestart)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) NodeRestart
NodeRestart)) (Gen [CoreNodeId] -> Gen (Map CoreNodeId NodeRestart))
-> Gen [CoreNodeId] -> Gen (Map CoreNodeId NodeRestart)
forall a b. (a -> b) -> a -> b
$
           [CoreNodeId] -> Gen [CoreNodeId]
forall a. [a] -> Gen [a]
sublistOf ([CoreNodeId] -> Gen [CoreNodeId])
-> [CoreNodeId] -> Gen [CoreNodeId]
forall a b. (a -> b) -> a -> b
$ Map CoreNodeId SlotNo -> [CoreNodeId]
forall k a. Map k a -> [k]
Map.keys (Map CoreNodeId SlotNo -> [CoreNodeId])
-> Map CoreNodeId SlotNo -> [CoreNodeId]
forall a b. (a -> b) -> a -> b
$ Map CoreNodeId SlotNo
candidates
  where
    isLeading :: CoreNodeId -> SlotNo -> Bool
isLeading (CoreNodeId Word64
i) SlotNo
s = Word64
i Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= SlotNo -> Word64
unSlotNo SlotNo
s Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
`mod` Word64
n
      where
        n :: Word64
n = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ Map CoreNodeId SlotNo -> Int
forall k a. Map k a -> Int
Map.size Map CoreNodeId SlotNo
m

    canRestartIn :: SlotNo -> CoreNodeId -> SlotNo -> Bool
canRestartIn SlotNo
s CoreNodeId
nid SlotNo
joinSlot =
        -- must be present
        SlotNo
joinSlot SlotNo -> SlotNo -> Bool
forall a. Ord a => a -> a -> Bool
<= SlotNo
s Bool -> Bool -> Bool
&&
        -- must not be leading (TODO relax this somehow?)
        Bool -> Bool
not (CoreNodeId -> SlotNo -> Bool
isLeading CoreNodeId
nid SlotNo
s)

shrinkNodeRestarts :: NodeRestarts -> [NodeRestarts]
shrinkNodeRestarts :: NodeRestarts -> [NodeRestarts]
shrinkNodeRestarts (NodeRestarts Map SlotNo (Map CoreNodeId NodeRestart)
m)
  | Map SlotNo (Map CoreNodeId NodeRestart) -> Bool
forall k a. Map k a -> Bool
Map.null Map SlotNo (Map CoreNodeId NodeRestart)
m = []  -- TODO better shrink
  | Bool
otherwise  = [NodeRestarts
noRestarts]