{-# LANGUAGE NamedFieldPuns #-}

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

import           Cardano.Ledger.BaseTypes (unNonZero)
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
> NonZero Word64 -> Word64
forall a. NonZero a -> a
unNonZero (SecurityParam -> NonZero 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'