Improve performance of the signing window check.

Author: nc6

nix/.stack.nix/ouroboros-consensus.nix

@@ -362,6 +362,16 @@ test-suite test-storage
   ghc-options:      -Wall
                     -fno-ignore-asserts
 
+test-suite doctests
+  type:             exitcode-stdio-1.0
+  hs-source-dirs:   test-doctest
+  ghc-options:      -threaded
+  main-is:          doctests.hs
+  build-depends:    base
+                  , doctest >= 0.8
+                  , ouroboros-consensus
+                  , QuickCheck
+
 executable byron-db-converter
    hs-source-dirs:     tools/db-convert
    build-depends:      base

ouroboros-consensus/ouroboros-consensus.cabal

@@ -19,6 +19,7 @@ module Ouroboros.Consensus.Node.ProtocolInfo.Byron (
 
 import           Control.Exception (Exception)
 import           Control.Monad.Except
+import qualified Data.Map.Strict as Map
 import qualified Data.Sequence as Seq
 import qualified Data.Set as Set
 import           Data.Maybe
@@ -160,7 +161,7 @@ protocolInfoByron genesisConfig@Genesis.Config {
                 blsCurrent   = initState
               , blsSnapshots = Seq.empty
               }
-          , ouroborosChainState = Seq.empty
+          , ouroborosChainState = Map.empty
           }
       , pInfoInitState  = ()
       }

ouroboros-consensus/src/Ouroboros/Consensus/Node/ProtocolInfo/Byron.hs

@@ -7,7 +7,7 @@ module Ouroboros.Consensus.Node.ProtocolInfo.Mock.PBFT (
 
 import           Codec.Serialise (Serialise (..))
 import qualified Data.Bimap as Bimap
-import qualified Data.Sequence as Seq
+import qualified Data.Map.Strict as Map
 
 import           Cardano.Crypto.DSIGN
 
@@ -41,7 +41,7 @@ protocolInfoMockPBFT (NumCoreNodes numCoreNodes) (CoreNodeId nid) params =
                                ]
           }
       , pInfoInitLedger = ExtLedgerState (genesisSimpleLedgerState addrDist)
-                                         Seq.empty
+                                         Map.empty
       , pInfoInitState  = ()
       }
   where

ouroboros-consensus/src/Ouroboros/Consensus/Node/ProtocolInfo/Mock/PBFT.hs

@@ -2,11 +2,13 @@
 {-# LANGUAGE DeriveGeneric           #-}
 {-# LANGUAGE FlexibleContexts        #-}
 {-# LANGUAGE FlexibleInstances       #-}
+{-# LANGUAGE LambdaCase              #-}
 {-# LANGUAGE MultiParamTypeClasses   #-}
 {-# LANGUAGE NamedFieldPuns          #-}
 {-# LANGUAGE RecordWildCards         #-}
 {-# LANGUAGE ScopedTypeVariables     #-}
 {-# LANGUAGE StandaloneDeriving      #-}
+{-# LANGUAGE TupleSections           #-}
 {-# LANGUAGE TypeFamilyDependencies  #-}
 {-# LANGUAGE TypeOperators           #-}
 {-# LANGUAGE UndecidableInstances    #-}
@@ -35,7 +37,11 @@ import           Crypto.Random (MonadRandom)
 import           Data.Bimap (Bimap)
 import qualified Data.Bimap as Bimap
 import           Data.Constraint
+import           Data.List (sortOn)
 import           Data.Reflection (Given (..), give)
+import           Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
+import           Data.Maybe (catMaybes, fromJust)
 import qualified Data.Set as Set
 import           Data.Sequence (Seq)
 import qualified Data.Sequence as Seq
@@ -167,11 +173,10 @@ instance (PBftCrypto c, Typeable c) => OuroborosTag (PBft c) where
 
   type IsLeader       (PBft c) = PBftIsLeader c
 
-  -- | Chain state consists of two things:
-  --   - a list of the last 'pbftSignatureWindow' signatures.
-  --   - The last seen block slot
+  -- | Chain state consists of a map from genesis keys to the list of blocks
+  -- which they have issued.
   type ChainState     (PBft c) =
-    Seq (PBftVerKeyHash c, SlotNo)
+    Map (PBftVerKeyHash c) (Seq SlotNo)
 
   protocolSecurityParam = pbftSecurityParam . pbftParams
 
@@ -201,9 +206,13 @@ instance (PBftCrypto c, Typeable c) => OuroborosTag (PBft c) where
         Right () -> return ()
         Left err -> throwError $ PBftInvalidSignature err
 
-      let (signers, lastSlot) = ( takeR winSize $ fst <$> chainState
-                                , maybe (SlotNo 0) snd $ Seq.lookup (Seq.length chainState) chainState
-                                )
+      -- We always include slot number 0 in case there are no signers yet.
+      let lastSlot = maximum . (SlotNo 0 :)
+                             $ (\case
+                                  _ Seq.:|> l -> l
+                                  _ -> SlotNo 0
+                               )
+                             <$> Map.elems chainState
 
       -- FIXME confirm that non-strict inequality is ok in general.
       -- It's here because EBBs have the same slot as the first block of their
@@ -214,26 +223,101 @@ instance (PBftCrypto c, Typeable c) => OuroborosTag (PBft c) where
       case Bimap.lookupR (hashVerKey pbftIssuer) dms of
         Nothing -> throwError $ PBftNotGenesisDelegate (hashVerKey pbftIssuer) lv
         Just gk -> do
-          let totalSigners = Seq.length signers
-              gkSigners = Seq.length (Seq.filter (== gk) signers)
+          let totalSigners = chainStateSize chainState
+              gkSigners = maybe 0 Seq.length $ Map.lookup gk chainState
           when (totalSigners >= winSize && gkSigners > wt)
             $ throwError (PBftExceededSignThreshold totalSigners gkSigners)
-          return $! takeR (winSize + 2*k) chainState Seq.|> (gk, blockSlot b)
+          return $! insertSigner gk (blockSlot b) $ pruneChainState (winSize + 2*k) chainState
     where
       PBftParams{..} = pbftParams
       PBftFields{..} = headerPBftFields cfg b
       winSize = fromIntegral pbftSignatureWindow
       SecurityParam (fromIntegral -> k) = pbftSecurityParam
       wt = floor $ pbftSignatureThreshold * fromIntegral winSize
-      -- Take the rightmost n elements of a sequence
-      takeR :: Integral i => i -> Seq a -> Seq a
-      takeR (fromIntegral -> n) s = Seq.drop (Seq.length s - n - 1) s
 
   rewindChainState _ cs mSlot = case mSlot of
-    Origin  -> Just Seq.empty
-    At slot -> case Seq.takeWhileL (\(_, s) -> s <= slot) cs of
-        _ Seq.:<| _ -> Just cs
-        _           -> Nothing
+    Origin  -> Just Map.empty
+    At slot -> let oldCs = Map.map (Seq.takeWhileL (\s -> s <= slot)) cs in
+      if (all Seq.null $ Map.elems oldCs) then Nothing else Just oldCs
+
+-- | Prune the chain state to the given size by dropping the signers in the
+-- oldest slots.
+--
+-- Examples
+--
+-- $setup
+-- >>> :set -XOverloadedLists
+-- >>> :set -XScopedTypeVariables
+-- >>> import Test.QuickCheck
+-- >>> :{
+--   newtype TestChainState = TestChainState (Map Int (Seq Int))
+--   instance Arbitrary TestChainState where
+--     arbitrary = do
+--       numElts <- arbitrary
+--       ub <- arbitrary
+--       return . TestChainState $ Map.fromList
+--         [(idx, elts) | idx <- [0..numElts]
+--                      , let elts = Seq.fromList [i | i <- [0..ub], i `mod` numElts == idx]
+--         ]
+-- :}
+
+-- prop> \cs -> pruneChainState (chainStateSize cs) cs == cs
+--
+-- prop> \n (TestChainState cs) -> n > 0 && (chainStateSize cs) > n ==> (chainStateSize (pruneChainState n cs)) == n
+--
+-- >>> :{
+--   let (cs :: Map Char (Seq Int)) = [('a', [1,2,5]), ('b', [3, 4, 6])]
+--   in pruneChainState 2 cs
+-- :}
+-- fromList [('a',fromList [5]),('b',fromList [6])]
+--
+pruneChainState :: forall k v. (Ord k, Ord v) => Int -> Map k (Seq v) -> Map k (Seq v)
+pruneChainState toSize cs = go
+    cs
+    (sortOn snd . catMaybes $ strengthr . (fmap (Seq.lookup 0)) <$> (Map.toAscList cs))
+    (max 0 $ chainStateSize cs - toSize)
+  where
+    go :: Map k (Seq v)
+       -> [(k, v)]
+       -> Int
+       -> Map k (Seq v)
+    go fromCS idx toDrop = if toDrop <= 0 then fromCS else case idx of
+      [] -> fromCS
+      (gk,_):xs@((_,nextLowest):_) ->
+        let (newSeq, numDropped) = fromJust $ dropWhileL (< nextLowest) <$> Map.lookup gk fromCS
+            newIdx = case newSeq of
+              x Seq.:<| _ -> sortOn snd $ (gk, x) : xs
+              _ -> xs
+        in go (Map.insert gk newSeq fromCS) newIdx (toDrop - numDropped)
+      -- Only one genesis key
+      (gk,_):[] ->
+        let newSeq = fromJust $ Seq.drop toDrop <$> Map.lookup gk fromCS
+        in Map.insert gk newSeq fromCS
+
+chainStateSize :: Map k (Seq v) -> Int
+chainStateSize cs = sum $ Seq.length <$> Map.elems cs
+
+-- | Functorial strength on the right
+strengthr :: Functor f => (a, f b) -> f (a, b)
+strengthr (a, fb) = fmap (a,) fb
+
+-- | Variant of 'dropWhileL' which also returns the number of elements dropped
+dropWhileL :: (a -> Bool) -> Seq a -> (Seq a, Int)
+dropWhileL f s = let res = Seq.dropWhileL f s in
+  (res, Seq.length s - Seq.length res)
+
+-- | Insert a signatory into the chain state.
+insertSigner
+  :: PBftCrypto c
+  => PBftVerKeyHash c
+  -> SlotNo
+  -> ChainState (PBft c)
+  -> ChainState (PBft c)
+insertSigner gk s =
+  Map.alter (\case
+      Just es -> Just $ es Seq.|> s
+      Nothing -> Just $ Seq.singleton s
+    ) gk
 
 {-------------------------------------------------------------------------------
   PBFT node order

ouroboros-consensus/src/Ouroboros/Consensus/Protocol/PBFT.hs

@@ -0,0 +1,3 @@
+import Test.DocTest
+
+main = doctest ["-isrc", "src/Ouroboros/Consensus/Protocol/PBFT.hs"]