{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableSuperClasses #-}
-- some GHC bug wrongfully complains about CanGetInstantStake constraint being redundant.
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Cardano.Ledger.State.Account (
  CanGetAccounts (..),
  CanSetAccounts (..),
  EraAccounts (..),
  lookupAccountState,
  updateLookupAccountState,
  isAccountRegistered,
  adjustAccountState,
  lookupStakePoolDelegation,
  sumBalancesAccounts,
  sumDepositsAccounts,
  addToBalanceAccounts,
  withdrawalsThatDoNotDrainAccounts,
  drainAccounts,
  removeStakePoolDelegations,
) where

import Cardano.Ledger.Address (RewardAccount (..), Withdrawals (..))
import Cardano.Ledger.BaseTypes (Network)
import Cardano.Ledger.Binary
import Cardano.Ledger.Coin
import Cardano.Ledger.Compactible
import Cardano.Ledger.Core
import Cardano.Ledger.Credential
import Control.DeepSeq (NFData)
import Control.Exception (assert)
import Data.Aeson (ToJSON)
import Data.Bifunctor (Bifunctor (..))
import Data.Default (Default)
import Data.Foldable (foldMap')
import Data.Kind (Type)
import qualified Data.Map.Merge.Strict as Map
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Lens.Micro
import NoThunks.Class (NoThunks)

class CanGetAccounts t where
  accountsG :: SimpleGetter (t era) (Accounts era)
  default accountsG :: CanSetAccounts t => SimpleGetter (t era) (Accounts era)
  accountsG = (Accounts era -> Const r (Accounts era))
-> t era -> Const r (t era)
forall era. Lens' (t era) (Accounts era)
forall (t :: * -> *) era.
CanSetAccounts t =>
Lens' (t era) (Accounts era)
accountsL
  {-# INLINE accountsG #-}

class CanGetAccounts t => CanSetAccounts t where
  accountsL :: Lens' (t era) (Accounts era)

class
  ( Era era
  , Eq (Accounts era)
  , Show (Accounts era)
  , Default (Accounts era)
  , EncCBOR (Accounts era)
  , DecShareCBOR (Accounts era)
  , Share (Accounts era)
      ~ ( Interns (Credential Staking)
        , Interns (KeyHash StakePool)
        , Interns (Credential DRepRole)
        )
  , ToJSON (Accounts era)
  , NFData (Accounts era)
  , NoThunks (Accounts era)
  , Eq (AccountState era)
  , Show (AccountState era)
  , NFData (AccountState era)
  , NoThunks (AccountState era)
  ) =>
  EraAccounts era
  where
  type AccountState era = (r :: Type) | r -> era
  type Accounts era = (r :: Type) | r -> era

  -- | Add `AccountState` to `Accounts`. There are no checks whether account is already registered
  -- or not.
  addAccountState :: Credential 'Staking -> AccountState era -> Accounts era -> Accounts era

  accountsMapL :: Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))

  balanceAccountStateL :: Lens' (AccountState era) (CompactForm Coin)

  depositAccountStateL :: Lens' (AccountState era) (CompactForm Coin)

  stakePoolDelegationAccountStateL :: Lens' (AccountState era) (Maybe (KeyHash 'StakePool))

  -- | Remove the account from the state. Note that it is not capable of affecting state for DReps
  -- and StakePools, those have to be handled separately.
  --
  -- There is no counterpart for registering an account, because different eras require different
  -- information. However for testing purposed there is
  -- `Test.Cardano.Ledger.Era.registerTestAccount` that can be used for all eras.
  unregisterAccount ::
    -- | Credential to unregister
    Credential 'Staking ->
    -- | `Accounts` to remove the account state from
    Accounts era ->
    -- | Returns `Just` whenever account was registered and `Nothing` otherwise. Produced `Accounts`
    -- will have the account state removed, if it was present there to begin with.
    (Maybe (AccountState era), Accounts era)

sumBalancesAccounts :: EraAccounts era => Accounts era -> Coin
sumBalancesAccounts :: forall era. EraAccounts era => Accounts era -> Coin
sumBalancesAccounts Accounts era
accounts =
  CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact (CompactForm Coin -> Coin) -> CompactForm Coin -> Coin
forall a b. (a -> b) -> a -> b
$ (AccountState era -> CompactForm Coin)
-> Map (Credential 'Staking) (AccountState era) -> CompactForm Coin
forall m a.
Monoid m =>
(a -> m) -> Map (Credential 'Staking) a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' (AccountState era
-> Getting (CompactForm Coin) (AccountState era) (CompactForm Coin)
-> CompactForm Coin
forall s a. s -> Getting a s a -> a
^. Getting (CompactForm Coin) (AccountState era) (CompactForm Coin)
forall era.
EraAccounts era =>
Lens' (AccountState era) (CompactForm Coin)
Lens' (AccountState era) (CompactForm Coin)
balanceAccountStateL) (Map (Credential 'Staking) (AccountState era) -> CompactForm Coin)
-> Map (Credential 'Staking) (AccountState era) -> CompactForm Coin
forall a b. (a -> b) -> a -> b
$ Accounts era
accounts Accounts era
-> Getting
     (Map (Credential 'Staking) (AccountState era))
     (Accounts era)
     (Map (Credential 'Staking) (AccountState era))
-> Map (Credential 'Staking) (AccountState era)
forall s a. s -> Getting a s a -> a
^. Getting
  (Map (Credential 'Staking) (AccountState era))
  (Accounts era)
  (Map (Credential 'Staking) (AccountState era))
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL

sumDepositsAccounts :: EraAccounts era => Accounts era -> Coin
sumDepositsAccounts :: forall era. EraAccounts era => Accounts era -> Coin
sumDepositsAccounts Accounts era
accounts =
  CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact (CompactForm Coin -> Coin) -> CompactForm Coin -> Coin
forall a b. (a -> b) -> a -> b
$ (AccountState era -> CompactForm Coin)
-> Map (Credential 'Staking) (AccountState era) -> CompactForm Coin
forall m a.
Monoid m =>
(a -> m) -> Map (Credential 'Staking) a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' (AccountState era
-> Getting (CompactForm Coin) (AccountState era) (CompactForm Coin)
-> CompactForm Coin
forall s a. s -> Getting a s a -> a
^. Getting (CompactForm Coin) (AccountState era) (CompactForm Coin)
forall era.
EraAccounts era =>
Lens' (AccountState era) (CompactForm Coin)
Lens' (AccountState era) (CompactForm Coin)
depositAccountStateL) (Map (Credential 'Staking) (AccountState era) -> CompactForm Coin)
-> Map (Credential 'Staking) (AccountState era) -> CompactForm Coin
forall a b. (a -> b) -> a -> b
$ Accounts era
accounts Accounts era
-> Getting
     (Map (Credential 'Staking) (AccountState era))
     (Accounts era)
     (Map (Credential 'Staking) (AccountState era))
-> Map (Credential 'Staking) (AccountState era)
forall s a. s -> Getting a s a -> a
^. Getting
  (Map (Credential 'Staking) (AccountState era))
  (Accounts era)
  (Map (Credential 'Staking) (AccountState era))
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL

-- | Top up balance in accounts
--
-- /Warning/ - it is an error to try to increase a balance of an account that is not present in
-- `Accounts`
addToBalanceAccounts ::
  EraAccounts era =>
  -- | Map containing amounts that the balance in the account should be increased by. It is
  -- important to ensure that all of the credentials in this Map are actually registered.
  Map (Credential 'Staking) (CompactForm Coin) ->
  -- | Accounts that will have their balance increased.
  Accounts era ->
  Accounts era
addToBalanceAccounts :: forall era.
EraAccounts era =>
Map (Credential 'Staking) (CompactForm Coin)
-> Accounts era -> Accounts era
addToBalanceAccounts Map (Credential 'Staking) (CompactForm Coin)
addBalanceMap Accounts era
accounts =
  Accounts era
accounts
    Accounts era -> (Accounts era -> Accounts era) -> Accounts era
forall a b. a -> (a -> b) -> b
& (Map (Credential 'Staking) (AccountState era)
 -> Identity (Map (Credential 'Staking) (AccountState era)))
-> Accounts era -> Identity (Accounts era)
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL
      ((Map (Credential 'Staking) (AccountState era)
  -> Identity (Map (Credential 'Staking) (AccountState era)))
 -> Accounts era -> Identity (Accounts era))
-> (Map (Credential 'Staking) (AccountState era)
    -> Map (Credential 'Staking) (AccountState era))
-> Accounts era
-> Accounts era
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ SimpleWhenMissing
  (Credential 'Staking) (CompactForm Coin) (AccountState era)
-> SimpleWhenMissing
     (Credential 'Staking) (AccountState era) (AccountState era)
-> SimpleWhenMatched
     (Credential 'Staking)
     (CompactForm Coin)
     (AccountState era)
     (AccountState era)
-> Map (Credential 'Staking) (CompactForm Coin)
-> Map (Credential 'Staking) (AccountState era)
-> Map (Credential 'Staking) (AccountState era)
forall k a c b.
Ord k =>
SimpleWhenMissing k a c
-> SimpleWhenMissing k b c
-> SimpleWhenMatched k a b c
-> Map k a
-> Map k b
-> Map k c
Map.merge
        -- We have an assert here, since this should never be the case that we try to add to a
        -- balance of a non-registered account
        ((Credential 'Staking
 -> CompactForm Coin -> Maybe (AccountState era))
-> SimpleWhenMissing
     (Credential 'Staking) (CompactForm Coin) (AccountState era)
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> Maybe y) -> WhenMissing f k x y
Map.mapMaybeMissing (\Credential 'Staking
_ CompactForm Coin
_ -> Bool -> Maybe (AccountState era) -> Maybe (AccountState era)
forall a. (?callStack::CallStack) => Bool -> a -> a
assert Bool
False Maybe (AccountState era)
forall a. Maybe a
Nothing))
        SimpleWhenMissing
  (Credential 'Staking) (AccountState era) (AccountState era)
forall (f :: * -> *) k x. Applicative f => WhenMissing f k x x
Map.preserveMissing
        ((Credential 'Staking
 -> CompactForm Coin -> AccountState era -> AccountState era)
-> SimpleWhenMatched
     (Credential 'Staking)
     (CompactForm Coin)
     (AccountState era)
     (AccountState era)
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> z) -> WhenMatched f k x y z
Map.zipWithMatched (\Credential 'Staking
_ CompactForm Coin
balanceToAdd -> (CompactForm Coin -> Identity (CompactForm Coin))
-> AccountState era -> Identity (AccountState era)
forall era.
EraAccounts era =>
Lens' (AccountState era) (CompactForm Coin)
Lens' (AccountState era) (CompactForm Coin)
balanceAccountStateL ((CompactForm Coin -> Identity (CompactForm Coin))
 -> AccountState era -> Identity (AccountState era))
-> CompactForm Coin -> AccountState era -> AccountState era
forall a s t. Monoid a => ASetter s t a a -> a -> s -> t
<>~ CompactForm Coin
balanceToAdd))
        Map (Credential 'Staking) (CompactForm Coin)
addBalanceMap

-- | Lookup an account state by its credential. Returns Nothing if such account is not registrered
lookupAccountState ::
  EraAccounts era => Credential 'Staking -> Accounts era -> Maybe (AccountState era)
lookupAccountState :: forall era.
EraAccounts era =>
Credential 'Staking -> Accounts era -> Maybe (AccountState era)
lookupAccountState Credential 'Staking
cred Accounts era
accounts = Credential 'Staking
-> Map (Credential 'Staking) (AccountState era)
-> Maybe (AccountState era)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Credential 'Staking
cred (Accounts era
accounts Accounts era
-> Getting
     (Map (Credential 'Staking) (AccountState era))
     (Accounts era)
     (Map (Credential 'Staking) (AccountState era))
-> Map (Credential 'Staking) (AccountState era)
forall s a. s -> Getting a s a -> a
^. Getting
  (Map (Credential 'Staking) (AccountState era))
  (Accounts era)
  (Map (Credential 'Staking) (AccountState era))
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL)

-- | Update account state. Returns Nothing if the value is not present and modified value otherwise
updateLookupAccountState ::
  EraAccounts era =>
  (AccountState era -> AccountState era) ->
  Credential 'Staking ->
  Accounts era ->
  (Maybe (AccountState era), Accounts era)
updateLookupAccountState :: forall era.
EraAccounts era =>
(AccountState era -> AccountState era)
-> Credential 'Staking
-> Accounts era
-> (Maybe (AccountState era), Accounts era)
updateLookupAccountState AccountState era -> AccountState era
f Credential 'Staking
cred Accounts era
accounts =
  case (Credential 'Staking
 -> AccountState era -> Maybe (AccountState era))
-> Credential 'Staking
-> Map (Credential 'Staking) (AccountState era)
-> (Maybe (AccountState era),
    Map (Credential 'Staking) (AccountState era))
forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
Map.updateLookupWithKey (\Credential 'Staking
_ -> AccountState era -> Maybe (AccountState era)
forall a. a -> Maybe a
Just (AccountState era -> Maybe (AccountState era))
-> (AccountState era -> AccountState era)
-> AccountState era
-> Maybe (AccountState era)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AccountState era -> AccountState era
f) Credential 'Staking
cred (Accounts era
accounts Accounts era
-> Getting
     (Map (Credential 'Staking) (AccountState era))
     (Accounts era)
     (Map (Credential 'Staking) (AccountState era))
-> Map (Credential 'Staking) (AccountState era)
forall s a. s -> Getting a s a -> a
^. Getting
  (Map (Credential 'Staking) (AccountState era))
  (Accounts era)
  (Map (Credential 'Staking) (AccountState era))
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL) of
    (Maybe (AccountState era)
res, Map (Credential 'Staking) (AccountState era)
accountsMap) -> (Maybe (AccountState era)
res, Accounts era
accounts Accounts era -> (Accounts era -> Accounts era) -> Accounts era
forall a b. a -> (a -> b) -> b
& (Map (Credential 'Staking) (AccountState era)
 -> Identity (Map (Credential 'Staking) (AccountState era)))
-> Accounts era -> Identity (Accounts era)
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL ((Map (Credential 'Staking) (AccountState era)
  -> Identity (Map (Credential 'Staking) (AccountState era)))
 -> Accounts era -> Identity (Accounts era))
-> Map (Credential 'Staking) (AccountState era)
-> Accounts era
-> Accounts era
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Map (Credential 'Staking) (AccountState era)
accountsMap)

-- | Check whether account for this staking credential is registered
isAccountRegistered :: EraAccounts era => Credential 'Staking -> Accounts era -> Bool
isAccountRegistered :: forall era.
EraAccounts era =>
Credential 'Staking -> Accounts era -> Bool
isAccountRegistered Credential 'Staking
cred Accounts era
accounts = Credential 'Staking
-> Map (Credential 'Staking) (AccountState era) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member Credential 'Staking
cred (Accounts era
accounts Accounts era
-> Getting
     (Map (Credential 'Staking) (AccountState era))
     (Accounts era)
     (Map (Credential 'Staking) (AccountState era))
-> Map (Credential 'Staking) (AccountState era)
forall s a. s -> Getting a s a -> a
^. Getting
  (Map (Credential 'Staking) (AccountState era))
  (Accounts era)
  (Map (Credential 'Staking) (AccountState era))
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL)

adjustAccountState ::
  EraAccounts era =>
  (AccountState era -> AccountState era) -> Credential 'Staking -> Accounts era -> Accounts era
adjustAccountState :: forall era.
EraAccounts era =>
(AccountState era -> AccountState era)
-> Credential 'Staking -> Accounts era -> Accounts era
adjustAccountState AccountState era -> AccountState era
cred Credential 'Staking
f = (Map (Credential 'Staking) (AccountState era)
 -> Identity (Map (Credential 'Staking) (AccountState era)))
-> Accounts era -> Identity (Accounts era)
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL ((Map (Credential 'Staking) (AccountState era)
  -> Identity (Map (Credential 'Staking) (AccountState era)))
 -> Accounts era -> Identity (Accounts era))
-> (Map (Credential 'Staking) (AccountState era)
    -> Map (Credential 'Staking) (AccountState era))
-> Accounts era
-> Accounts era
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (AccountState era -> AccountState era)
-> Credential 'Staking
-> Map (Credential 'Staking) (AccountState era)
-> Map (Credential 'Staking) (AccountState era)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust AccountState era -> AccountState era
cred Credential 'Staking
f

-- | In case when account state is registered and it is delegated to a stake pool this function
-- will return that delegation.
lookupStakePoolDelegation ::
  EraAccounts era =>
  Credential 'Staking ->
  Accounts era ->
  Maybe (KeyHash 'StakePool)
lookupStakePoolDelegation :: forall era.
EraAccounts era =>
Credential 'Staking -> Accounts era -> Maybe (KeyHash 'StakePool)
lookupStakePoolDelegation Credential 'Staking
cred Accounts era
accounts =
  Credential 'Staking -> Accounts era -> Maybe (AccountState era)
forall era.
EraAccounts era =>
Credential 'Staking -> Accounts era -> Maybe (AccountState era)
lookupAccountState Credential 'Staking
cred Accounts era
accounts
    Maybe (AccountState era)
-> (AccountState era -> Maybe (KeyHash 'StakePool))
-> Maybe (KeyHash 'StakePool)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (AccountState era
-> Getting
     (Maybe (KeyHash 'StakePool))
     (AccountState era)
     (Maybe (KeyHash 'StakePool))
-> Maybe (KeyHash 'StakePool)
forall s a. s -> Getting a s a -> a
^. Getting
  (Maybe (KeyHash 'StakePool))
  (AccountState era)
  (Maybe (KeyHash 'StakePool))
forall era.
EraAccounts era =>
Lens' (AccountState era) (Maybe (KeyHash 'StakePool))
Lens' (AccountState era) (Maybe (KeyHash 'StakePool))
stakePoolDelegationAccountStateL)

-- | This function returns `Nothing` iff all of the accounts that withdrawals
-- are trying to drain are indeed registered and all of the amounts in the
-- withdrawals match the respective balances exactly. It returns a 2-tuple where
-- the `fst` is withdrawals with missing reward accounts or the wrong network,
-- and `snd` is incomplete withdrawals.
--
-- NOTE: We simply `checkBadWithdrawals` to avoid allocating new variables for
-- the most likely case.
withdrawalsThatDoNotDrainAccounts ::
  EraAccounts era =>
  Withdrawals ->
  Network ->
  Accounts era ->
  -- | invalid withdrawal = that which does not have a reward account or is in
  -- the wrong network.
  -- incomplete withdrawal = that which does not withdraw the exact account
  -- balance.
  Maybe (Withdrawals, Withdrawals)
withdrawalsThatDoNotDrainAccounts :: forall era.
EraAccounts era =>
Withdrawals
-> Network -> Accounts era -> Maybe (Withdrawals, Withdrawals)
withdrawalsThatDoNotDrainAccounts (Withdrawals Map RewardAccount Coin
withdrawals) Network
networkId Accounts era
accounts
  -- @withdrawals@ is small and @accounts@ big, better to traverse the former than the latter.
  | (RewardAccount -> Coin -> Bool -> Bool)
-> Bool -> Map RewardAccount Coin -> Bool
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey RewardAccount -> Coin -> Bool -> Bool
checkBadWithdrawals Bool
True Map RewardAccount Coin
withdrawals = Maybe (Withdrawals, Withdrawals)
forall a. Maybe a
Nothing
  | Bool
otherwise =
      (Withdrawals, Withdrawals) -> Maybe (Withdrawals, Withdrawals)
forall a. a -> Maybe a
Just ((Withdrawals, Withdrawals) -> Maybe (Withdrawals, Withdrawals))
-> (Withdrawals, Withdrawals) -> Maybe (Withdrawals, Withdrawals)
forall a b. (a -> b) -> a -> b
$
        (Map RewardAccount Coin -> Withdrawals)
-> (Map RewardAccount Coin -> Withdrawals)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> (Withdrawals, Withdrawals)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Map RewardAccount Coin -> Withdrawals
Withdrawals Map RewardAccount Coin -> Withdrawals
Withdrawals ((Map RewardAccount Coin, Map RewardAccount Coin)
 -> (Withdrawals, Withdrawals))
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> (Withdrawals, Withdrawals)
forall a b. (a -> b) -> a -> b
$
          (RewardAccount
 -> Coin
 -> (Map RewardAccount Coin, Map RewardAccount Coin)
 -> (Map RewardAccount Coin, Map RewardAccount Coin))
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> Map RewardAccount Coin
-> (Map RewardAccount Coin, Map RewardAccount Coin)
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey RewardAccount
-> Coin
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
collectBadWithdrawals (Map RewardAccount Coin
forall k a. Map k a
Map.empty, Map RewardAccount Coin
forall k a. Map k a
Map.empty) Map RewardAccount Coin
withdrawals
  where
    checkBadWithdrawals :: RewardAccount -> Coin -> Bool -> Bool
checkBadWithdrawals RewardAccount
rewardAccount Coin
withdrawalAmount Bool
noBadWithdrawals =
      Bool
noBadWithdrawals Bool -> Bool -> Bool
&& RewardAccount -> Coin -> Bool
isGoodWithdrawal RewardAccount
rewardAccount Coin
withdrawalAmount
    collectBadWithdrawals :: RewardAccount
-> Coin
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
collectBadWithdrawals RewardAccount
rewardAccount Coin
withdrawalAmount accum :: (Map RewardAccount Coin, Map RewardAccount Coin)
accum@(!Map RewardAccount Coin
_, !Map RewardAccount Coin
_) =
      case RewardAccount -> Maybe (AccountState era)
lookupAccount RewardAccount
rewardAccount of
        Maybe (AccountState era)
Nothing -> (Map RewardAccount Coin -> Map RewardAccount Coin)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (RewardAccount
-> Coin -> Map RewardAccount Coin -> Map RewardAccount Coin
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert RewardAccount
rewardAccount Coin
withdrawalAmount) (Map RewardAccount Coin, Map RewardAccount Coin)
accum
        Just AccountState era
account
          | Coin -> AccountState era -> Bool
forall {era}.
(Assert
   (OrdCond
      (CmpNat (ProtVerLow era) (ProtVerHigh era)) 'True 'True 'False)
   (TypeError ...),
 Assert
   (OrdCond (CmpNat 0 (ProtVerLow era)) 'True 'True 'False)
   (TypeError ...),
 Assert
   (OrdCond (CmpNat 0 (ProtVerHigh era)) 'True 'True 'False)
   (TypeError ...),
 EraAccounts era) =>
Coin -> AccountState era -> Bool
isBalanceZero Coin
withdrawalAmount AccountState era
account -> (Map RewardAccount Coin, Map RewardAccount Coin)
accum
          | Bool
otherwise -> (Map RewardAccount Coin -> Map RewardAccount Coin)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
-> (Map RewardAccount Coin, Map RewardAccount Coin)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (RewardAccount
-> Coin -> Map RewardAccount Coin -> Map RewardAccount Coin
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert RewardAccount
rewardAccount Coin
withdrawalAmount) (Map RewardAccount Coin, Map RewardAccount Coin)
accum
    isGoodWithdrawal :: RewardAccount -> Coin -> Bool
isGoodWithdrawal RewardAccount
rewardAccount Coin
withdrawalAmount =
      Bool
-> (AccountState era -> Bool) -> Maybe (AccountState era) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Coin -> AccountState era -> Bool
forall {era}.
(Assert
   (OrdCond
      (CmpNat (ProtVerLow era) (ProtVerHigh era)) 'True 'True 'False)
   (TypeError ...),
 Assert
   (OrdCond (CmpNat 0 (ProtVerLow era)) 'True 'True 'False)
   (TypeError ...),
 Assert
   (OrdCond (CmpNat 0 (ProtVerHigh era)) 'True 'True 'False)
   (TypeError ...),
 EraAccounts era) =>
Coin -> AccountState era -> Bool
isBalanceZero Coin
withdrawalAmount) (RewardAccount -> Maybe (AccountState era)
lookupAccount RewardAccount
rewardAccount)
    isBalanceZero :: Coin -> AccountState era -> Bool
isBalanceZero Coin
withdrawalAmount AccountState era
accountState =
      Coin
withdrawalAmount Coin -> Coin -> Bool
forall a. Eq a => a -> a -> Bool
== CompactForm Coin -> Coin
forall a. Compactible a => CompactForm a -> a
fromCompact (AccountState era
accountState AccountState era
-> Getting (CompactForm Coin) (AccountState era) (CompactForm Coin)
-> CompactForm Coin
forall s a. s -> Getting a s a -> a
^. Getting (CompactForm Coin) (AccountState era) (CompactForm Coin)
forall era.
EraAccounts era =>
Lens' (AccountState era) (CompactForm Coin)
Lens' (AccountState era) (CompactForm Coin)
balanceAccountStateL)
    lookupAccount :: RewardAccount -> Maybe (AccountState era)
lookupAccount RewardAccount {Credential 'Staking
raCredential :: Credential 'Staking
raCredential :: RewardAccount -> Credential 'Staking
raCredential, Network
raNetwork :: Network
raNetwork :: RewardAccount -> Network
raNetwork}
      | Network
raNetwork Network -> Network -> Bool
forall a. Eq a => a -> a -> Bool
== Network
networkId = Credential 'Staking -> Accounts era -> Maybe (AccountState era)
forall era.
EraAccounts era =>
Credential 'Staking -> Accounts era -> Maybe (AccountState era)
lookupAccountState Credential 'Staking
raCredential Accounts era
accounts
      | Bool
otherwise = Maybe (AccountState era)
forall a. Maybe a
Nothing

-- | Reset balances to zero for all accounts that are specified in the supplied `Withdrawals`.
--
-- /Note/ - There are no checks that withdrawals mention only registered accounts with correct
-- `NetworkId`. Nor there are any checks that amounts in withdrawals match up the balance in the
-- corresponding accounts. Use `withdrawalsThatDoNotDrainAccounts` to verify that calling
-- `drainAccounts` is actually safe on the supplied arguments
drainAccounts ::
  EraAccounts era =>
  Withdrawals ->
  Accounts era ->
  Accounts era
drainAccounts :: forall era.
EraAccounts era =>
Withdrawals -> Accounts era -> Accounts era
drainAccounts (Withdrawals Map RewardAccount Coin
withdrawalsMap) Accounts era
accounts =
  Accounts era
accounts
    Accounts era -> (Accounts era -> Accounts era) -> Accounts era
forall a b. a -> (a -> b) -> b
& (Map (Credential 'Staking) (AccountState era)
 -> Identity (Map (Credential 'Staking) (AccountState era)))
-> Accounts era -> Identity (Accounts era)
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL ((Map (Credential 'Staking) (AccountState era)
  -> Identity (Map (Credential 'Staking) (AccountState era)))
 -> Accounts era -> Identity (Accounts era))
-> (Map (Credential 'Staking) (AccountState era)
    -> Map (Credential 'Staking) (AccountState era))
-> Accounts era
-> Accounts era
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ \Map (Credential 'Staking) (AccountState era)
accountsMap ->
      (RewardAccount
 -> Coin
 -> Map (Credential 'Staking) (AccountState era)
 -> Map (Credential 'Staking) (AccountState era))
-> Map (Credential 'Staking) (AccountState era)
-> Map RewardAccount Coin
-> Map (Credential 'Staking) (AccountState era)
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey'
        (\RewardAccount
ra Coin
_withdrawalAmount -> (AccountState era -> AccountState era)
-> Credential 'Staking
-> Map (Credential 'Staking) (AccountState era)
-> Map (Credential 'Staking) (AccountState era)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust ((CompactForm Coin -> Identity (CompactForm Coin))
-> AccountState era -> Identity (AccountState era)
forall era.
EraAccounts era =>
Lens' (AccountState era) (CompactForm Coin)
Lens' (AccountState era) (CompactForm Coin)
balanceAccountStateL ((CompactForm Coin -> Identity (CompactForm Coin))
 -> AccountState era -> Identity (AccountState era))
-> CompactForm Coin -> AccountState era -> AccountState era
forall s t a b. ASetter s t a b -> b -> s -> t
.~ CompactForm Coin
forall a. Monoid a => a
mempty) (RewardAccount -> Credential 'Staking
raCredential RewardAccount
ra))
        Map (Credential 'Staking) (AccountState era)
accountsMap
        Map RewardAccount Coin
withdrawalsMap

-- TODO: This is an expensive operation, since it iterates over the whole accountsMap. We need to
-- start keeping track of all delegations to the stake pool in its state, then we would be able to
-- switch from `Set (KeyHash 'StakePool)` to `Map (KeyHash 'StakePool) (Set (Credential Staking))`
-- and drastically speed up this operation.

-- | Remove delegations for the supplied Stake
removeStakePoolDelegations ::
  EraAccounts era => Set (KeyHash 'StakePool) -> Accounts era -> Accounts era
removeStakePoolDelegations :: forall era.
EraAccounts era =>
Set (KeyHash 'StakePool) -> Accounts era -> Accounts era
removeStakePoolDelegations Set (KeyHash 'StakePool)
stakeDelegationsToRemove Accounts era
accounts =
  Accounts era
accounts Accounts era -> (Accounts era -> Accounts era) -> Accounts era
forall a b. a -> (a -> b) -> b
& (Map (Credential 'Staking) (AccountState era)
 -> Identity (Map (Credential 'Staking) (AccountState era)))
-> Accounts era -> Identity (Accounts era)
forall era.
EraAccounts era =>
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
Lens' (Accounts era) (Map (Credential 'Staking) (AccountState era))
accountsMapL ((Map (Credential 'Staking) (AccountState era)
  -> Identity (Map (Credential 'Staking) (AccountState era)))
 -> Accounts era -> Identity (Accounts era))
-> (Map (Credential 'Staking) (AccountState era)
    -> Map (Credential 'Staking) (AccountState era))
-> Accounts era
-> Accounts era
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (AccountState era -> AccountState era)
-> Map (Credential 'Staking) (AccountState era)
-> Map (Credential 'Staking) (AccountState era)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map AccountState era -> AccountState era
clearAccountStateDelegation
  where
    clearAccountStateDelegation :: AccountState era -> AccountState era
clearAccountStateDelegation =
      (Maybe (KeyHash 'StakePool)
 -> Identity (Maybe (KeyHash 'StakePool)))
-> AccountState era -> Identity (AccountState era)
forall era.
EraAccounts era =>
Lens' (AccountState era) (Maybe (KeyHash 'StakePool))
Lens' (AccountState era) (Maybe (KeyHash 'StakePool))
stakePoolDelegationAccountStateL ((Maybe (KeyHash 'StakePool)
  -> Identity (Maybe (KeyHash 'StakePool)))
 -> AccountState era -> Identity (AccountState era))
-> (Maybe (KeyHash 'StakePool) -> Maybe (KeyHash 'StakePool))
-> AccountState era
-> AccountState era
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ \case
        Just KeyHash 'StakePool
poolId | KeyHash 'StakePool
poolId KeyHash 'StakePool -> Set (KeyHash 'StakePool) -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set (KeyHash 'StakePool)
stakeDelegationsToRemove -> Maybe (KeyHash 'StakePool)
forall a. Maybe a
Nothing
        Maybe (KeyHash 'StakePool)
delegation -> Maybe (KeyHash 'StakePool)
delegation