Use a Data.OMap.Strict to replace ProposalsSnapshot (#3791)

* Add Data.OMap.Strict to cardano-data

* Make ProposalsSnapshot an OMap

eras/conway/impl/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 1.10.1.0
 
+* Switch to using `OMap` for `ProposalsSnapshot` #3791
 * Add `VotingOnExpiredGovAction` predicate failure in `GOV` #3825
 
 ### `testlib`

eras/conway/impl/cardano-ledger-conway.cabal

@@ -75,7 +75,7 @@ library
         aeson >=2.2,
         data-default-class,
         cardano-crypto-class,
-        cardano-data >=1.1.1.0,
+        cardano-data >=1.1.2.0,
         cardano-ledger-binary >=1.2,
         cardano-ledger-allegra >=1.1,
         cardano-ledger-alonzo ^>=1.5,

eras/conway/impl/src/Cardano/Ledger/Conway/Governance/Procedures.hs

@@ -3,9 +3,11 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards #-}
@@ -14,6 +16,7 @@
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
 
 module Cardano.Ledger.Conway.Governance.Procedures (
@@ -105,6 +108,7 @@ import Data.Aeson.Types (toJSONKeyText)
 import Data.Default.Class
 import Data.Map.Strict (Map)
 import Data.Maybe.Strict (StrictMaybe (..))
+import qualified Data.OMap.Strict as OMap
 import qualified Data.OSet.Strict as OSet
 import qualified Data.Sequence as Seq
 import Data.Set (Set)
@@ -264,6 +268,13 @@ instance EraPParams era => EncCBOR (GovActionState era) where
         !> To gasProposedIn
         !> To gasExpiresAfter
 
+-- Ref: https://gitlab.haskell.org/ghc/ghc/-/issues/14046
+instance
+  c ~ EraCrypto era =>
+  OMap.HasOKey (GovActionId c) (GovActionState era)
+  where
+  okeyL = lens gasId $ \gas gi -> gas {gasId = gi}
+
 data Voter c
   = CommitteeVoter !(Credential 'HotCommitteeRole c)
   | DRepVoter !(Credential 'DRepRole c)

eras/conway/impl/src/Cardano/Ledger/Conway/Governance/Snapshots.hs

@@ -1,7 +1,10 @@
 {-# LANGUAGE DeriveGeneric #-}
-{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE UndecidableInstances #-}
 
 module Cardano.Ledger.Conway.Governance.Snapshots (
   ProposalsSnapshot,
@@ -32,25 +35,20 @@ import Cardano.Ledger.TreeDiff (ToExpr)
 import Control.DeepSeq (NFData)
 import Data.Aeson (ToJSON)
 import Data.Default.Class (Default (..))
-import Data.Foldable (Foldable (..))
-import Data.List (sort)
 import qualified Data.Map as Map
 import Data.Map.Strict (Map)
-import Data.MapExtras (extractKeys)
-import Data.Maybe (fromMaybe)
+import qualified Data.OMap.Strict as OMap
 import Data.Sequence.Strict (StrictSeq (..))
 import Data.Set (Set)
-import qualified Data.Set as Set
 import GHC.Generics (Generic)
 import Lens.Micro (Lens', (%~))
 import NoThunks.Class (NoThunks)
 
-data ProposalsSnapshot era = ProposalsSnapshot
-  { psGovActionStates :: !(Map (GovActionId (EraCrypto era)) (GovActionState era))
-  , psProposalOrder :: !(StrictSeq (GovActionId (EraCrypto era)))
-  -- ^ Newer actions are near the end
-  }
-  deriving (Generic, Eq, Show)
+newtype ProposalsSnapshot era
+  = ProposalsSnapshot
+      (OMap.OMap (GovActionId (EraCrypto era)) (GovActionState era))
+  deriving newtype (Show, Eq)
+  deriving stock (Generic)
 
 instance EraPParams era => ToExpr (ProposalsSnapshot era)
 
@@ -61,7 +59,7 @@ instance EraPParams era => NFData (ProposalsSnapshot era)
 instance EraPParams era => NoThunks (ProposalsSnapshot era)
 
 instance Default (ProposalsSnapshot era) where
-  def = ProposalsSnapshot def def
+  def = ProposalsSnapshot def
 
 instance EraPParams era => EncCBOR (ProposalsSnapshot era) where
   encCBOR = encCBOR . snapshotActions
@@ -73,49 +71,43 @@ instance EraPParams era => DecCBOR (ProposalsSnapshot era) where
 instance EraPParams era => DecShareCBOR (ProposalsSnapshot era) where
   decShareCBOR _ = fromGovActionStateSeq <$> decCBOR
 
+-- | Insert a `GovActionState`, overwriting an entry of it if the
+-- corresponding `GovActionId` already exists.
 snapshotInsertGovAction ::
   GovActionState era ->
   ProposalsSnapshot era ->
   ProposalsSnapshot era
-snapshotInsertGovAction gas@GovActionState {gasId} ps@ProposalsSnapshot {..}
-  | Map.member gasId psGovActionStates =
-      ps {psGovActionStates = Map.insert gasId gas psGovActionStates}
-  | otherwise =
-      ProposalsSnapshot
-        { psGovActionStates = Map.insert gasId gas psGovActionStates
-        , psProposalOrder = psProposalOrder :|> gasId
-        }
+snapshotInsertGovAction gas (ProposalsSnapshot omap) =
+  ProposalsSnapshot (omap OMap.||> gas)
 
+-- | Get the sequence of `GovActionState`s
 snapshotActions ::
   ProposalsSnapshot era ->
   StrictSeq (GovActionState era)
-snapshotActions ProposalsSnapshot {..} = toGovAction <$> psProposalOrder
-  where
-    toGovAction gaId =
-      fromMaybe
-        (error $ "Impossible: ProposalsSnapshot invariant is not maintained: " <> show gaId)
-        (Map.lookup gaId psGovActionStates)
+snapshotActions (ProposalsSnapshot omap) = OMap.toStrictSeq omap
 
+-- | Get the sequence of `GovActionId`s
 snapshotIds ::
   ProposalsSnapshot era ->
   StrictSeq (GovActionId (EraCrypto era))
-snapshotIds = psProposalOrder
+snapshotIds (ProposalsSnapshot omap) = OMap.toStrictSeqOKeys omap
 
+-- | Get the unordered map of `GovActionId`s and `GovActionState`s
 snapshotGovActionStates ::
   ProposalsSnapshot era ->
   Map (GovActionId (EraCrypto era)) (GovActionState era)
-snapshotGovActionStates = psGovActionStates
+snapshotGovActionStates (ProposalsSnapshot omap) = OMap.toMap omap
 
+-- | Add a vote to an existing `GovActionState` This is a no-op if the .
+-- provided `GovActionId` does not already exist                       .
 snapshotAddVote ::
   Voter (EraCrypto era) ->
   Vote ->
   GovActionId (EraCrypto era) ->
   ProposalsSnapshot era ->
   ProposalsSnapshot era
-snapshotAddVote voter vote gId ps@ProposalsSnapshot {..} =
-  ps
-    { psGovActionStates = Map.update (Just . updateVote) gId psGovActionStates
-    }
+snapshotAddVote voter vote gai (ProposalsSnapshot omap) =
+  ProposalsSnapshot $ OMap.adjust updateVote gai omap
   where
     insertVote ::
       Ord k =>
@@ -129,39 +121,29 @@ snapshotAddVote voter vote gId ps@ProposalsSnapshot {..} =
       StakePoolVoter kh -> insertVote gasStakePoolVotesL kh
       CommitteeVoter c -> insertVote gasCommitteeVotesL c
 
+-- | Extract `GovActionState`s for the given set of `GovActionId`s from the `Proposals`
 snapshotRemoveIds ::
   Set (GovActionId (EraCrypto era)) ->
   ProposalsSnapshot era ->
   (ProposalsSnapshot era, Map.Map (GovActionId (EraCrypto era)) (GovActionState era))
-snapshotRemoveIds gIds (ProposalsSnapshot {..}) = (retainedProposals, removedGovActionStates)
-  where
-    (retainedGovActionStates, removedGovActionStates) = psGovActionStates `extractKeys` gIds
-    retainedProposals =
-      ProposalsSnapshot
-        { psGovActionStates = retainedGovActionStates
-        , psProposalOrder =
-            foldl' (\s x -> if x `Set.member` gIds then s else x :<| s) mempty psProposalOrder
-        }
+snapshotRemoveIds gais (ProposalsSnapshot omap) =
+  let (retained, removed) = OMap.extractKeys gais omap
+   in (ProposalsSnapshot retained, removed)
 
 snapshotLookupId ::
   GovActionId (EraCrypto era) ->
   ProposalsSnapshot era ->
   Maybe (GovActionState era)
-snapshotLookupId gId ProposalsSnapshot {psGovActionStates} =
-  Map.lookup gId psGovActionStates
+snapshotLookupId gai (ProposalsSnapshot omap) = OMap.lookup gai omap
 
 -- | Converts a sequence of `GovActionState`s to a `ProposalsSnapshot`.
 --
 -- /Warning/ - This function expects `GovActionState`'s to have unique
 -- `GovActionId`s, because duplicate Ids will result in `GovActionStates`
 -- to be dropped.
-fromGovActionStateSeq ::
-  StrictSeq (GovActionState era) ->
-  ProposalsSnapshot era
-fromGovActionStateSeq = foldl' (flip snapshotInsertGovAction) def
+fromGovActionStateSeq :: StrictSeq (GovActionState era) -> ProposalsSnapshot era
+fromGovActionStateSeq = ProposalsSnapshot . OMap.fromFoldable
 
 -- | Internal function for checking if the invariants are maintained
 isConsistent_ :: ProposalsSnapshot era -> Bool
-isConsistent_ (ProposalsSnapshot {psGovActionStates, psProposalOrder}) =
-  Map.keys psGovActionStates == sort (toList psProposalOrder)
-    && all (\(k, GovActionState {gasId}) -> k == gasId) (Map.toList psGovActionStates)
+isConsistent_ (ProposalsSnapshot omap) = OMap.invariantHolds' omap

libs/cardano-data/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Version history for `cardano-data`
 
+## 1.1.2.0
+
+- Add Data.OMap.Strict #3791
+
 ## 1.1.1.0
 
 - Add Data.OSet.Strict #3779

libs/cardano-data/cardano-data.cabal

@@ -1,6 +1,6 @@
 cabal-version:      3.0
 name:               cardano-data
-version:            1.1.1.0
+version:            1.1.2.0
 license:            Apache-2.0
 maintainer:         operations@iohk.io
 author:             IOHK
@@ -23,6 +23,7 @@ library
         Data.ListMap
         Data.Universe
         Data.OSet.Strict
+        Data.OMap.Strict
 
     hs-source-dirs:   src
     default-language: Haskell2010
@@ -36,10 +37,12 @@ library
         cardano-ledger-binary >=1.2,
         cardano-strict-containers >=0.1.2.1,
         containers,
+        data-default-class,
         deepseq,
         mtl,
         nothunks,
-        vector
+        vector,
+        microlens
 
 library testlib
     exposed-modules:
@@ -55,11 +58,14 @@ library testlib
 
     build-depends:
         base,
+        cardano-data,
+        cardano-ledger-binary:testlib,
         containers,
         hspec,
         QuickCheck,
         cardano-ledger-binary:testlib,
-        cardano-data
+        cardano-data,
+        microlens
 
 test-suite cardano-data-tests
     type:             exitcode-stdio-1.0
@@ -68,6 +74,7 @@ test-suite cardano-data-tests
     other-modules:
         Test.Cardano.Data.MapExtrasSpec
         Test.Cardano.Data.OSet.StrictSpec
+        Test.Cardano.Data.OMap.StrictSpec
 
     default-language: Haskell2010
     ghc-options:
@@ -77,11 +84,12 @@ test-suite cardano-data-tests
 
     build-depends:
         base,
+        cardano-strict-containers,
         containers,
         hspec,
         cardano-data,
         cardano-ledger-binary:testlib,
-        cardano-strict-containers,
         testlib,
         QuickCheck,
-        quickcheck-classes-base
+        quickcheck-classes-base,
+        microlens