{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}

module Test.ThreadNet.Rekeying (
    Rekeying (..)
  , fromRekeyingToRekeyM
  ) where

import           Data.Functor ((<&>))
import           Ouroboros.Consensus.Block
import           Ouroboros.Consensus.Node.ProtocolInfo
import           Ouroboros.Consensus.NodeId
import           Ouroboros.Consensus.Util.IOLike
import           Test.ThreadNet.Network
import           Test.Util.Stream

-- | Functionality used by test node in order to update its operational key
--
-- This is the conceptual interface demanded from the test-specific logic. It
-- is used to define 'tnaRekeyM', which the test infrastructure invokes per the
-- 'NodeRestarts' schedule.
data Rekeying m blk = forall opKey. Rekeying
  { forall (m :: * -> *) blk.
Rekeying m blk -> CoreNodeId -> SlotNo -> Maybe SlotNo
rekeyOracle
      :: CoreNodeId -> SlotNo -> Maybe SlotNo
    -- ^ The first /nominal/ slot after the given slot, assuming the given core
    -- node cannot lead.
    --
    -- IE the first slot that will result in a block successfully being forged
    -- and diffused (eg no @PBftExceededSignThreshold@).
  , ()
rekeyUpd ::
         CoreNodeId
      -> ProtocolInfo blk
      -> m [BlockForging m blk]
      -> EpochNo
      -> opKey
      -> m (Maybe (TestNodeInitialization m blk))
     -- ^ new config and any corresponding delegation certificate transactions
     --
     -- The given epoch contains the first nominal slot whose block will
     -- include the redelegation certificate transaction.
     --
     -- The 'TestNodeInitialization' includes the new 'ProtocolInfo' used when
     -- the node completes restarting.
  , ()
rekeyFreshSKs :: Stream opKey
     -- ^ a stream that only repeats itself after an *effectively* *infinite*
     -- number of iterations and also never includes an operational key from
     -- the genesis configuration
  }

fromRekeyingToRekeyM :: IOLike m => Rekeying m blk -> m (RekeyM m blk)
fromRekeyingToRekeyM :: forall (m :: * -> *) blk.
IOLike m =>
Rekeying m blk -> m (RekeyM m blk)
fromRekeyingToRekeyM Rekeying{Stream opKey
rekeyFreshSKs :: ()
rekeyFreshSKs :: Stream opKey
rekeyFreshSKs, CoreNodeId -> SlotNo -> Maybe SlotNo
rekeyOracle :: forall (m :: * -> *) blk.
Rekeying m blk -> CoreNodeId -> SlotNo -> Maybe SlotNo
rekeyOracle :: CoreNodeId -> SlotNo -> Maybe SlotNo
rekeyOracle, CoreNodeId
-> ProtocolInfo blk
-> m [BlockForging m blk]
-> EpochNo
-> opKey
-> m (Maybe (TestNodeInitialization m blk))
rekeyUpd :: ()
rekeyUpd :: CoreNodeId
-> ProtocolInfo blk
-> m [BlockForging m blk]
-> EpochNo
-> opKey
-> m (Maybe (TestNodeInitialization m blk))
rekeyUpd} = do
    StrictTVar m (Stream opKey)
rekeyVar <- Stream opKey -> m (StrictTVar m (Stream opKey))
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
uncheckedNewTVarM Stream opKey
rekeyFreshSKs
    RekeyM m blk -> m (RekeyM m blk)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RekeyM m blk -> m (RekeyM m blk))
-> RekeyM m blk -> m (RekeyM m blk)
forall a b. (a -> b) -> a -> b
$ \CoreNodeId
cid ProtocolInfo blk
pInfo m [BlockForging m blk]
blockForging SlotNo
s SlotNo -> m EpochNo
mkEno -> case CoreNodeId -> SlotNo -> Maybe SlotNo
rekeyOracle CoreNodeId
cid SlotNo
s of
      Maybe SlotNo
Nothing -> TestNodeInitialization m blk -> m (TestNodeInitialization m blk)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TestNodeInitialization m blk -> m (TestNodeInitialization m blk))
-> TestNodeInitialization m blk -> m (TestNodeInitialization m blk)
forall a b. (a -> b) -> a -> b
$ ProtocolInfo blk
-> m [BlockForging m blk] -> TestNodeInitialization m blk
forall blk (m :: * -> *).
ProtocolInfo blk
-> m [BlockForging m blk] -> TestNodeInitialization m blk
plainTestNodeInitialization ProtocolInfo blk
pInfo m [BlockForging m blk]
blockForging
      Just SlotNo
s' -> do
        opKey
x <- STM m opKey -> m opKey
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m opKey -> m opKey) -> STM m opKey -> m opKey
forall a b. (a -> b) -> a -> b
$ do
          opKey
x :< Stream opKey
xs <- StrictTVar m (Stream opKey) -> STM m (Stream opKey)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m (Stream opKey)
rekeyVar
          opKey
x opKey -> STM m () -> STM m opKey
forall a b. a -> STM m b -> STM m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ StrictTVar m (Stream opKey) -> Stream opKey -> STM m ()
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
StrictTVar m a -> a -> STM m ()
writeTVar StrictTVar m (Stream opKey)
rekeyVar Stream opKey
xs
        EpochNo
eno <- SlotNo -> m EpochNo
mkEno SlotNo
s'
        CoreNodeId
-> ProtocolInfo blk
-> m [BlockForging m blk]
-> EpochNo
-> opKey
-> m (Maybe (TestNodeInitialization m blk))
rekeyUpd CoreNodeId
cid ProtocolInfo blk
pInfo m [BlockForging m blk]
blockForging EpochNo
eno opKey
x m (Maybe (TestNodeInitialization m blk))
-> (Maybe (TestNodeInitialization m blk)
    -> TestNodeInitialization m blk)
-> m (TestNodeInitialization m blk)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
          Maybe (TestNodeInitialization m blk)
Nothing  -> ProtocolInfo blk
-> m [BlockForging m blk] -> TestNodeInitialization m blk
forall blk (m :: * -> *).
ProtocolInfo blk
-> m [BlockForging m blk] -> TestNodeInitialization m blk
plainTestNodeInitialization ProtocolInfo blk
pInfo m [BlockForging m blk]
blockForging
          Just TestNodeInitialization m blk
tni -> TestNodeInitialization m blk
tni