{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}

#if __GLASGOW_HASKELL__ >= 908
{-# OPTIONS_GHC -Wno-x-partial #-}
#endif

module Test.ThreadNet.Util.NodeJoinPlan
  ( -- * Node Join Plan
    NodeJoinPlan (..)
  , coreNodeIdJoinSlot
  , genNodeJoinPlan
  , nodeIdJoinSlot
  , shrinkNodeJoinPlan
  , trivialNodeJoinPlan
  ) where

import qualified Data.List as List
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import GHC.Stack (HasCallStack)
import Ouroboros.Consensus.Block
import Ouroboros.Consensus.Node.ProtocolInfo
import Ouroboros.Consensus.NodeId
import Ouroboros.Consensus.Util.Condense
import Ouroboros.Consensus.Util.Orphans ()
import Test.QuickCheck
import Test.Util.Slots (NumSlots (..))

{-------------------------------------------------------------------------------
  Node Join Plans
-------------------------------------------------------------------------------}

-- | In which slot each node joins the network
newtype NodeJoinPlan = NodeJoinPlan (Map CoreNodeId SlotNo)
  deriving Int -> NodeJoinPlan -> ShowS
[NodeJoinPlan] -> ShowS
NodeJoinPlan -> String
(Int -> NodeJoinPlan -> ShowS)
-> (NodeJoinPlan -> String)
-> ([NodeJoinPlan] -> ShowS)
-> Show NodeJoinPlan
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NodeJoinPlan -> ShowS
showsPrec :: Int -> NodeJoinPlan -> ShowS
$cshow :: NodeJoinPlan -> String
show :: NodeJoinPlan -> String
$cshowList :: [NodeJoinPlan] -> ShowS
showList :: [NodeJoinPlan] -> ShowS
Show

instance Condense NodeJoinPlan where
  condense :: NodeJoinPlan -> String
condense (NodeJoinPlan Map CoreNodeId SlotNo
m) =
    [(NodeId, SlotNo)] -> String
forall a. Condense a => a -> String
condense
      [(CoreNodeId -> NodeId
fromCoreNodeId CoreNodeId
nid, SlotNo
slot) | (CoreNodeId
nid, SlotNo
slot) <- Map CoreNodeId SlotNo -> [(CoreNodeId, SlotNo)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map CoreNodeId SlotNo
m]

-- | All nodes join immediately
trivialNodeJoinPlan :: NumCoreNodes -> NodeJoinPlan
trivialNodeJoinPlan :: NumCoreNodes -> NodeJoinPlan
trivialNodeJoinPlan NumCoreNodes
numCoreNodes =
  Map CoreNodeId SlotNo -> NodeJoinPlan
NodeJoinPlan (Map CoreNodeId SlotNo -> NodeJoinPlan)
-> Map CoreNodeId SlotNo -> NodeJoinPlan
forall a b. (a -> b) -> a -> b
$
    [(CoreNodeId, SlotNo)] -> Map CoreNodeId SlotNo
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(CoreNodeId, SlotNo)] -> Map CoreNodeId SlotNo)
-> [(CoreNodeId, SlotNo)] -> Map CoreNodeId SlotNo
forall a b. (a -> b) -> a -> b
$
      [(CoreNodeId
nid, Word64 -> SlotNo
SlotNo Word64
0) | CoreNodeId
nid <- NumCoreNodes -> [CoreNodeId]
enumCoreNodes NumCoreNodes
numCoreNodes]

-- | Generate a 'NodeJoinPlan' consistent with the given properties
--
-- INVARIANT: Nodes with higher Ids will not join before nodes with lower Ids.
-- This eliminates some uninteresting symmetry and makes the counter-examples
-- easier for humans to interpret.
genNodeJoinPlan ::
  -- | PRECONDITION: non-negative
  NumCoreNodes ->
  -- | PRECONDITION: positive
  NumSlots ->
  Gen NodeJoinPlan
genNodeJoinPlan :: NumCoreNodes -> NumSlots -> Gen NodeJoinPlan
genNodeJoinPlan numCoreNodes :: NumCoreNodes
numCoreNodes@(NumCoreNodes Word64
n) numSlots :: NumSlots
numSlots@(NumSlots Word64
t)
  | Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
0 Bool -> Bool -> Bool
|| Word64
t Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
1 =
      String -> Gen NodeJoinPlan
forall a. HasCallStack => String -> a
error (String -> Gen NodeJoinPlan) -> String -> Gen NodeJoinPlan
forall a b. (a -> b) -> a -> b
$
        String
"Cannot generate TestConfig: "
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ (NumCoreNodes, NumSlots) -> String
forall a. Show a => a -> String
show (NumCoreNodes
numCoreNodes, NumSlots
numSlots)
  | Bool
otherwise = do
      let genJoinSlot :: Gen SlotNo
genJoinSlot = do
            let lastSlot :: Word64
lastSlot = Word64
t Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1
            Word64 -> SlotNo
SlotNo (Word64 -> SlotNo) -> Gen Word64 -> Gen SlotNo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word64, Word64) -> Gen Word64
forall a. Random a => (a, a) -> Gen a
choose (Word64
0, Word64
lastSlot)

      let nids :: [CoreNodeId]
nids = NumCoreNodes -> [CoreNodeId]
enumCoreNodes NumCoreNodes
numCoreNodes :: [CoreNodeId]
      schedules <- Int -> Gen SlotNo -> Gen [SlotNo]
forall a. Int -> Gen a -> Gen [a]
vectorOf (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n) Gen SlotNo
genJoinSlot
      -- without loss of generality, the nodes start initializing in order of
      -- their Ids; this merely makes it easer to interpret the counterexamples
      pure $ NodeJoinPlan $ Map.fromList $ zip nids $ List.sort schedules

-- | Shrink a node join plan
--
-- INVARIANT no inter-join delay increases
--
-- Specifically, we shrink by setting some of the delays to 0.
shrinkNodeJoinPlan :: NodeJoinPlan -> [NodeJoinPlan]
shrinkNodeJoinPlan :: NodeJoinPlan -> [NodeJoinPlan]
shrinkNodeJoinPlan (NodeJoinPlan Map CoreNodeId SlotNo
m0) =
  [NodeJoinPlan] -> [NodeJoinPlan]
forall a. HasCallStack => [a] -> [a]
init ([NodeJoinPlan] -> [NodeJoinPlan])
-> [NodeJoinPlan] -> [NodeJoinPlan]
forall a b. (a -> b) -> a -> b
$ -- the last one is the same as the input
    (((CoreNodeId, SlotNo), Map CoreNodeId SlotNo) -> NodeJoinPlan)
-> [((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)]
-> [NodeJoinPlan]
forall a b. (a -> b) -> [a] -> [b]
map (Map CoreNodeId SlotNo -> NodeJoinPlan
NodeJoinPlan (Map CoreNodeId SlotNo -> NodeJoinPlan)
-> (((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)
    -> Map CoreNodeId SlotNo)
-> ((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)
-> NodeJoinPlan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)
-> Map CoreNodeId SlotNo
forall a b. (a, b) -> b
snd) ([((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)] -> [NodeJoinPlan])
-> [((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)]
-> [NodeJoinPlan]
forall a b. (a -> b) -> a -> b
$
      [SlotNo] -> [((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)]
go [SlotNo]
diffs0
 where
  slots :: [SlotNo]
slots = ((CoreNodeId, SlotNo) -> SlotNo)
-> [(CoreNodeId, SlotNo)] -> [SlotNo]
forall a b. (a -> b) -> [a] -> [b]
map (CoreNodeId, SlotNo) -> SlotNo
forall a b. (a, b) -> b
snd (Map CoreNodeId SlotNo -> [(CoreNodeId, SlotNo)]
forall k a. Map k a -> [(k, a)]
Map.toDescList Map CoreNodeId SlotNo
m0) [SlotNo] -> [SlotNo] -> [SlotNo]
forall a. [a] -> [a] -> [a]
++ [SlotNo
0]
  diffs0 :: [SlotNo]
diffs0 = (SlotNo -> SlotNo -> SlotNo) -> [SlotNo] -> [SlotNo] -> [SlotNo]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\SlotNo
j2 SlotNo
j1 -> SlotNo
j2 SlotNo -> SlotNo -> SlotNo
forall a. Num a => a -> a -> a
- SlotNo
j1) [SlotNo]
slots ([SlotNo] -> [SlotNo]
forall a. HasCallStack => [a] -> [a]
tail [SlotNo]
slots)

  go :: [SlotNo] -> [((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)]
go = \case
    [] -> [((Word64 -> CoreNodeId
CoreNodeId Word64
0, SlotNo
0), Map CoreNodeId SlotNo
forall k a. Map k a
Map.empty)]
    SlotNo
d : [SlotNo]
ds -> do
      ((CoreNodeId i, mx), m) <- [SlotNo] -> [((CoreNodeId, SlotNo), Map CoreNodeId SlotNo)]
go [SlotNo]
ds
      let f SlotNo
s = ((Word64 -> CoreNodeId
CoreNodeId (Word64 -> Word64
forall a. Enum a => a -> a
succ Word64
i), SlotNo
s), CoreNodeId
-> SlotNo -> Map CoreNodeId SlotNo -> Map CoreNodeId SlotNo
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Word64 -> CoreNodeId
CoreNodeId Word64
i) SlotNo
s Map CoreNodeId SlotNo
m)
      [f mx] ++ [f (mx + d) | d > 0]

-- | Partial; @error@ for a node not in the plan
coreNodeIdJoinSlot ::
  HasCallStack =>
  NodeJoinPlan -> CoreNodeId -> SlotNo
coreNodeIdJoinSlot :: HasCallStack => NodeJoinPlan -> CoreNodeId -> SlotNo
coreNodeIdJoinSlot (NodeJoinPlan Map CoreNodeId SlotNo
m) CoreNodeId
nid =
  SlotNo -> CoreNodeId -> Map CoreNodeId SlotNo -> SlotNo
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault
    (String -> SlotNo
forall a. HasCallStack => String -> a
error (String -> SlotNo) -> String -> SlotNo
forall a b. (a -> b) -> a -> b
$ String
"not found: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> (CoreNodeId, [(CoreNodeId, SlotNo)]) -> String
forall a. Condense a => a -> String
condense (CoreNodeId
nid, Map CoreNodeId SlotNo -> [(CoreNodeId, SlotNo)]
forall k a. Map k a -> [(k, a)]
Map.toList Map CoreNodeId SlotNo
m))
    CoreNodeId
nid
    Map CoreNodeId SlotNo
m

-- | Partial; @error@ for a node not in the plan
nodeIdJoinSlot ::
  HasCallStack =>
  NodeJoinPlan -> NodeId -> SlotNo
nodeIdJoinSlot :: HasCallStack => NodeJoinPlan -> NodeId -> SlotNo
nodeIdJoinSlot nodeJoinPlan :: NodeJoinPlan
nodeJoinPlan@(NodeJoinPlan Map CoreNodeId SlotNo
m) NodeId
ni = case NodeId
ni of
  CoreId CoreNodeId
cni -> HasCallStack => NodeJoinPlan -> CoreNodeId -> SlotNo
NodeJoinPlan -> CoreNodeId -> SlotNo
coreNodeIdJoinSlot NodeJoinPlan
nodeJoinPlan CoreNodeId
cni
  NodeId
_ -> String -> SlotNo
forall a. HasCallStack => String -> a
error (String -> SlotNo) -> String -> SlotNo
forall a b. (a -> b) -> a -> b
$ String
"not found: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> (NodeId, [(CoreNodeId, SlotNo)]) -> String
forall a. Condense a => a -> String
condense (NodeId
ni, Map CoreNodeId SlotNo -> [(CoreNodeId, SlotNo)]
forall k a. Map k a -> [(k, a)]
Map.toList Map CoreNodeId SlotNo
m)