{-# LANGUAGE NamedFieldPuns #-}

module Test.ThreadNet.Util.Expectations (
    NumBlocks (..)
  , determineForkLength
  ) where

import           Data.Foldable as Foldable (foldl')
import qualified Data.Map.Strict as Map
import           Data.Word (Word64)
import           Ouroboros.Consensus.Block
import           Ouroboros.Consensus.Config.SecurityParam
import           Ouroboros.Consensus.NodeId (CoreNodeId (..))
import           Ouroboros.Consensus.Protocol.LeaderSchedule
import           Test.ThreadNet.Util.NodeJoinPlan

newtype NumBlocks = NumBlocks Word64
  deriving (NumBlocks -> NumBlocks -> Bool
(NumBlocks -> NumBlocks -> Bool)
-> (NumBlocks -> NumBlocks -> Bool) -> Eq NumBlocks
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NumBlocks -> NumBlocks -> Bool
== :: NumBlocks -> NumBlocks -> Bool
$c/= :: NumBlocks -> NumBlocks -> Bool
/= :: NumBlocks -> NumBlocks -> Bool
Eq, Int -> NumBlocks -> ShowS
[NumBlocks] -> ShowS
NumBlocks -> String
(Int -> NumBlocks -> ShowS)
-> (NumBlocks -> String)
-> ([NumBlocks] -> ShowS)
-> Show NumBlocks
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NumBlocks -> ShowS
showsPrec :: Int -> NumBlocks -> ShowS
$cshow :: NumBlocks -> String
show :: NumBlocks -> String
$cshowList :: [NumBlocks] -> ShowS
showList :: [NumBlocks] -> ShowS
Show)

-- | Internal accumulator of 'determineForkLength'
data Acc = Acc
  { Acc -> Word64
maxChainLength :: !Word64
    -- ^ Upper bound on length of the longest chain in the network
  , Acc -> Word64
maxForkLength  :: !Word64
    -- ^ Upper bound on length of the longest fork in the network, excluding
    -- the common prefix
    --
    -- Note that @0@ corresponds to /consensus/.
  }

-- | Compute a bound for the length of the final forks
--
-- At the end of the test, the nodes' final chains will share a common prefix
-- (\"intersection\"), though it may be empty. Each node's current chain is a
-- fork off that prefix. No such fork will be longer than the result of this
-- function.
--
-- ASSUMPTION: The network connectivity is such that -- barring other
-- non-network considerations -- any block forged at the start of a slot will
-- be fetched and possibly selected by every other node before the end of the
-- slot.
--
-- == How 'LeaderSchedule' affects this function
--
-- A round-robin 'LeaderSchedule' will always reach consensus, so the fork
-- length will be @0@. For other 'LeaderSchedule's -- whether known /a priori/
-- (see "Test.ThreadNet.LeaderSchedule") or /a posteriori/ (see
-- "Test.ThreadNet.Praos") -- we often will not expect consensus. The key
-- difference is that a round-robin schedule always has exactly one leader per
-- slot. Arbitrary schedules instead may have 0 or multiple leaders.
--
-- A sequence of slots in which no slot has exactly one leader can drive a
-- network away from consensus. This function bounds how far from consensus the
-- network could be after the last simulated slot. It determines a conservative
-- upper bound on the length of the contended suffix of the nodes' final
-- chains. The bound may be less than, equal, or greater than @k@; regardless
-- of which, it should be used to test the nodes' final chains.
--
-- If such a sequence is long enough, it might even prevent the network from
-- every reaching consensus again: the nodes at the end of the sequence may
-- have chains that disagree by more than @k@ blocks, thereby exceeding the
-- security parameter @k@. In particular, if the bound determined by a
-- 'LeaderSchedule' is greater than @k@, then no extension of that
-- 'LeaderSchedule' can determine a lesser bound.
--
-- Multiple leaders create divergent chains as follows. Assume that every
-- leader of @s + 1@ begins the slot with a chain of length @n@, and that no
-- chain in the network has a greater length. Each leader forges a unique
-- block. A race condition between ChainSync/BlockFetch and forging makes it
-- possible, though relatively unlikely, that a leader would have received
-- another leader's new block before it forges its own. In that case, the new
-- leader will forged a block but then not select it, since it already selected
-- the other leader's new block. This function assumes the \"worst-case\" in
-- which both leaders make separate chains, thereby breaking consensus. Hence
-- it computes an upper bound. Each leader thus selects a new @n+1@-length
-- chain, and each non-leader will adopt one of them because they previously
-- had a chain of length @<= n@. Thus the network is likely not in consensus at
-- the end of a multi-leader slot.
--
-- The network will reestablish consensus at the end of a single-leader slot,
-- because the leader creates the only chain of length @n+1@ and all other
-- nodes thus select it.
--
-- In a slot with no leaders, all nodes will simply retain their current
-- chains.
--
-- == How 'NodeJoinPlan' affects this function
--
-- Because the race condition between forging and ChainSync/BlockFetch is
-- consistently won by forging, a node that leads in the same slot it joins
-- always attempts to make a chain of length 1 (ASSUMPTION: new nodes start
-- with an empty chain). Except for the edge cases when the longest chains in
-- the network are of length 0 or 1, this means that leaders who are joining
-- can be disregarded.
--
determineForkLength ::
     SecurityParam
  -> NodeJoinPlan
  -> LeaderSchedule
  -> NumBlocks
determineForkLength :: SecurityParam -> NodeJoinPlan -> LeaderSchedule -> NumBlocks
determineForkLength SecurityParam
k (NodeJoinPlan Map CoreNodeId SlotNo
joinPlan) (LeaderSchedule Map SlotNo [CoreNodeId]
sched) =
    Acc -> NumBlocks
prj (Acc -> NumBlocks) -> Acc -> NumBlocks
forall a b. (a -> b) -> a -> b
$ (Acc -> (SlotNo, [CoreNodeId]) -> Acc)
-> Acc -> [(SlotNo, [CoreNodeId])] -> Acc
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' Acc -> (SlotNo, [CoreNodeId]) -> Acc
step Acc
initial (Map SlotNo [CoreNodeId] -> [(SlotNo, [CoreNodeId])]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map SlotNo [CoreNodeId]
sched)
  where
    prj :: Acc -> NumBlocks
prj Acc{Word64
maxForkLength :: Acc -> Word64
maxForkLength :: Word64
maxForkLength} = Word64 -> NumBlocks
NumBlocks Word64
maxForkLength

    -- assume the network begins in consensus (eg all nodes start at Genesis)
    initial :: Acc
initial = Acc
      { maxChainLength :: Word64
maxChainLength = Word64
0
      , maxForkLength :: Word64
maxForkLength = Word64
0
      }

    -- this logic focuses on the new chains made in this slot that are longer
    -- than the longest chains from the previous slot
    step :: Acc -> (SlotNo, [CoreNodeId]) -> Acc
step Acc{Word64
maxChainLength :: Acc -> Word64
maxChainLength :: Word64
maxChainLength, Word64
maxForkLength :: Acc -> Word64
maxForkLength :: Word64
maxForkLength} (SlotNo
slot, [CoreNodeId]
leaders) =
        Acc
          { maxChainLength :: Word64
maxChainLength = Word64 -> Word64
grow    Word64
maxChainLength
          , maxForkLength :: Word64
maxForkLength  = Word64 -> Word64
update  Word64
maxForkLength
          }
      where
        grow :: Word64 -> Word64
grow = if Int
0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
pullingAhead then Word64 -> Word64
forall a. a -> a
id else (Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1)

        update :: Word64 -> Word64
update
            -- too late to reach consensus, so further diverge
          | Word64
maxForkLength Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> SecurityParam -> Word64
maxRollbacks SecurityParam
k = Word64 -> Word64
grow

            -- assume (common) worst-case: each leader creates a unique longer
            -- chain
          | Int
pullingAhead Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>  Int
1              = (Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1)

            -- the sole leader creates the sole longer chain, bringing the
            -- network into consensus
          | Int
pullingAhead Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1              = Word64 -> Word64 -> Word64
forall a b. a -> b -> a
const Word64
0

            -- there will be multiple longest chains that disagree on at least
            -- the latest block
            --
            -- Note that pullingAhead == 0 by the preceding guards
          | Int
pullingEven  Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>  Int
0              = Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
max Word64
1

            -- no nodes are extending their chain, so the longest chains are
            -- the same as in the previous slot
          | Bool
otherwise                      = Word64 -> Word64
forall a. a -> a
id

        -- how many leaders are forging a block onto a @maxForkLength@-chain
        pullingAhead :: Int
pullingAhead = Int
nlOld Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Bool -> Int
nlNew (Word64
maxChainLength Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0)
        -- how many leaders are forging a block onto a @maxForkLength - 1@-chain
        pullingEven :: Int
pullingEven  = Bool -> Int
nlNew (Word64
maxChainLength Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
1)

        -- how many leaders joined before this slot
        nlOld :: Int
nlOld = [CoreNodeId] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([CoreNodeId] -> Int) -> [CoreNodeId] -> Int
forall a b. (a -> b) -> a -> b
$ (CoreNodeId -> Bool) -> [CoreNodeId] -> [CoreNodeId]
forall a. (a -> Bool) -> [a] -> [a]
filter ((SlotNo -> SlotNo -> Bool
forall a. Ord a => a -> a -> Bool
< SlotNo
slot) (SlotNo -> Bool) -> (CoreNodeId -> SlotNo) -> CoreNodeId -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreNodeId -> SlotNo
joinSlot) [CoreNodeId]
leaders
        nlNew :: Bool -> Int
nlNew Bool
b
          | Bool -> Bool
not Bool
b     = Int
0
            -- how many leaders are joining during this slot; these might not
            -- be relevant
            --
            -- A node leading the same slot it joins always forges a 1-block
            -- chain; there's actually a race-condition between block forging
            -- and BlockFetch/ChainSync, but forging always wins in the current
            -- test framework implementation.
          | Bool
otherwise = [CoreNodeId] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([CoreNodeId] -> Int) -> [CoreNodeId] -> Int
forall a b. (a -> b) -> a -> b
$ (CoreNodeId -> Bool) -> [CoreNodeId] -> [CoreNodeId]
forall a. (a -> Bool) -> [a] -> [a]
filter ((SlotNo -> SlotNo -> Bool
forall a. Eq a => a -> a -> Bool
== SlotNo
slot) (SlotNo -> Bool) -> (CoreNodeId -> SlotNo) -> CoreNodeId -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreNodeId -> SlotNo
joinSlot) [CoreNodeId]
leaders

        -- slot in which the node joins the network
        joinSlot :: CoreNodeId -> SlotNo
        joinSlot :: CoreNodeId -> SlotNo
joinSlot CoreNodeId
nid = case CoreNodeId -> Map CoreNodeId SlotNo -> Maybe SlotNo
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup CoreNodeId
nid Map CoreNodeId SlotNo
joinPlan of
            Maybe SlotNo
Nothing    -> String -> SlotNo
forall a. HasCallStack => String -> a
error String
"determineForkLength: incomplete node join plan"
            Just SlotNo
slot' -> SlotNo
slot'