content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE NoImplicitPrelude #-}
  
  module Cardano.DbSync.Era.Shelley.Generic.StakeDist (
    StakeSliceRes (..),
    StakeSlice (..),
    getSecurityParameter,
    getStakeSlice,
    getPoolDistr,
  ) where
  
  import Cardano.DbSync.Types
  import Cardano.Ledger.BaseTypes.NonZero (NonZero (..))
  import Cardano.Ledger.Coin (Coin (..))
  import qualified Cardano.Ledger.Compactible as Ledger
  import Cardano.Ledger.Credential (Credential)
  import Cardano.Ledger.Keys (KeyHash (..), KeyRole (..))
  import qualified Cardano.Ledger.Shelley.LedgerState as Shelley
  import qualified Cardano.Ledger.State as Ledger
  import Cardano.Ledger.Val ((<+>))
  import Cardano.Prelude
  import qualified Data.Map.Strict as Map
  import Data.VMap (VB, VMap (..), VP)
  import qualified Data.VMap as VMap
  import qualified Data.Vector.Generic as VG
  import Lens.Micro
  import Ouroboros.Consensus.Block
  import Ouroboros.Consensus.Cardano.Block (LedgerState (..))
  import Ouroboros.Consensus.Config
  import Ouroboros.Consensus.HardFork.Combinator
  import Ouroboros.Consensus.Ledger.Extended (ExtLedgerState (..))
  import Ouroboros.Consensus.Node.ProtocolInfo
  import Ouroboros.Consensus.Protocol.Abstract
  import Ouroboros.Consensus.Shelley.Ledger
  import qualified Ouroboros.Consensus.Shelley.Ledger.Ledger as Consensus
  import Prelude (id)
  
  data StakeSliceRes
    = Slice !StakeSlice !Bool -- True if this is the final slice for this epoch. Can be used for logging.
    | NoSlices
  
  data StakeSlice = StakeSlice
    { sliceEpochNo :: !EpochNo
    , sliceDistr :: !(Map StakeCred (Coin, PoolKeyHash))
    }
    deriving (Eq)
  
  emptySlice :: EpochNo -> StakeSlice
  emptySlice epoch = StakeSlice epoch Map.empty
  
  getSecurityParameter ::
    ConsensusProtocol (BlockProtocol blk) =>
    ProtocolInfo blk ->
    NonZero Word64
  getSecurityParameter = maxRollbacks . configSecurityParam . pInfoConfig
  
  -- 'sliceIndex' can match the epochBlockNo for every block.
  --
  -- 'minSliceSize' has to be constant or it could cause missing data.
  -- If this value is too small it will be adjusted to a 'defaultEpochSliceSize'
  -- which is big enough to cover all delegations.
  -- On mainnet, for a value minSliceSize = 2000, it will be used as the actual size of slices
  -- until the size of delegations grows up to 8.6M, in which case, the size of slices
  -- will be adjusted.
  getStakeSlice ::
    ConsensusProtocol (BlockProtocol blk) =>
    ProtocolInfo blk ->
    Word64 ->
    ExtLedgerState CardanoBlock ->
    Bool ->
    StakeSliceRes
  getStakeSlice pInfo !epochBlockNo els isMigration =
    case ledgerState els of
      LedgerStateByron _ -> NoSlices
      LedgerStateShelley sls -> genericStakeSlice pInfo epochBlockNo sls isMigration
      LedgerStateAllegra als -> genericStakeSlice pInfo epochBlockNo als isMigration
      LedgerStateMary mls -> genericStakeSlice pInfo epochBlockNo mls isMigration
      LedgerStateAlonzo als -> genericStakeSlice pInfo epochBlockNo als isMigration
      LedgerStateBabbage bls -> genericStakeSlice pInfo epochBlockNo bls isMigration
      LedgerStateConway cls -> genericStakeSlice pInfo epochBlockNo cls isMigration
  
  genericStakeSlice ::
    forall era blk p.
    ConsensusProtocol (BlockProtocol blk) =>
    ProtocolInfo blk ->
    Word64 ->
    LedgerState (ShelleyBlock p era) ->
    Bool ->
    StakeSliceRes
  genericStakeSlice pInfo epochBlockNo lstate isMigration
    | index > delegationsLen = NoSlices
    | index == delegationsLen = Slice (emptySlice epoch) True
    | index + size > delegationsLen = Slice (mkSlice (delegationsLen - index)) True
    | otherwise = Slice (mkSlice size) False
    where
      epoch :: EpochNo
      epoch = EpochNo $ 1 + unEpochNo (Shelley.nesEL (Consensus.shelleyLedgerState lstate))
  
      minSliceSize :: Word64
      minSliceSize = 2000
  
      -- On mainnet this is 2160
      k :: Word64
      k = unNonZero $ getSecurityParameter pInfo
  
      -- We use 'ssStakeMark' here. That means that when these values
      -- are added to the database, the epoch number where they become active is the current
      -- epoch plus one.
      stakeSnapshot :: Ledger.SnapShot
      stakeSnapshot =
        Ledger.ssStakeMark . Shelley.esSnapshots . Shelley.nesEs $
          Consensus.shelleyLedgerState lstate
  
      delegations :: VMap.KVVector VB VB (Credential 'Staking, KeyHash 'StakePool)
      delegations = VMap.unVMap $ Ledger.ssDelegations stakeSnapshot
  
      delegationsLen :: Word64
      delegationsLen = fromIntegral $ VG.length delegations
  
      stakes :: VMap VB VP (Credential 'Staking) (Ledger.CompactForm Coin)
      stakes = Ledger.unStake $ Ledger.ssStake stakeSnapshot
  
      lookupStake :: Credential 'Staking -> Maybe Coin
      lookupStake cred = Ledger.fromCompact <$> VMap.lookup cred stakes
  
      -- This is deterministic for the whole epoch and is the constant size of slices
      -- until the data are over. This means the last slice could be of smaller size and slices
      -- after that will be empty.
      epochSliceSize :: Word64
      epochSliceSize =
        max minSliceSize defaultEpochSliceSize
        where
          -- On mainnet this is 2160
          expectedBlocks :: Word64
          expectedBlocks = 10 * k
  
          -- This size of slices is enough to cover the whole list, even if only
          -- the 20% of the expected blocks appear in an epoch.
          defaultEpochSliceSize :: Word64
          defaultEpochSliceSize = 1 + div (delegationsLen * 5) expectedBlocks
  
      -- The starting index of the data in the delegation vector.
      index :: Word64
      index
        | isMigration = 0
        | epochBlockNo < k = delegationsLen + 1 -- so it creates the empty Slice.
        | otherwise = (epochBlockNo - k) * epochSliceSize
  
      size :: Word64
      size
        | isMigration, epochBlockNo + 1 < k = 0
        | isMigration = (epochBlockNo + 1 - k) * epochSliceSize
        | otherwise = epochSliceSize
  
      mkSlice :: Word64 -> StakeSlice
      mkSlice actualSize =
        StakeSlice
          { sliceEpochNo = epoch
          , sliceDistr = distribution
          }
        where
          delegationsSliced :: VMap VB VB (Credential 'Staking) (KeyHash 'StakePool)
          delegationsSliced = VMap $ VG.slice (fromIntegral index) (fromIntegral actualSize) delegations
  
          distribution :: Map StakeCred (Coin, PoolKeyHash)
          distribution =
            VMap.toMap $
              VMap.mapMaybe id $
                VMap.mapWithKey (\a p -> (,p) <$> lookupStake a) delegationsSliced
  
  getPoolDistr ::
    ExtLedgerState CardanoBlock ->
    Maybe (Map PoolKeyHash (Coin, Word64), Map PoolKeyHash Natural)
  getPoolDistr els =
    case ledgerState els of
      LedgerStateByron _ -> Nothing
      LedgerStateShelley sls -> Just $ genericPoolDistr sls
      LedgerStateAllegra als -> Just $ genericPoolDistr als
      LedgerStateMary mls -> Just $ genericPoolDistr mls
      LedgerStateAlonzo als -> Just $ genericPoolDistr als
      LedgerStateBabbage bls -> Just $ genericPoolDistr bls
      LedgerStateConway cls -> Just $ genericPoolDistr cls
  
  genericPoolDistr ::
    forall era p.
    LedgerState (ShelleyBlock p era) ->
    (Map PoolKeyHash (Coin, Word64), Map PoolKeyHash Natural)
  genericPoolDistr lstate =
    (stakePerPool, blocksPerPool)
    where
      nes :: Shelley.NewEpochState era
      nes = Consensus.shelleyLedgerState lstate
  
      stakeMark :: Ledger.SnapShot
      stakeMark = Ledger.ssStakeMark $ Shelley.esSnapshots $ Shelley.nesEs nes
  
      stakePerPool = countStakePerPool (Ledger.ssDelegations stakeMark) (Ledger.ssStake stakeMark)
      blocksPerPool = nes ^. Shelley.nesBprevL
  
  countStakePerPool ::
    VMap VB VB StakeCred PoolKeyHash ->
    Ledger.Stake ->
    Map PoolKeyHash (Coin, Word64)
  countStakePerPool delegs (Ledger.Stake stake) = VMap.foldlWithKey accum Map.empty stake
    where
      accum !acc cred compactCoin =
        case VMap.lookup cred delegs of
          Nothing -> acc
          Just kh -> Map.insertWith addDel kh (Ledger.fromCompact compactCoin, 1) acc
  
      addDel (c, n) (c', n') = (c <+> c', n + n')
  