Update conway Transition to set the pool delegations

and add a test for it

Author: teodanciu

eras/conway/impl/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 1.21.0.0
 
+* Expose `conwayRegisterInitialAccounts`
 * Move withdrawal-validation and DRep expiry updates from `CERTS` to `LEDGER` starting protocol version 11.
   - Add `ConwayWithdrawalsMissingAccounts` and `ConwayIncompleteWithdrawals` to `ConwayLedgerPredFailure`.
   - Add `hardforkConwayMoveWithdrawalsAndDRepChecksToLedgerRule` to `Conway.Era`.

eras/conway/impl/src/Cardano/Ledger/Conway/Transition.hs

@@ -16,6 +16,7 @@ module Cardano.Ledger.Conway.Transition (
   TransitionConfig (..),
   toConwayTransitionConfigPairs,
   registerDRepsThenDelegs,
+  conwayRegisterInitialAccounts,
   conwayRegisterInitialFundsThenStaking,
 ) where
 
@@ -44,6 +45,7 @@ import Data.Aeson (KeyValue (..))
 import Data.ListMap (ListMap)
 import qualified Data.ListMap as ListMap
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 import GHC.Generics
 import GHC.Stack
 import Lens.Micro
@@ -127,14 +129,20 @@ conwayRegisterInitialAccounts ::
   NewEpochState era
 conwayRegisterInitialAccounts ShelleyGenesisStaking {sgsStake} nes =
   nes
-    & nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL %~ \initAccounts ->
-      foldr registerAndDelegate initAccounts $ ListMap.toList sgsStake
+    & nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL .~ updatedAccounts
+    & nesEsL . esLStateL . lsCertStateL . certPStateL . psStakePoolsL .~ updatedStakePoolStates
   where
     stakePools = nes ^. nesEsL . esLStateL . lsCertStateL . certPStateL . psStakePoolsL
+    initialAccounts = nes ^. nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL
     deposit = compactCoinOrError $ nes ^. nesEsL . curPParamsEpochStateL . ppKeyDepositL
-    registerAndDelegate (stakeKeyHash, stakePool) !accounts
+
+    !(!updatedAccounts, !updatedStakePoolStates) =
+      foldr registerAndDelegate (initialAccounts, stakePools) (ListMap.toList sgsStake)
+    registerAndDelegate (stakeKeyHash, stakePool) (!accounts, !stakePoolMap)
       | stakePool `Map.member` stakePools =
-          registerConwayAccount (KeyHashObj stakeKeyHash) deposit (Just (DelegStake stakePool)) accounts
+          ( (registerConwayAccount (KeyHashObj stakeKeyHash) deposit (Just (DelegStake stakePool)) accounts)
+          , Map.adjust (spsDelegsL %~ Set.insert (KeyHashObj stakeKeyHash)) stakePool stakePoolMap
+          )
       | otherwise =
           error $
             "Invariant of a delegation of "

eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/Imp/DelegSpec.hs

@@ -28,15 +28,18 @@ import Cardano.Ledger.Conway.Core
 import Cardano.Ledger.Conway.Governance
 import Cardano.Ledger.Conway.Rules (ConwayDelegPredFailure (..))
 import Cardano.Ledger.Conway.State hiding (balance)
+import Cardano.Ledger.Conway.Transition (conwayRegisterInitialAccounts)
 import Cardano.Ledger.Conway.TxCert
 import Cardano.Ledger.Credential (Credential (..))
 import Cardano.Ledger.DRep
 import Cardano.Ledger.Plutus (
   SLanguage (..),
   hashPlutusScript,
  )
+import Cardano.Ledger.Shelley.Genesis (ShelleyGenesisStaking (..))
 import Cardano.Ledger.Shelley.LedgerState
 import Cardano.Ledger.Val (Val (..))
+import qualified Data.ListMap as LM
 import qualified Data.Map.Strict as Map
 import qualified Data.Sequence.Strict as SSeq
 import qualified Data.Set as Set
@@ -531,6 +534,29 @@ spec = do
       expectDelegatedToPool cred poolKh'
       expectNotDelegatedToPool cred poolKh
       expectDelegatedVote cred (DRepCredential drepCred)
+
+  it "Transition creates the delegations correctly" $ do
+    pool1 <- freshKeyHash >>= \kh -> kh <$ registerPool kh
+    pool2 <- freshKeyHash >>= \kh -> kh <$ registerPool kh
+    pool3 <- freshKeyHash >>= \kh -> kh <$ registerPool kh
+    poolParams <- freshKeyHash >>= \kh -> registerRewardAccount >>= freshPoolParams kh
+    deleg1 <- freshKeyHash >>= \kh -> kh <$ registerStakeCredential (KeyHashObj kh)
+    deleg2 <- freshKeyHash >>= \kh -> kh <$ registerStakeCredential (KeyHashObj kh)
+    deleg3 <- freshKeyHash >>= \kh -> kh <$ registerStakeCredential (KeyHashObj kh)
+    nes <- getsNES id
+    let sgs =
+          ShelleyGenesisStaking
+            { sgsPools = LM.ListMap [(pool1, poolParams), (pool2, poolParams), (pool3, poolParams)]
+            , sgsStake = LM.ListMap [(deleg1, pool1), (deleg2, pool1), (deleg3, pool2)]
+            }
+    let updatedNES = conwayRegisterInitialAccounts sgs nes
+    delegateStake (KeyHashObj deleg1) pool1
+    delegateStake (KeyHashObj deleg2) pool1
+    delegateStake (KeyHashObj deleg3) pool2
+    getPoolsState <$> (getsNES id) `shouldReturn` getPoolsState updatedNES
+    getDelegs deleg1 updatedNES `shouldReturn` Just pool1
+    getDelegs deleg2 updatedNES `shouldReturn` Just pool1
+    getDelegs deleg3 updatedNES `shouldReturn` Just pool2
   where
     expectDelegatedVote :: HasCallStack => Credential 'Staking -> DRep -> ImpTestM era ()
     expectDelegatedVote cred drep = do
@@ -561,6 +587,10 @@ spec = do
         assertBool
           ("Expected no drep state delegation to contain the stake credential: " <> show cred)
           (all (Set.notMember cred . drepDelegs) (Map.elems dreps))
+    getDelegs kh nes = do
+      let accounts = nes ^. nesEsL . esLStateL . lsCertStateL . certDStateL . accountsL . accountsMapL
+      pure $ Map.lookup (KeyHashObj kh) accounts >>= (^. stakePoolDelegationAccountStateL)
+    getPoolsState nes = nes ^. nesEsL . esLStateL . lsCertStateL . certPStateL . psStakePoolsL
 
 conwayEraSpecificSpec :: SpecWith (ImpInit (LedgerSpec ConwayEra))
 conwayEraSpecificSpec = do