{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Ouroboros.Consensus.Mock.Ledger.Stake
  ( -- * Stakeholders
    StakeHolder (..)

    -- * Address distribution
  , AddrDist

    -- * Stake distribution
  , StakeDist (..)
  , equalStakeDist
  , genesisStakeDist
  , relativeStakes
  , stakeWithDefault
  , totalStakes

    -- * Type family instances
  , Ticked (..)
  ) where

import Codec.Serialise (Serialise)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (mapMaybe)
import NoThunks.Class (NoThunks)
import Ouroboros.Consensus.Mock.Ledger.Address
import Ouroboros.Consensus.Mock.Ledger.UTxO
import Ouroboros.Consensus.NodeId (CoreNodeId (..), NodeId (..))
import Ouroboros.Consensus.Ticked

{-------------------------------------------------------------------------------
  Stakeholders
-------------------------------------------------------------------------------}

data StakeHolder
  = -- | Stake of a core node
    StakeCore CoreNodeId
  | -- | Stake for everybody else (we don't need to distinguish)
    StakeEverybodyElse
  deriving (Int -> StakeHolder -> ShowS
[StakeHolder] -> ShowS
StakeHolder -> String
(Int -> StakeHolder -> ShowS)
-> (StakeHolder -> String)
-> ([StakeHolder] -> ShowS)
-> Show StakeHolder
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StakeHolder -> ShowS
showsPrec :: Int -> StakeHolder -> ShowS
$cshow :: StakeHolder -> String
show :: StakeHolder -> String
$cshowList :: [StakeHolder] -> ShowS
showList :: [StakeHolder] -> ShowS
Show, StakeHolder -> StakeHolder -> Bool
(StakeHolder -> StakeHolder -> Bool)
-> (StakeHolder -> StakeHolder -> Bool) -> Eq StakeHolder
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StakeHolder -> StakeHolder -> Bool
== :: StakeHolder -> StakeHolder -> Bool
$c/= :: StakeHolder -> StakeHolder -> Bool
/= :: StakeHolder -> StakeHolder -> Bool
Eq, Eq StakeHolder
Eq StakeHolder =>
(StakeHolder -> StakeHolder -> Ordering)
-> (StakeHolder -> StakeHolder -> Bool)
-> (StakeHolder -> StakeHolder -> Bool)
-> (StakeHolder -> StakeHolder -> Bool)
-> (StakeHolder -> StakeHolder -> Bool)
-> (StakeHolder -> StakeHolder -> StakeHolder)
-> (StakeHolder -> StakeHolder -> StakeHolder)
-> Ord StakeHolder
StakeHolder -> StakeHolder -> Bool
StakeHolder -> StakeHolder -> Ordering
StakeHolder -> StakeHolder -> StakeHolder
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: StakeHolder -> StakeHolder -> Ordering
compare :: StakeHolder -> StakeHolder -> Ordering
$c< :: StakeHolder -> StakeHolder -> Bool
< :: StakeHolder -> StakeHolder -> Bool
$c<= :: StakeHolder -> StakeHolder -> Bool
<= :: StakeHolder -> StakeHolder -> Bool
$c> :: StakeHolder -> StakeHolder -> Bool
> :: StakeHolder -> StakeHolder -> Bool
$c>= :: StakeHolder -> StakeHolder -> Bool
>= :: StakeHolder -> StakeHolder -> Bool
$cmax :: StakeHolder -> StakeHolder -> StakeHolder
max :: StakeHolder -> StakeHolder -> StakeHolder
$cmin :: StakeHolder -> StakeHolder -> StakeHolder
min :: StakeHolder -> StakeHolder -> StakeHolder
Ord)

{-------------------------------------------------------------------------------
  Stake distribution
-------------------------------------------------------------------------------}

-- | In the mock setup, only core nodes have stake
--
-- INVARIANT: The rationals should sum to 1.
newtype StakeDist = StakeDist {StakeDist -> Map CoreNodeId (Ratio Integer)
stakeDistToMap :: Map CoreNodeId Rational}
  deriving (Int -> StakeDist -> ShowS
[StakeDist] -> ShowS
StakeDist -> String
(Int -> StakeDist -> ShowS)
-> (StakeDist -> String)
-> ([StakeDist] -> ShowS)
-> Show StakeDist
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StakeDist -> ShowS
showsPrec :: Int -> StakeDist -> ShowS
$cshow :: StakeDist -> String
show :: StakeDist -> String
$cshowList :: [StakeDist] -> ShowS
showList :: [StakeDist] -> ShowS
Show, StakeDist -> StakeDist -> Bool
(StakeDist -> StakeDist -> Bool)
-> (StakeDist -> StakeDist -> Bool) -> Eq StakeDist
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StakeDist -> StakeDist -> Bool
== :: StakeDist -> StakeDist -> Bool
$c/= :: StakeDist -> StakeDist -> Bool
/= :: StakeDist -> StakeDist -> Bool
Eq, [StakeDist] -> Encoding
StakeDist -> Encoding
(StakeDist -> Encoding)
-> (forall s. Decoder s StakeDist)
-> ([StakeDist] -> Encoding)
-> (forall s. Decoder s [StakeDist])
-> Serialise StakeDist
forall s. Decoder s [StakeDist]
forall s. Decoder s StakeDist
forall a.
(a -> Encoding)
-> (forall s. Decoder s a)
-> ([a] -> Encoding)
-> (forall s. Decoder s [a])
-> Serialise a
$cencode :: StakeDist -> Encoding
encode :: StakeDist -> Encoding
$cdecode :: forall s. Decoder s StakeDist
decode :: forall s. Decoder s StakeDist
$cencodeList :: [StakeDist] -> Encoding
encodeList :: [StakeDist] -> Encoding
$cdecodeList :: forall s. Decoder s [StakeDist]
decodeList :: forall s. Decoder s [StakeDist]
Serialise, Context -> StakeDist -> IO (Maybe ThunkInfo)
Proxy StakeDist -> String
(Context -> StakeDist -> IO (Maybe ThunkInfo))
-> (Context -> StakeDist -> IO (Maybe ThunkInfo))
-> (Proxy StakeDist -> String)
-> NoThunks StakeDist
forall a.
(Context -> a -> IO (Maybe ThunkInfo))
-> (Context -> a -> IO (Maybe ThunkInfo))
-> (Proxy a -> String)
-> NoThunks a
$cnoThunks :: Context -> StakeDist -> IO (Maybe ThunkInfo)
noThunks :: Context -> StakeDist -> IO (Maybe ThunkInfo)
$cwNoThunks :: Context -> StakeDist -> IO (Maybe ThunkInfo)
wNoThunks :: Context -> StakeDist -> IO (Maybe ThunkInfo)
$cshowTypeOf :: Proxy StakeDist -> String
showTypeOf :: Proxy StakeDist -> String
NoThunks)

stakeWithDefault :: Rational -> CoreNodeId -> StakeDist -> Rational
stakeWithDefault :: Ratio Integer -> CoreNodeId -> StakeDist -> Ratio Integer
stakeWithDefault Ratio Integer
d CoreNodeId
n = Ratio Integer
-> CoreNodeId -> Map CoreNodeId (Ratio Integer) -> Ratio Integer
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Ratio Integer
d CoreNodeId
n (Map CoreNodeId (Ratio Integer) -> Ratio Integer)
-> (StakeDist -> Map CoreNodeId (Ratio Integer))
-> StakeDist
-> Ratio Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StakeDist -> Map CoreNodeId (Ratio Integer)
stakeDistToMap

relativeStakes :: Map StakeHolder Amount -> StakeDist
relativeStakes :: Map StakeHolder Word -> StakeDist
relativeStakes Map StakeHolder Word
m =
  Map CoreNodeId (Ratio Integer) -> StakeDist
StakeDist (Map CoreNodeId (Ratio Integer) -> StakeDist)
-> Map CoreNodeId (Ratio Integer) -> StakeDist
forall a b. (a -> b) -> a -> b
$
    let totalStake :: Ratio Integer
totalStake = Word -> Ratio Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Ratio Integer) -> Word -> Ratio Integer
forall a b. (a -> b) -> a -> b
$ [Word] -> Word
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Word] -> Word) -> [Word] -> Word
forall a b. (a -> b) -> a -> b
$ Map StakeHolder Word -> [Word]
forall k a. Map k a -> [a]
Map.elems Map StakeHolder Word
m
     in [(CoreNodeId, Ratio Integer)] -> Map CoreNodeId (Ratio Integer)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
          [ (CoreNodeId
nid, Word -> Ratio Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
stake Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Fractional a => a -> a -> a
/ Ratio Integer
totalStake)
          | (StakeCore CoreNodeId
nid, Word
stake) <- Map StakeHolder Word -> [(StakeHolder, Word)]
forall k a. Map k a -> [(k, a)]
Map.toList Map StakeHolder Word
m
          ]

-- | Compute stakes of all nodes
--
-- The 'Nothing' value holds the total stake of all addresses that don't
-- get mapped to a NodeId.
totalStakes :: Map Addr NodeId -> Utxo -> Map StakeHolder Amount
totalStakes :: Map Addr NodeId -> Utxo -> Map StakeHolder Word
totalStakes Map Addr NodeId
addrDist = (Map StakeHolder Word -> TxOut -> Map StakeHolder Word)
-> Map StakeHolder Word -> Utxo -> Map StakeHolder Word
forall b a. (b -> a -> b) -> b -> Map TxIn a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map StakeHolder Word -> TxOut -> Map StakeHolder Word
f Map StakeHolder Word
forall k a. Map k a
Map.empty
 where
  f :: Map StakeHolder Amount -> TxOut -> Map StakeHolder Amount
  f :: Map StakeHolder Word -> TxOut -> Map StakeHolder Word
f Map StakeHolder Word
m (Addr
a, Word
stake) = case Addr -> Map Addr NodeId -> Maybe NodeId
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Addr
a Map Addr NodeId
addrDist of
    Just (CoreId CoreNodeId
nid) -> (Word -> Word -> Word)
-> StakeHolder
-> Word
-> Map StakeHolder Word
-> Map StakeHolder Word
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Word -> Word -> Word
forall a. Num a => a -> a -> a
(+) (CoreNodeId -> StakeHolder
StakeCore CoreNodeId
nid) Word
stake Map StakeHolder Word
m
    Maybe NodeId
_ -> (Word -> Word -> Word)
-> StakeHolder
-> Word
-> Map StakeHolder Word
-> Map StakeHolder Word
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Word -> Word -> Word
forall a. Num a => a -> a -> a
(+) StakeHolder
StakeEverybodyElse Word
stake Map StakeHolder Word
m

-- | Stake distribution where every address has equal state
equalStakeDist :: AddrDist -> StakeDist
equalStakeDist :: Map Addr NodeId -> StakeDist
equalStakeDist Map Addr NodeId
ad =
  Map CoreNodeId (Ratio Integer) -> StakeDist
StakeDist (Map CoreNodeId (Ratio Integer) -> StakeDist)
-> Map CoreNodeId (Ratio Integer) -> StakeDist
forall a b. (a -> b) -> a -> b
$
    [(CoreNodeId, Ratio Integer)] -> Map CoreNodeId (Ratio Integer)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(CoreNodeId, Ratio Integer)] -> Map CoreNodeId (Ratio Integer))
-> [(CoreNodeId, Ratio Integer)] -> Map CoreNodeId (Ratio Integer)
forall a b. (a -> b) -> a -> b
$
      ((Addr, NodeId) -> Maybe (CoreNodeId, Ratio Integer))
-> [(Addr, NodeId)] -> [(CoreNodeId, Ratio Integer)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (NodeId -> Maybe (CoreNodeId, Ratio Integer)
nodeStake (NodeId -> Maybe (CoreNodeId, Ratio Integer))
-> ((Addr, NodeId) -> NodeId)
-> (Addr, NodeId)
-> Maybe (CoreNodeId, Ratio Integer)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Addr, NodeId) -> NodeId
forall a b. (a, b) -> b
snd) ([(Addr, NodeId)] -> [(CoreNodeId, Ratio Integer)])
-> [(Addr, NodeId)] -> [(CoreNodeId, Ratio Integer)]
forall a b. (a -> b) -> a -> b
$
        Map Addr NodeId -> [(Addr, NodeId)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Addr NodeId
ad
 where
  nodeStake :: NodeId -> Maybe (CoreNodeId, Rational)
  nodeStake :: NodeId -> Maybe (CoreNodeId, Ratio Integer)
nodeStake (RelayId Word64
_) = Maybe (CoreNodeId, Ratio Integer)
forall a. Maybe a
Nothing
  nodeStake (CoreId CoreNodeId
i) = (CoreNodeId, Ratio Integer) -> Maybe (CoreNodeId, Ratio Integer)
forall a. a -> Maybe a
Just (CoreNodeId
i, Ratio Integer -> Ratio Integer
forall a. Fractional a => a -> a
recip (Int -> Ratio Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n))

  n :: Int
n = [NodeId] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([NodeId] -> Int) -> [NodeId] -> Int
forall a b. (a -> b) -> a -> b
$ (NodeId -> Bool) -> [NodeId] -> [NodeId]
forall a. (a -> Bool) -> [a] -> [a]
filter NodeId -> Bool
isCore ([NodeId] -> [NodeId]) -> [NodeId] -> [NodeId]
forall a b. (a -> b) -> a -> b
$ Map Addr NodeId -> [NodeId]
forall k a. Map k a -> [a]
Map.elems Map Addr NodeId
ad

  isCore :: NodeId -> Bool
  isCore :: NodeId -> Bool
isCore CoreId{} = Bool
True
  isCore RelayId{} = Bool
False

-- | Genesis stake distribution
genesisStakeDist :: AddrDist -> StakeDist
genesisStakeDist :: Map Addr NodeId -> StakeDist
genesisStakeDist Map Addr NodeId
addrDist =
  Map StakeHolder Word -> StakeDist
relativeStakes (Map Addr NodeId -> Utxo -> Map StakeHolder Word
totalStakes Map Addr NodeId
addrDist (Map Addr NodeId -> Utxo
genesisUtxo Map Addr NodeId
addrDist))