{-# 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 ::
     NumCoreNodes
     -- ^ PRECONDITION: non-negative
  -> NumSlots
     -- ^ PRECONDITION: positive
  -> 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)