O.C.Peras.Weight: add totalWeightForFragment/takeVolatileSuffix

ouroboros-consensus/ouroboros-consensus.cabal

@@ -598,6 +598,7 @@ test-suite consensus-test
     Test.Consensus.MiniProtocol.ChainSync.CSJ
     Test.Consensus.MiniProtocol.ChainSync.Client
     Test.Consensus.MiniProtocol.LocalStateQuery.Server
+    Test.Consensus.Peras.WeightSnapshot
     Test.Consensus.Util.MonadSTM.NormalForm
     Test.Consensus.Util.Versioned
 

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

@@ -21,19 +21,26 @@ import Data.Word (Word64)
 import GHC.Generics (Generic)
 import NoThunks.Class
 import Ouroboros.Consensus.Block.Abstract
+import Ouroboros.Consensus.Util.Condense
 import Quiet (Quiet (..))
 
 newtype PerasRoundNo = PerasRoundNo {unPerasRoundNo :: Word64}
   deriving Show via Quiet PerasRoundNo
   deriving stock Generic
   deriving newtype (Eq, Ord, NoThunks)
 
+instance Condense PerasRoundNo where
+  condense = show . unPerasRoundNo
+
 newtype PerasWeight = PerasWeight {unPerasWeight :: Word64}
   deriving Show via Quiet PerasWeight
   deriving stock Generic
   deriving newtype (Eq, Ord, NoThunks)
   deriving (Semigroup, Monoid) via Sum Word64
 
+instance Condense PerasWeight where
+  condense = show . unPerasWeight
+
 -- | TODO this will become a Ledger protocol parameter
 boostPerCert :: PerasWeight
 boostPerCert = PerasWeight 15

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Config/SecurityParam.hs

@@ -3,24 +3,39 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
-module Ouroboros.Consensus.Config.SecurityParam (SecurityParam (..)) where
+module Ouroboros.Consensus.Config.SecurityParam
+  ( SecurityParam (..)
+  , maxRollbackWeight
+  ) where
 
 import Cardano.Binary
 import Cardano.Ledger.BaseTypes.NonZero
 import Data.Word
 import GHC.Generics (Generic)
 import NoThunks.Class (NoThunks)
+import Ouroboros.Consensus.Block.SupportsPeras (PerasWeight (..))
 import Quiet
 
 -- | Protocol security parameter
 --
--- We interpret this as the number of rollbacks we support.
+-- In longest-chain protocols, we interpret this as the number of rollbacks we
+-- support.
 --
 -- i.e., k == 1: we can roll back at most one block
 --       k == 2: we can roll back at most two blocks, etc
 --
 -- NOTE: This talks about the number of /blocks/ we can roll back, not
 -- the number of /slots/.
+--
+-- In weightiest-chain protocols (Ouroboros Peras), we interpret this as the
+-- maximum amount of weight we can roll back.
+--
+-- i.e. k == 30: we can roll back at most 30 unweighted blocks, or two blocks
+-- each having additional weight 14.
 newtype SecurityParam = SecurityParam {maxRollbacks :: NonZero Word64}
   deriving (Eq, Generic, NoThunks, ToCBOR, FromCBOR)
   deriving Show via Quiet SecurityParam
+
+-- | The maximum amount of weight we can roll back.
+maxRollbackWeight :: SecurityParam -> PerasWeight
+maxRollbackWeight = PerasWeight . unNonZero . maxRollbacks

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

@@ -26,14 +26,18 @@ module Ouroboros.Consensus.Peras.Weight
     -- * Query
   , weightBoostOfPoint
   , weightBoostOfFragment
+  , totalWeightOfFragment
+  , takeVolatileSuffix
   ) where
 
 import Data.Foldable as Foldable (foldl')
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
+import Data.Word (Word64)
 import GHC.Generics (Generic)
 import NoThunks.Class
 import Ouroboros.Consensus.Block
+import Ouroboros.Consensus.Config.SecurityParam
 import Ouroboros.Network.AnchoredFragment (AnchoredFragment)
 import qualified Ouroboros.Network.AnchoredFragment as AF
 
@@ -236,11 +240,140 @@ weightBoostOfFragment weightSnap frag =
     (weightBoostOfPoint weightSnap . castPoint . blockPoint)
     (AF.toOldestFirst frag)
 
+-- | Get the total weight for a fragment, ie the length plus the weight boost
+-- ('weightBoostOfFragment') of the fragment.
+--
+-- Note that this quantity is relative to the anchor of the fragment, so it
+-- should only be compared against other fragments with the same anchor.
+--
+-- >>> :{
+-- weights :: [(Point Blk, PerasWeight)]
+-- weights =
+--   [ (BlockPoint 2 "foo", PerasWeight 2)
+--   , (GenesisPoint,       PerasWeight 3)
+--   , (BlockPoint 3 "bar", PerasWeight 2)
+--   , (BlockPoint 2 "foo", PerasWeight 2)
+--   ]
+-- :}
+--
+-- >>> :{
+-- snap = mkPerasWeightSnapshot weights
+-- foo = HeaderFields (SlotNo 2) (BlockNo 1) "foo"
+-- bar = HeaderFields (SlotNo 3) (BlockNo 2) "bar"
+-- frag0 :: AnchoredFragment (HeaderFields Blk)
+-- frag0 = Empty AnchorGenesis :> foo :> bar
+-- :}
+--
+-- >>> totalWeightOfFragment snap frag0
+-- PerasWeight 8
+--
+-- Only keeping the last block from @frag0@:
+--
+-- >>> frag1 = AF.anchorNewest 1 frag0
+-- >>> totalWeightOfFragment snap frag1
+-- PerasWeight 3
+--
+-- Dropping the head from @frag0@, and instead adding an unboosted point:
+--
+-- >>> frag2 = AF.dropNewest 1 frag0 :> HeaderFields (SlotNo 4) (BlockNo 2) "baz"
+-- >>> totalWeightOfFragment snap frag2
+-- PerasWeight 6
+totalWeightOfFragment ::
+  forall blk h.
+  (StandardHash blk, HasHeader h, HeaderHash blk ~ HeaderHash h) =>
+  PerasWeightSnapshot blk ->
+  AnchoredFragment h ->
+  PerasWeight
+totalWeightOfFragment weightSnap frag =
+  weightLength <> weightBoost
+ where
+  weightLength = PerasWeight $ fromIntegral $ AF.length frag
+  weightBoost = weightBoostOfFragment weightSnap frag
+
+-- | Take the longest suffix of the given fragment with total weight
+-- ('totalWeightOfFragment') at most @k@. This is the volatile suffix of blocks
+-- which are subject to rollback.
+--
+-- If the total weight of the input fragment is at least @k@, then the anchor of
+-- the output fragment is the most recent point on the input fragment that is
+-- buried under at least weight @k@ (also counting the weight boost of that
+-- point).
+--
+-- >>> :{
+-- weights :: [(Point Blk, PerasWeight)]
+-- weights =
+--   [ (BlockPoint 2 "foo", PerasWeight 2)
+--   , (GenesisPoint,       PerasWeight 3)
+--   , (BlockPoint 3 "bar", PerasWeight 2)
+--   , (BlockPoint 2 "foo", PerasWeight 2)
+--   ]
+-- snap = mkPerasWeightSnapshot weights
+-- foo = HeaderFields (SlotNo 2) (BlockNo 1) "foo"
+-- bar = HeaderFields (SlotNo 3) (BlockNo 2) "bar"
+-- frag :: AnchoredFragment (HeaderFields Blk)
+-- frag = Empty AnchorGenesis :> foo :> bar
+-- :}
+--
+-- >>> k1 = SecurityParam $ knownNonZeroBounded @1
+-- >>> k3 = SecurityParam $ knownNonZeroBounded @3
+-- >>> k6 = SecurityParam $ knownNonZeroBounded @6
+-- >>> k9 = SecurityParam $ knownNonZeroBounded @9
+--
+-- >>> AF.toOldestFirst $ takeVolatileSuffix snap k1 frag
+-- []
+--
+-- >>> AF.toOldestFirst $ takeVolatileSuffix snap k3 frag
+-- [HeaderFields {headerFieldSlot = SlotNo 3, headerFieldBlockNo = BlockNo 2, headerFieldHash = "bar"}]
+--
+-- >>> AF.toOldestFirst $ takeVolatileSuffix snap k6 frag
+-- [HeaderFields {headerFieldSlot = SlotNo 3, headerFieldBlockNo = BlockNo 2, headerFieldHash = "bar"}]
+--
+-- >>> AF.toOldestFirst $ takeVolatileSuffix snap k9 frag
+-- [HeaderFields {headerFieldSlot = SlotNo 2, headerFieldBlockNo = BlockNo 1, headerFieldHash = "foo"},HeaderFields {headerFieldSlot = SlotNo 3, headerFieldBlockNo = BlockNo 2, headerFieldHash = "bar"}]
+takeVolatileSuffix ::
+  forall blk h.
+  (StandardHash blk, HasHeader h, HeaderHash blk ~ HeaderHash h) =>
+  PerasWeightSnapshot blk ->
+  -- | The security parameter @k@ is interpreted as a weight.
+  SecurityParam ->
+  AnchoredFragment h ->
+  AnchoredFragment h
+takeVolatileSuffix snap secParam frag
+  | Map.null $ getPerasWeightSnapshot snap =
+      -- Optimize the case where Peras is disabled.
+      AF.anchorNewest (unPerasWeight k) frag
+  | hasAtMostWeightK frag = frag
+  | otherwise = go 0 lenFrag (AF.Empty $ AF.headAnchor frag)
+ where
+  k :: PerasWeight
+  k = maxRollbackWeight secParam
+
+  hasAtMostWeightK :: AnchoredFragment h -> Bool
+  hasAtMostWeightK f = totalWeightOfFragment snap f <= k
+
+  lenFrag = fromIntegral $ AF.length frag
+
+  -- Binary search for the longest suffix of @frag@ which 'hasAtMostWeightK'.
+  go ::
+    Word64 -> -- lb. The length lb suffix satisfies 'hasAtMostWeightK'.
+    Word64 -> -- ub. The length ub suffix does not satisfy 'hasAtMostWeightK'.
+    AnchoredFragment h -> -- The length lb suffix.
+    AnchoredFragment h
+  go lb ub lbFrag
+    | lb + 1 == ub = lbFrag
+    | hasAtMostWeightK midFrag = go mid ub midFrag
+    | otherwise = go lb mid lbFrag
+   where
+    mid = (lb + ub) `div` 2
+    midFrag = AF.anchorNewest mid frag
+
 -- $setup
+-- >>> import Cardano.Ledger.BaseTypes
 -- >>> import Ouroboros.Consensus.Block
+-- >>> import Ouroboros.Consensus.Config.SecurityParam
 -- >>> import Ouroboros.Network.AnchoredFragment (AnchoredFragment, AnchoredSeq(..), Anchor(..))
 -- >>> import qualified Ouroboros.Network.AnchoredFragment as AF
--- >>> :set -XTypeFamilies
+-- >>> :set -XDataKinds -XTypeApplications -XTypeFamilies
 -- >>> data Blk = Blk
 -- >>> type instance HeaderHash Blk = String
 -- >>> instance StandardHash Blk

ouroboros-consensus/test/consensus-test/Main.hs

@@ -16,6 +16,7 @@ import qualified Test.Consensus.MiniProtocol.BlockFetch.Client (tests)
 import qualified Test.Consensus.MiniProtocol.ChainSync.CSJ (tests)
 import qualified Test.Consensus.MiniProtocol.ChainSync.Client (tests)
 import qualified Test.Consensus.MiniProtocol.LocalStateQuery.Server (tests)
+import qualified Test.Consensus.Peras.WeightSnapshot (tests)
 import qualified Test.Consensus.Util.MonadSTM.NormalForm (tests)
 import qualified Test.Consensus.Util.Versioned (tests)
 import Test.Tasty
@@ -43,6 +44,7 @@ tests =
         , Test.Consensus.Mempool.Fairness.tests
         , Test.Consensus.Mempool.StateMachine.tests
         ]
+    , Test.Consensus.Peras.WeightSnapshot.tests
     , Test.Consensus.Util.MonadSTM.NormalForm.tests
     , Test.Consensus.Util.Versioned.tests
     , testGroup