{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}

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 Rational
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 :: Rational -> CoreNodeId -> StakeDist -> Rational
stakeWithDefault Rational
d CoreNodeId
n = Rational -> CoreNodeId -> Map CoreNodeId Rational -> Rational
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Rational
d CoreNodeId
n (Map CoreNodeId Rational -> Rational)
-> (StakeDist -> Map CoreNodeId Rational) -> StakeDist -> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StakeDist -> Map CoreNodeId Rational
stakeDistToMap

relativeStakes :: Map StakeHolder Amount -> StakeDist
relativeStakes :: Map StakeHolder Amount -> StakeDist
relativeStakes Map StakeHolder Amount
m = Map CoreNodeId Rational -> StakeDist
StakeDist (Map CoreNodeId Rational -> StakeDist)
-> Map CoreNodeId Rational -> StakeDist
forall a b. (a -> b) -> a -> b
$
   let totalStake :: Rational
totalStake    = Amount -> Rational
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Amount -> Rational) -> Amount -> Rational
forall a b. (a -> b) -> a -> b
$ [Amount] -> Amount
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Amount] -> Amount) -> [Amount] -> Amount
forall a b. (a -> b) -> a -> b
$ Map StakeHolder Amount -> [Amount]
forall k a. Map k a -> [a]
Map.elems Map StakeHolder Amount
m
   in  [(CoreNodeId, Rational)] -> Map CoreNodeId Rational
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (CoreNodeId
nid, Amount -> Rational
forall a b. (Integral a, Num b) => a -> b
fromIntegral Amount
stake Rational -> Rational -> Rational
forall a. Fractional a => a -> a -> a
/ Rational
totalStake)
                    | (StakeCore CoreNodeId
nid, Amount
stake) <- Map StakeHolder Amount -> [(StakeHolder, Amount)]
forall k a. Map k a -> [(k, a)]
Map.toList Map StakeHolder Amount
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 Amount
totalStakes Map Addr NodeId
addrDist = (Map StakeHolder Amount -> TxOut -> Map StakeHolder Amount)
-> Map StakeHolder Amount -> Utxo -> Map StakeHolder Amount
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 Amount -> TxOut -> Map StakeHolder Amount
f Map StakeHolder Amount
forall k a. Map k a
Map.empty
 where
   f :: Map StakeHolder Amount -> TxOut -> Map StakeHolder Amount
   f :: Map StakeHolder Amount -> TxOut -> Map StakeHolder Amount
f Map StakeHolder Amount
m (Addr
a, Amount
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) -> (Amount -> Amount -> Amount)
-> StakeHolder
-> Amount
-> Map StakeHolder Amount
-> Map StakeHolder Amount
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Amount -> Amount -> Amount
forall a. Num a => a -> a -> a
(+) (CoreNodeId -> StakeHolder
StakeCore CoreNodeId
nid)    Amount
stake Map StakeHolder Amount
m
       Maybe NodeId
_                 -> (Amount -> Amount -> Amount)
-> StakeHolder
-> Amount
-> Map StakeHolder Amount
-> Map StakeHolder Amount
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Amount -> Amount -> Amount
forall a. Num a => a -> a -> a
(+) StakeHolder
StakeEverybodyElse Amount
stake Map StakeHolder Amount
m

-- | Stake distribution where every address has equal state
equalStakeDist :: AddrDist -> StakeDist
equalStakeDist :: Map Addr NodeId -> StakeDist
equalStakeDist Map Addr NodeId
ad =
    Map CoreNodeId Rational -> StakeDist
StakeDist (Map CoreNodeId Rational -> StakeDist)
-> Map CoreNodeId Rational -> StakeDist
forall a b. (a -> b) -> a -> b
$
    [(CoreNodeId, Rational)] -> Map CoreNodeId Rational
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(CoreNodeId, Rational)] -> Map CoreNodeId Rational)
-> [(CoreNodeId, Rational)] -> Map CoreNodeId Rational
forall a b. (a -> b) -> a -> b
$
    ((Addr, NodeId) -> Maybe (CoreNodeId, Rational))
-> [(Addr, NodeId)] -> [(CoreNodeId, Rational)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (NodeId -> Maybe (CoreNodeId, Rational)
nodeStake (NodeId -> Maybe (CoreNodeId, Rational))
-> ((Addr, NodeId) -> NodeId)
-> (Addr, NodeId)
-> Maybe (CoreNodeId, Rational)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Addr, NodeId) -> NodeId
forall a b. (a, b) -> b
snd) ([(Addr, NodeId)] -> [(CoreNodeId, Rational)])
-> [(Addr, NodeId)] -> [(CoreNodeId, Rational)]
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, Rational)
nodeStake (RelayId Word64
_) = Maybe (CoreNodeId, Rational)
forall a. Maybe a
Nothing
    nodeStake (CoreId CoreNodeId
i)  = (CoreNodeId, Rational) -> Maybe (CoreNodeId, Rational)
forall a. a -> Maybe a
Just (CoreNodeId
i, Rational -> Rational
forall a. Fractional a => a -> a
recip (Int -> Rational
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 Amount -> StakeDist
relativeStakes (Map Addr NodeId -> Utxo -> Map StakeHolder Amount
totalStakes Map Addr NodeId
addrDist (Map Addr NodeId -> Utxo
genesisUtxo Map Addr NodeId
addrDist))