mempool: adapt for generalized validation

Author: crocodile-dentist

ouroboros-network/src/Ouroboros/Network/TxSubmission/Mempool/Simple.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE NamedFieldPuns      #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections       #-}
 
 -- | The module should be imported qualified.
 --
@@ -16,13 +17,11 @@ module Ouroboros.Network.TxSubmission.Mempool.Simple
 import Prelude hiding (read, seq)
 
 import Control.Concurrent.Class.MonadSTM.Strict
-import Control.Monad (filterM)
-import Control.Monad.Class.MonadTime.SI
-
+import Data.Bifunctor (bimap)
+import Data.Either
 import Data.Foldable (toList)
 import Data.Foldable qualified as Foldable
-import Data.Function (on)
-import Data.List (find, nubBy)
+import Data.List (find)
 import Data.Maybe (isJust)
 import Data.Sequence (Seq)
 import Data.Sequence qualified as Seq
@@ -102,9 +101,8 @@ getReader getTxId getTxSize (Mempool mempool) =
 
 -- | A simple mempool writer.
 --
-getWriter :: forall tx txid ctx metadata m.
+getWriter :: forall tx txid ctx failure m.
              ( MonadSTM m
-             , MonadTime m
              , Ord txid
              )
           => (tx -> txid)
@@ -121,38 +119,22 @@ getWriter getTxId acquireCtx validateTx failureFilterFn (Mempool mempool) =
     TxSubmissionMempoolWriter {
         txId = getTxId,
 
-          -- let metaTxs = getMetaData <$> txs
-          -- now <- getCurrentTime
-          -- atomically $ do
-          --   mempoolTxs <- readTVar mempool
-          --   let currentIds = Set.fromList (map getTxId (toList mempoolTxs))
-          --   validTxs <-   fmap (fmap fst <$> nubBy ((==) `on` getTxId . fst))
-          --               . filterM
-          --                  (\(tx, meta) -> do
-          --                    valid <- validateTx now meta
-          --                    return $ valid && getTxId tx `Set.notMember` currentIds)
-          --               $ zip txs metaTxs
-          --   let mempoolTxs' = Foldable.foldl' (Seq.|>) mempoolTxs validTxs
-          --   writeTVar mempool mempoolTxs'
         mempoolAddTxs = \txs -> do
           ctx <- acquireCtx
-          (invalidTxIds, validTxs) <- atomically $ do
+          atomically $ do
             MempoolSeq { mempoolSet, mempoolSeq } <- readTVar mempool
-            let (invalidTxIds, validTxs) =
-                    first (filter (failureFilterFn . snd))
-                  . partitionEithers
-                  . map (\tx -> case validateTx ctx tx of
-                                 Left  e -> Left (getTxId tx, e)
-                                 Right _ -> Right tx
-                        )
-                  . filter (\tx -> getTxId tx `Set.notMember` mempoolSet)
-                  $ txs
-                mempoolTxs' = MempoolSeq {
-                    mempoolSet = Foldable.foldl' (\s tx -> getTxId tx `Set.insert` s)
-                                                 mempoolSet
-                                                 validTxs,
-                    mempoolSeq = Foldable.foldl' (Seq.|>) mempoolSeq validTxs
-                  }
+            (invalidTxIds, (validIds, validTxs)) <-
+                bimap (filter (failureFilterFn . snd)) unzip
+              . partitionEithers <$>
+                  sequence [bimap (txid,) (const (txid, tx)) <$> validateTx ctx tx
+                           | tx <- txs
+                           , let txid = getTxId tx
+                           , txid `Set.notMember` mempoolSet
+                           ]
+            let mempoolTxs' = MempoolSeq {
+                  mempoolSet = Set.union mempoolSet (Set.fromList validIds),
+                  mempoolSeq = Foldable.foldl' (Seq.|>) mempoolSeq validTxs
+                }
             writeTVar mempool mempoolTxs'
-            return (map getTxId validTxs)
+            return validIds
       }