Move a couple of mempool operations to work on ledgerstates instead

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/HardFork/Combinator/Mempool.hs

@@ -29,6 +29,7 @@ module Ouroboros.Consensus.HardFork.Combinator.Mempool (
 
 import           Control.Arrow (first, (+++))
 import           Control.Monad.Except
+import           Data.Functor.Identity
 import           Data.Functor.Product
 import           Data.Kind (Type)
 import qualified Data.Measure as Measure
@@ -40,6 +41,7 @@ import           Data.SOP.InPairs (InPairs)
 import qualified Data.SOP.InPairs as InPairs
 import qualified Data.SOP.Match as Match
 import           Data.SOP.Strict
+import qualified Data.SOP.Telescope as Tele
 import           Data.Typeable (Typeable)
 import           GHC.Generics (Generic)
 import           NoThunks.Class (NoThunks)
@@ -224,6 +226,50 @@ instance ( CanHardFork xs
         -> K (LedgerTables (LedgerState (HardForkBlock xs)) KeysMK) x
       f idx tx = K $ injectLedgerTables idx $ getTransactionKeySets tx
 
+  -- This optimization is worthwile because we can save the projection and
+  -- injection of ledger tables.
+  --
+  -- These operations are used when adding new transactions to the mempool,
+  -- which is _not_ in the critical path for the forging loop but still will
+  -- make adoption of new transactions faster. As adding a transaction takes a
+  -- TMVar, it is interesting to hold it for as short of a time as possible.
+  prependMempoolDiffs
+    (TickedHardForkLedgerState _ (State.HardForkState st1))
+    (TickedHardForkLedgerState tr (State.HardForkState st2))
+    = TickedHardForkLedgerState
+          tr
+        $ State.HardForkState
+        $ runIdentity
+         (Tele.alignExtend
+         (InPairs.hpure (error "When prepending mempool diffs we used to un-aligned states, this should be impossible!"))
+         (hcpure proxySingle $ fn_2 $ \(State.Current _ a) (State.Current start b) -> State.Current start $
+                         FlipTickedLedgerState
+                       $ prependMempoolDiffs
+                          (getFlipTickedLedgerState a)
+                          (getFlipTickedLedgerState b)
+                      )
+         st1
+         st2)
+
+  -- This optimization is worthwile because we can save the projection and
+  -- injection of ledger tables.
+  --
+  -- These operations are used when adding new transactions to the mempool,
+  -- which is _not_ in the critical path for the forging loop but still will
+  -- make adoption of new transactions faster. As adding a transaction takes a
+  -- TMVar, it is interesting to hold it for as short of a time as possible.
+  applyMempoolDiffs
+    vals keys (TickedHardForkLedgerState tr (State.HardForkState st)) =
+    TickedHardForkLedgerState tr $ State.HardForkState $ hcimap
+      proxySingle
+      (\idx (State.Current start (FlipTickedLedgerState a)) ->
+         State.Current start $ FlipTickedLedgerState
+         $ applyMempoolDiffs
+           (ejectLedgerTables idx vals)
+           (ejectLedgerTables idx keys) a )
+      st
+
+
 instance CanHardFork xs => TxLimits (HardForkBlock xs) where
   type TxMeasure (HardForkBlock xs) = HardForkTxMeasure xs
 

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Ledger/SupportsMempool.hs

@@ -198,6 +198,33 @@ class ( UpdateLedger blk
   -- transaction size.
   getTransactionKeySets :: GenTx blk -> LedgerTables (LedgerState blk) KeysMK
 
+  -- Mempools live in a single slot so in the hard fork block case
+  -- it is cheaper to perform these operations on LedgerStates, saving
+  -- the time of projecting and injecting ledger tables.
+  --
+  -- The cost of this when adding transactions is very small compared
+  -- to eg the networking costs of mempool synchronization, but still
+  -- it is worthwile locking the mempool for as short as possible.
+  --
+  -- Eventually the Ledger will provide these diffs, so we might even
+  -- be able to remove this optimization altogether.
+
+  -- | Prepend diffs on ledger states
+  prependMempoolDiffs ::
+       TickedLedgerState blk DiffMK
+    -> TickedLedgerState blk DiffMK
+    -> TickedLedgerState blk DiffMK
+  prependMempoolDiffs = prependDiffs
+
+  -- | Apply diffs on ledger states
+  applyMempoolDiffs ::
+       LedgerTables (LedgerState blk) ValuesMK
+    -> LedgerTables (LedgerState blk) KeysMK
+    -> TickedLedgerState blk DiffMK
+    -> TickedLedgerState blk ValuesMK
+  applyMempoolDiffs = applyDiffForKeysOnTables
+
+
 data ReapplyTxsResult extra blk =
   ReapplyTxsResult {
       -- | txs that are now invalid. Order doesn't matter

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Mempool/Impl/Common.hs

@@ -305,7 +305,7 @@ validateNewTransaction cfg wti tx txsz origValues st is =
              , isTxKeys       = isTxKeys <> getTransactionKeySets tx
              , isTxValues     = ltliftA2 unionValues isTxValues origValues
              , isTxIds        = Set.insert (txId tx) isTxIds
-             , isLedgerState  = prependDiffs isLedgerState st'
+             , isLedgerState  = prependMempoolDiffs isLedgerState st'
              , isLastTicketNo = nextTicketNo
              }
         )
@@ -347,7 +347,7 @@ revalidateTxsFor capacityOverride cfg slot st values lastTicketNo txTickets =
       unwrap = (\(tx, (tk, tz)) -> TxTicket tx tk tz)
       ReapplyTxsResult err val st' =
         reapplyTxs ComputeDiffs cfg slot theTxs
-        $ applyDiffForKeysOnTables
+        $ applyMempoolDiffs
               values
               (Foldable.foldMap' (getTransactionKeySets . txForgetValidated . fst) theTxs)
               st
@@ -395,7 +395,7 @@ computeSnapshot capacityOverride cfg slot st values lastTicketNo txTickets =
       unwrap = (\(tx, (tk, tz)) -> TxTicket tx tk tz)
       ReapplyTxsResult _ val st' =
         reapplyTxs IgnoreDiffs cfg slot theTxs
-        $ applyDiffForKeysOnTables
+        $ applyMempoolDiffs
               values
               (Foldable.foldMap' (getTransactionKeySets . txForgetValidated . fst) theTxs)
               st

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Mempool/Update.hs

@@ -24,7 +24,6 @@ import           Data.Void
 import           Ouroboros.Consensus.HeaderValidation
 import           Ouroboros.Consensus.Ledger.Abstract
 import           Ouroboros.Consensus.Ledger.SupportsMempool
-import           Ouroboros.Consensus.Ledger.Tables.Utils
 import           Ouroboros.Consensus.Mempool.API
 import           Ouroboros.Consensus.Mempool.Capacity
 import           Ouroboros.Consensus.Mempool.Impl.Common
@@ -214,7 +213,7 @@ pureTryAddTx ::
   -> LedgerTables (LedgerState blk) ValuesMK
   -> TriedToAddTx blk
 pureTryAddTx cfg wti tx is values =
-  let st = applyDiffForKeysOnTables values (getTransactionKeySets tx) (isLedgerState is) in
+  let st = applyMempoolDiffs values (getTransactionKeySets tx) (isLedgerState is) in
   case runExcept $ txMeasure cfg st tx of
     Left err ->
       -- The transaction does not have a valid measure (eg its ExUnits is