Make PerasWeightSnapshot opaque

Author: amesgen

ouroboros-consensus/bench/PerasCertDB-bench/Main.hs

@@ -12,12 +12,12 @@
 module Main (main) where
 
 import Data.List (iterate')
-import Data.Map.Strict qualified as Map
 import Numeric.Natural (Natural)
 import Ouroboros.Consensus.Block (PerasWeight (PerasWeight), SlotNo (..))
 import Ouroboros.Consensus.Peras.Weight
-  ( PerasWeightSnapshot (..)
+  ( PerasWeightSnapshot
   , boostedWeightForFragment
+  , mkPerasWeightSnapshot
   )
 import Ouroboros.Network.AnchoredFragment qualified as AF
 import Test.Ouroboros.Storage.TestBlock (TestBlock (..), TestBody (..), TestHeader (..))
@@ -96,7 +96,7 @@ uniformWeightSnapshot fragment =
           . AF.toOldestFirst
           $ fragment
       weights = repeat (boostWeight benchmarkParams)
-   in PerasWeightSnapshot{getPerasWeightSnapshot = Map.fromList $ zip pointsToBoost weights}
+   in mkPerasWeightSnapshot $ pointsToBoost `zip` weights
 
 getEveryN :: Natural -> [(Natural, a)] -> [(Natural, a)]
 getEveryN n = filter (\(i, _) -> (i `mod` n) == 0)

ouroboros-consensus/ouroboros-consensus.cabal

@@ -846,7 +846,6 @@ benchmark PerasCertDB-bench
   other-modules:
   build-depends:
     base,
-    containers,
     ouroboros-consensus,
     ouroboros-network-api,
     tasty-bench,

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Peras/Weight.hs

@@ -6,11 +6,18 @@
 {-# LANGUAGE TypeOperators #-}
 
 module Ouroboros.Consensus.Peras.Weight
-  ( PerasWeightSnapshot (..)
+  ( -- * 'PerasWeightSnapshot'
+    PerasWeightSnapshot
+  , emptyPerasWeightSnapshot
+  , mkPerasWeightSnapshot
+  , perasWeightSnapshotToList
+  , addToPerasWeightSnapshot
+  , removeFromPerasWeightSnapshot
   , boostedWeightForPoint
   , boostedWeightForFragment
   ) where
 
+import Data.Foldable as Foldable (foldl')
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import GHC.Generics (Generic)
@@ -26,6 +33,44 @@ newtype PerasWeightSnapshot blk = PerasWeightSnapshot
   deriving Generic
   deriving newtype NoThunks
 
+emptyPerasWeightSnapshot :: PerasWeightSnapshot blk
+emptyPerasWeightSnapshot = PerasWeightSnapshot Map.empty
+
+mkPerasWeightSnapshot ::
+  StandardHash blk =>
+  [(Point blk, PerasWeight)] ->
+  PerasWeightSnapshot blk
+mkPerasWeightSnapshot =
+  Foldable.foldl'
+    (\s (pt, weight) -> addToPerasWeightSnapshot pt weight s)
+    emptyPerasWeightSnapshot
+
+perasWeightSnapshotToList :: PerasWeightSnapshot blk -> [(Point blk, PerasWeight)]
+perasWeightSnapshotToList = Map.toList . getPerasWeightSnapshot
+
+addToPerasWeightSnapshot ::
+  StandardHash blk =>
+  Point blk ->
+  PerasWeight ->
+  PerasWeightSnapshot blk ->
+  PerasWeightSnapshot blk
+addToPerasWeightSnapshot pt weight =
+  PerasWeightSnapshot . Map.insertWith (<>) pt weight . getPerasWeightSnapshot
+
+removeFromPerasWeightSnapshot ::
+  StandardHash blk =>
+  Point blk ->
+  PerasWeight ->
+  PerasWeightSnapshot blk ->
+  PerasWeightSnapshot blk
+removeFromPerasWeightSnapshot pt (PerasWeight weight) =
+  PerasWeightSnapshot . Map.update subtractWeight pt . getPerasWeightSnapshot
+ where
+  subtractWeight :: PerasWeight -> Maybe PerasWeight
+  subtractWeight (PerasWeight w)
+    | w > weight = Just $ PerasWeight (w - weight)
+    | otherwise = Nothing
+
 boostedWeightForPoint ::
   forall blk.
   StandardHash blk =>

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/PerasCertDB/Impl.hs

@@ -21,8 +21,8 @@ module Ouroboros.Consensus.Storage.PerasCertDB.Impl
   ) where
 
 import Control.Tracer (Tracer, nullTracer, traceWith)
+import Data.Foldable as Foldable (foldl')
 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 GHC.Generics (Generic)
@@ -165,7 +165,7 @@ implAddCert env cert = do
                   Map.insert roundNo cert pvcsCerts
               , -- Note that the same block might be boosted by multiple points.
                 pvcsWeightByPoint =
-                  Map.insertWith (<>) boostedPt boostPerCert pvcsWeightByPoint
+                  addToPerasWeightSnapshot boostedPt boostPerCert pvcsWeightByPoint
               }
             (succ fp)
         pure AddedPerasCertToDB
@@ -186,7 +186,7 @@ implGetWeightSnapshot ::
   IOLike m =>
   PerasCertDbEnv m blk -> STM m (WithFingerprint (PerasWeightSnapshot blk))
 implGetWeightSnapshot PerasCertDbEnv{pcdbVolatileState} =
-  fmap (PerasWeightSnapshot . pvcsWeightByPoint) <$> readTVar pcdbVolatileState
+  fmap pvcsWeightByPoint <$> readTVar pcdbVolatileState
 
 implGarbageCollect ::
   forall m blk.
@@ -202,33 +202,21 @@ implGarbageCollect PerasCertDbEnv{pcdbVolatileState} slot =
     PerasVolatileCertState
       { pvcsCerts = certsToKeep
       , pvcsWeightByPoint =
-          Map.merge
-            -- Do not touch weight of boosted blocks that we do not subtract any
-            -- weight from.
-            Map.preserveMissing
-            -- Irrelevant, the key set of @weightToRemove@ is a subset of the
-            -- key set of @pvcsWeightByPoint@.
-            Map.dropMissing
-            (Map.zipWithMaybeMatched $ \_pt -> subtractWeight)
+          Foldable.foldl'
+            ( \s cert ->
+                removeFromPerasWeightSnapshot
+                  (perasCertBoostedBlock cert)
+                  boostPerCert
+                  s
+            )
             pvcsWeightByPoint
-            weightToRemove
+            certsToRemove
       }
    where
     (certsToRemove, certsToKeep) =
       Map.partition isTooOld pvcsCerts
     isTooOld cert =
       pointSlot (perasCertBoostedBlock cert) < NotOrigin slot
-    weightToRemove =
-      Map.fromListWith
-        (<>)
-        [ (perasCertBoostedBlock cert, boostPerCert)
-        | cert <- Map.elems certsToRemove
-        ]
-
-  subtractWeight :: PerasWeight -> PerasWeight -> Maybe PerasWeight
-  subtractWeight (PerasWeight w1) (PerasWeight w2)
-    | w1 > w2 = Just $ PerasWeight (w1 - w2)
-    | otherwise = Nothing
 
 {-------------------------------------------------------------------------------
   Implementation-internal types
@@ -239,7 +227,7 @@ implGarbageCollect PerasCertDbEnv{pcdbVolatileState} slot =
 data PerasVolatileCertState blk = PerasVolatileCertState
   { pvcsCerts :: !(Map PerasRoundNo (PerasCert blk))
   -- ^ The boosted blocks by 'RoundNo' of all certificates currently in the db.
-  , pvcsWeightByPoint :: !(Map (Point blk) PerasWeight)
+  , pvcsWeightByPoint :: !(PerasWeightSnapshot blk)
   -- ^ The weight of boosted blocks w.r.t. the certificates currently in the db.
   --
   -- INVARIANT: In sync with 'pvcsCerts'.
@@ -252,7 +240,7 @@ initialPerasVolatileCertState =
   WithFingerprint
     PerasVolatileCertState
       { pvcsCerts = Map.empty
-      , pvcsWeightByPoint = Map.empty
+      , pvcsWeightByPoint = emptyPerasWeightSnapshot
       }
     (Fingerprint 0)
 

ouroboros-consensus/test/storage-test/Test/Ouroboros/Storage/PerasCertDB/Model.hs

@@ -14,12 +14,14 @@ module Test.Ouroboros.Storage.PerasCertDB.Model
   , garbageCollect
   ) where
 
-import qualified Data.Map as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
 import GHC.Generics (Generic)
 import Ouroboros.Consensus.Block
-import Ouroboros.Consensus.Peras.Weight (PerasWeightSnapshot (..))
+import Ouroboros.Consensus.Peras.Weight
+  ( PerasWeightSnapshot
+  , mkPerasWeightSnapshot
+  )
 
 data Model blk = Model
   { certs :: Set (PerasCert blk)
@@ -47,16 +49,9 @@ addCert model@Model{certs} cert =
 getWeightSnapshot ::
   StandardHash blk =>
   Model blk -> PerasWeightSnapshot blk
-getWeightSnapshot Model{certs} = snap
- where
-  snap =
-    PerasWeightSnapshot
-      { getPerasWeightSnapshot =
-          Set.fold
-            (\cert acc -> Map.insertWith (<>) (perasCertBoostedBlock cert) boostPerCert acc)
-            Map.empty
-            certs
-      }
+getWeightSnapshot Model{certs} =
+  mkPerasWeightSnapshot
+    [(perasCertBoostedBlock cert, boostPerCert) | cert <- Set.toList certs]
 
 garbageCollect :: StandardHash blk => SlotNo -> Model blk -> Model blk
 garbageCollect slot model@Model{certs} =