Change PerasVoteAggregate{,Status} datatype to a GADT instead of a record

Author: tbagrel1

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Block/SupportsPeras.hs

@@ -12,6 +12,8 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
+{-# OPTIONS_GHC -Wno-partial-fields #-}
+{-# OPTIONS_GHC -Wno-unused-imports #-}
 
 module Ouroboros.Consensus.Block.SupportsPeras
   ( PerasRoundNo (..)
@@ -26,9 +28,11 @@ module Ouroboros.Consensus.Block.SupportsPeras
   , ValidatedPerasCert (..)
   , ValidatedPerasVote (..)
   , PerasVoteAggregate (..)
-  , emptyPerasVoteAggregate
+  , PerasVoteAggregateStatus (..)
+  , pvasMaybeCert
+  , emptyPerasVoteAggregateStatus
   , updatePerasVoteAggregate
-  , UpdatePerasVoteAggregateResult (..)
+  , updatePerasVoteAggregateStatus
   , makePerasCfg
   , HasId (..)
   , HasPerasCertRound (..)
@@ -51,6 +55,8 @@ import Codec.Serialise.Decoding (decodeListLenOf)
 import Codec.Serialise.Encoding (encodeListLen)
 import Control.Applicative ((<|>))
 import Data.Coerce (coerce)
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
 import Data.Maybe (isNothing)
 import Data.Monoid (Sum (..))
 import Data.Proxy (Proxy (..))
@@ -212,69 +218,89 @@ const_PERAS_QUORUM_THRESHOLD = PerasVoteStake 0.75
 
 data PerasVoteAggregate blk = PerasVoteAggregate
   { pvaTarget :: !(PerasVoteTarget blk)
-  , pvaVotes :: !(Set (WithArrivalTime (ValidatedPerasVote blk)))
+  , pvaVotes :: !(Map (IdOf (PerasVote blk)) (WithArrivalTime (ValidatedPerasVote blk)))
   , pvaTotalStake :: !PerasVoteStake
-  , pvaMaybeCert :: !(Maybe (ValidatedPerasCert blk))
   }
   deriving stock (Generic, Eq, Ord, Show)
   deriving anyclass NoThunks
 
-emptyPerasVoteAggregate :: PerasVoteTarget blk -> PerasVoteAggregate blk
-emptyPerasVoteAggregate target =
-  PerasVoteAggregate
-    { pvaTotalStake = PerasVoteStake 0
-    , pvaTarget = target
-    , pvaVotes = Set.empty
-    , pvaMaybeCert = Nothing
-    }
-
-data UpdatePerasVoteAggregateResult blk
-  = IncorrectPerasVoteTarget
-  | AddedPerasVoteButDidntGenerateNewCert (PerasVoteAggregate blk)
-  | AddedPerasVoteAndGeneratedNewCert (PerasVoteAggregate blk) (ValidatedPerasCert blk)
+data PerasVoteAggregateStatus blk
+  = PerasVoteAggregateQuorumNotReached {pvasVoteAggregate :: !(PerasVoteAggregate blk)}
+  | PerasVoteAggregateQuorumReachedAlready
+      {pvasVoteAggregate :: !(PerasVoteAggregate blk), pvasCert :: ValidatedPerasCert blk}
   deriving stock (Generic, Eq, Ord, Show)
   deriving anyclass NoThunks
 
+pvasMaybeCert :: PerasVoteAggregateStatus blk -> Maybe (ValidatedPerasCert blk)
+pvasMaybeCert aggr = case aggr of
+  PerasVoteAggregateQuorumNotReached{} -> Nothing
+  PerasVoteAggregateQuorumReachedAlready{pvasCert} -> Just pvasCert
+
+emptyPerasVoteAggregateStatus :: PerasVoteTarget blk -> PerasVoteAggregateStatus blk
+emptyPerasVoteAggregateStatus target =
+  PerasVoteAggregateQuorumNotReached $
+    PerasVoteAggregate
+      { pvaTotalStake = PerasVoteStake 0
+      , pvaTarget = target
+      , pvaVotes = Map.empty
+      }
+
+-- | Add a vote to an existing aggregate if it isn't already present, and update
+-- the stake accordingly.
+-- PRECONDITION: the vote's target must match the aggregate's target.
 updatePerasVoteAggregate ::
   StandardHash blk =>
   PerasVoteAggregate blk ->
   WithArrivalTime (ValidatedPerasVote blk) ->
-  UpdatePerasVoteAggregateResult blk
+  PerasVoteAggregate blk
 updatePerasVoteAggregate
   pva@PerasVoteAggregate
     { pvaTarget = (roundNo, point)
     , pvaVotes = existingVotes
     , pvaTotalStake = initialStake
-    , pvaMaybeCert = mExistingCert
     }
   vote =
-    if getPerasVoteRound vote == roundNo && getPerasVoteVotedBlock vote == point
-      then
-        let pvaTotalStake = initialStake + vpvVoteStake (forgetArrivalTime vote)
-            pvaVotes = Set.insert vote existingVotes
-            mNewCert =
-              if isNothing mExistingCert && pvaTotalStake >= const_PERAS_QUORUM_THRESHOLD
-                then
-                  Just $
-                    ValidatedPerasCert
-                      { vpcCertBoost = boostPerCert
-                      , vpcCert =
-                          PerasCert
-                            { pcCertRound = roundNo
-                            , pcCertBoostedBlock = point
-                            }
-                      }
-                else Nothing
-            pva' =
-              pva
-                { pvaVotes
-                , pvaTotalStake
-                , pvaMaybeCert = mExistingCert <|> mNewCert
-                }
-         in case mNewCert of
-              Just cert -> AddedPerasVoteAndGeneratedNewCert pva' cert
-              Nothing -> AddedPerasVoteButDidntGenerateNewCert pva'
-      else IncorrectPerasVoteTarget
+    if not (getPerasVoteRound vote == roundNo && getPerasVoteVotedBlock vote == point)
+      then error "updatePerasVoteAggregate: vote target does not match aggregate target"
+      else
+        let (pvaVotes', pvaTotalStake') =
+              case Map.insertLookupWithKey
+                (\_k old _new -> old)
+                (getId vote)
+                vote
+                existingVotes of
+                (Nothing, votes') ->
+                  -- key was NOT present → inserted and stake updated
+                  (votes', initialStake + vpvVoteStake (forgetArrivalTime vote))
+                (Just _, _) ->
+                  -- key WAS already present → votes and stake unchanged
+                  (existingVotes, initialStake)
+         in pva{pvaVotes = pvaVotes', pvaTotalStake = pvaTotalStake'}
+
+updatePerasVoteAggregateStatus ::
+  StandardHash blk =>
+  PerasVoteAggregateStatus blk ->
+  WithArrivalTime (ValidatedPerasVote blk) ->
+  PerasVoteAggregateStatus blk
+updatePerasVoteAggregateStatus aggr vote = case aggr of
+  PerasVoteAggregateQuorumNotReached{pvasVoteAggregate} ->
+    let aggr' = updatePerasVoteAggregate pvasVoteAggregate vote
+     in if pvaTotalStake aggr' >= const_PERAS_QUORUM_THRESHOLD
+          then
+            PerasVoteAggregateQuorumReachedAlready
+              { pvasVoteAggregate = aggr'
+              , pvasCert =
+                  ValidatedPerasCert
+                    { vpcCertBoost = boostPerCert
+                    , vpcCert = uncurry PerasCert (pvaTarget aggr')
+                    }
+              }
+          else PerasVoteAggregateQuorumNotReached{pvasVoteAggregate = aggr'}
+  PerasVoteAggregateQuorumReachedAlready{pvasVoteAggregate, pvasCert} ->
+    PerasVoteAggregateQuorumReachedAlready
+      { pvasVoteAggregate = updatePerasVoteAggregate pvasVoteAggregate vote
+      , pvasCert
+      }
 
 type PerasVoteTarget blk = (PerasRoundNo, Point blk)
 

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/PerasVoteDB/API.hs

@@ -56,7 +56,7 @@ data AddPerasVoteResult blk
   deriving anyclass NoThunks
 
 newtype PerasStakeSnapshot blk = PerasStakeSnapshot
-  {unPerasStakeSnapshot :: WithFingerprint (Map (PerasVoteTarget blk) (PerasVoteAggregate blk))}
+  {unPerasStakeSnapshot :: WithFingerprint (Map (PerasVoteTarget blk) (PerasVoteAggregateStatus blk))}
   deriving Generic
   deriving newtype NoThunks
 
@@ -65,7 +65,7 @@ getPerasCertsFromStakeSnapshot ::
   PerasStakeSnapshot blk ->
   Set (ValidatedPerasCert blk)
 getPerasCertsFromStakeSnapshot (PerasStakeSnapshot mp) =
-  Set.fromList $ Map.elems $ Map.mapMaybe pvaMaybeCert (forgetFingerprint mp)
+  Set.fromList $ Map.elems $ Map.mapMaybe pvasMaybeCert (forgetFingerprint mp)
 
 data PerasVoteSnapshot blk = PerasVoteSnapshot
   { containsVote :: IdOf (PerasVote blk) -> Bool

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

@@ -36,7 +36,7 @@ import Data.Set (Set)
 import Data.Set qualified as Set
 import GHC.Generics (Generic)
 import NoThunks.Class
-import Ouroboros.Consensus.Block hiding (UpdatePerasVoteAggregateResult (..))
+import Ouroboros.Consensus.Block
 import Ouroboros.Consensus.Block.SupportsPeras qualified as UPVAR
 import Ouroboros.Consensus.BlockchainTime (WithArrivalTime (forgetArrivalTime))
 import Ouroboros.Consensus.Peras.Weight
@@ -182,11 +182,14 @@ implAddVote env vote = do
             (res, pvsVotesByTarget') =
               Map.alterF
                 ( \mExistingAggr ->
-                    let aggr = fromMaybe (emptyPerasVoteAggregate voteTarget) mExistingAggr
-                     in case updatePerasVoteAggregate aggr vote of
-                          UPVAR.IncorrectPerasVoteTarget -> error "The aggregate should match the vote target."
-                          UPVAR.AddedPerasVoteButDidntGenerateNewCert aggr' -> (AddedPerasVoteButDidntGenerateNewCert, Just aggr')
-                          UPVAR.AddedPerasVoteAndGeneratedNewCert aggr' cert -> (AddedPerasVoteAndGeneratedNewCert cert, Just aggr')
+                    let aggr = fromMaybe (emptyPerasVoteAggregateStatus voteTarget) mExistingAggr
+                        aggr' = updatePerasVoteAggregateStatus aggr vote
+                     in case (aggr, aggr') of
+                          -- if we observe a state transition, it means a certificate was emitted
+                          (PerasVoteAggregateQuorumNotReached{}, PerasVoteAggregateQuorumReachedAlready{pvasCert}) ->
+                            (AddedPerasVoteAndGeneratedNewCert pvasCert, Just aggr')
+                          _ ->
+                            (AddedPerasVoteButDidntGenerateNewCert, Just aggr')
                 )
                 voteTarget
                 pvsVotesByTarget
@@ -281,7 +284,7 @@ implGarbageCollect PerasVoteDbEnv{pvdbPerasVoteStateVar} roundNo =
 
 data PerasVoteState blk = PerasVoteState
   { pvsVoteIds :: !(Set (IdOf (PerasVote blk)))
-  , pvsVotesByTarget :: !(Map (PerasVoteTarget blk) (PerasVoteAggregate blk))
+  , pvsVotesByTarget :: !(Map (PerasVoteTarget blk) (PerasVoteAggregateStatus blk))
   , pvsVotesByTicket :: !(Map PerasVoteTicketNo (WithArrivalTime (ValidatedPerasVote blk)))
   -- ^ The votes by 'PerasVoteTicketNo'.
   --