feature(cardano-db-sync): Create a new HasNewEpochState class

This class can be used to get/update the underlying extended
ledger state

Author: sgillespie

cardano-db-sync/src/Cardano/DbSync/Ledger/Types.hs

@@ -5,6 +5,7 @@
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE NoImplicitPrelude #-}
 
 module Cardano.DbSync.Ledger.Types where
@@ -28,6 +29,7 @@ import Cardano.Ledger.Coin (Coin)
 import Cardano.Ledger.Conway.Governance
 import Cardano.Ledger.Credential (Credential (..))
 import Cardano.Ledger.Keys (KeyRole (..))
+import Cardano.Ledger.Shelley.LedgerState (NewEpochState ())
 import Cardano.Prelude hiding (atomically)
 import Cardano.Slotting.Slot (
   EpochNo (..),
@@ -39,13 +41,17 @@ import Control.Concurrent.Class.MonadSTM.Strict (
  )
 import Control.Concurrent.STM.TBQueue (TBQueue)
 import qualified Data.Map.Strict as Map
+import Data.SOP.Strict
 import qualified Data.Set as Set
 import qualified Data.Strict.Maybe as Strict
+import Lens.Micro (Traversal')
 import Ouroboros.Consensus.BlockchainTime.WallClock.Types (SystemStart (..))
-import Ouroboros.Consensus.Cardano.Block (StandardConway, StandardCrypto)
+import Ouroboros.Consensus.Cardano.Block hiding (CardanoBlock, CardanoLedgerState)
+import Ouroboros.Consensus.HardFork.Combinator.Basics (LedgerState (..))
 import Ouroboros.Consensus.Ledger.Abstract (getTipSlot)
 import Ouroboros.Consensus.Ledger.Extended (ExtLedgerState (..))
 import qualified Ouroboros.Consensus.Node.ProtocolInfo as Consensus
+import Ouroboros.Consensus.Shelley.Ledger (LedgerState (..), ShelleyBlock)
 import Ouroboros.Network.AnchoredSeq (Anchorable (..), AnchoredSeq (..))
 import Prelude (fail, id)
 
@@ -190,3 +196,61 @@ instance Anchorable (WithOrigin SlotNo) CardanoLedgerState CardanoLedgerState wh
   getAnchorMeasure _ = getTipSlot . clsState
 
 data SnapshotPoint = OnDisk LedgerStateFile | InMemory CardanoPoint
+
+-- | Per-era pure getters and setters on @NewEpochState@. Note this is a bit of an abuse
+-- of the cardano-ledger/ouroboros-consensus public APIs, because ledger state is not
+-- designed to be updated this way. We are only replaying the chain, so this should be
+-- safe.
+class HasNewEpochState era where
+  getNewEpochState :: ExtLedgerState CardanoBlock -> Maybe (NewEpochState era)
+
+  applyNewEpochState ::
+    NewEpochState era ->
+    ExtLedgerState CardanoBlock ->
+    ExtLedgerState CardanoBlock
+
+instance HasNewEpochState StandardShelley where
+  getNewEpochState st = case ledgerState st of
+    LedgerStateShelley shelley -> Just (shelleyLedgerState shelley)
+    _ -> Nothing
+
+  applyNewEpochState st =
+    hApplyExtLedgerState $
+      fn (applyNewEpochState' st) :* fn id :* fn id :* fn id :* fn id :* fn id :* Nil
+
+instance HasNewEpochState StandardConway where
+  getNewEpochState st = case ledgerState st of
+    LedgerStateConway shelley -> Just (shelleyLedgerState shelley)
+    _ -> Nothing
+
+  applyNewEpochState st =
+    hApplyExtLedgerState $
+      fn id :* fn id :* fn id :* fn id :* fn id :* fn (applyNewEpochState' st) :* Nil
+
+hApplyExtLedgerState ::
+  NP (LedgerState -.-> LedgerState) (CardanoShelleyEras StandardCrypto) ->
+  ExtLedgerState CardanoBlock ->
+  ExtLedgerState CardanoBlock
+hApplyExtLedgerState f ledger =
+  case ledgerState ledger of
+    HardForkLedgerState hfState ->
+      let newHfState = hap (fn id :* f) hfState
+       in updateLedgerState $ HardForkLedgerState newHfState
+  where
+    updateLedgerState st = ledger {ledgerState = st}
+
+applyNewEpochState' ::
+  NewEpochState era ->
+  LedgerState (ShelleyBlock proto era) ->
+  LedgerState (ShelleyBlock proto era)
+applyNewEpochState' newEpochState' ledger =
+  ledger {shelleyLedgerState = newEpochState'}
+
+-- | A @Traversal@ that targets the @NewEpochState@ from the extended ledger state
+newEpochStateT ::
+  HasNewEpochState era =>
+  Traversal' (ExtLedgerState CardanoBlock) (NewEpochState era)
+newEpochStateT f ledger =
+  case getNewEpochState ledger of
+    Just newEpochState' -> flip applyNewEpochState ledger <$> f newEpochState'
+    Nothing -> pure ledger

cardano-db-sync/src/Cardano/DbSync/Types.hs

@@ -1,6 +1,8 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module Cardano.DbSync.Types (
   BlockDetails (..),
@@ -53,6 +55,7 @@ import qualified Cardano.Ledger.Credential as Ledger
 import Cardano.Ledger.Crypto (StandardCrypto)
 import qualified Cardano.Ledger.Hashes as Ledger
 import Cardano.Ledger.Keys
+
 import Cardano.Prelude hiding (Meta, show)
 import Cardano.Slotting.Slot (EpochNo (..), EpochSize (..), SlotNo (..))
 import qualified Data.Text as Text