mempool: adapt for generalized validation

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

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

Author: crocodile-dentist

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

@@ -1,5 +1,7 @@
-{-# LANGUAGE NamedFieldPuns      #-}
-{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE DisambiguateRecordFields #-}
+{-# LANGUAGE NamedFieldPuns           #-}
+{-# LANGUAGE ScopedTypeVariables      #-}
+{-# LANGUAGE TupleSections            #-}
 
 -- | The module should be imported qualified.
 --
@@ -16,19 +18,19 @@ module Ouroboros.Network.TxSubmission.Mempool.Simple
 import Prelude hiding (read, seq)
 
 import Control.Concurrent.Class.MonadSTM.Strict
-
+import Data.Bitraversable
+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
 import Data.Set (Set)
 import Data.Set qualified as Set
 
+import Ouroboros.Network.Protocol.LocalTxSubmission.Type (SubmitResult (..))
 import Ouroboros.Network.SizeInBytes
-import Ouroboros.Network.TxSubmission.Inbound.V2.Types
 import Ouroboros.Network.TxSubmission.Mempool.Reader
 
 
@@ -98,31 +100,67 @@ getReader getTxId getTxSize (Mempool mempool) =
     f idx tx = (getTxId tx, idx, getTxSize tx)
 
 
--- | A simple mempool writer.
+-- | A mempool writer which generalizes the tx submission mempool writer
+--
+data MempoolWriter txid tx failure idx m =
+     MempoolWriter {
+
+       -- | Compute the transaction id from a transaction.
+       --
+       -- This is used in the protocol handler to verify a full transaction
+       -- matches a previously given transaction id.
+       --
+       txId          :: tx -> txid,
+
+       -- | Supply a batch of transactions to the mempool. They are either
+       -- accepted or rejected individually, but in the order supplied.
+       --
+       -- The 'txid's of all transactions that were added successfully are
+       -- returned.
+       mempoolAddTxs :: [tx] -> m [SubmitResult failure]
+    }
+
+
+-- | A mempool writer with validation harness
+-- PRECONDITION: no duplicates given to mempoolAddTxs
 --
-getWriter :: forall tx txid m.
+getWriter :: forall tx txid ctx tx' failure m.
              ( MonadSTM m
              , Ord txid
              )
           => (tx -> txid)
-          -> (tx -> Bool)
-          -- ^ validate a tx
-          -> Mempool m tx
-          -> TxSubmissionMempoolWriter txid tx Int m
-getWriter getTxId validateTx (Mempool mempool) =
-    TxSubmissionMempoolWriter {
-        txId = getTxId,
-
-        mempoolAddTxs = \txs -> do
-          atomically $ do
-            mempoolTxs <- readTVar mempool
-            let currentIds = Set.fromList (map getTxId (toList mempoolTxs))
-                validTxs = nubBy (on (==) getTxId)
-                         $ filter
-                            (\tx -> validateTx tx
-                                 && getTxId tx `Set.notMember` currentIds)
-                           txs
-                mempoolTxs' = Foldable.foldl' (Seq.|>) mempoolTxs validTxs
-            writeTVar mempool mempoolTxs'
-            return (map getTxId validTxs)
-      }
+          -- ^ get txid of a tx
+          -> m ctx
+          -- ^ monadic validation context, acquired once prior to all work
+          -> (ctx -> tx -> tx')
+          -- ^ pre-process every transanction with the context
+          -> (tx' -> Bool -> Either failure tx)
+          -- ^ validate a tx in an atomic block, any failing `tx` throws an exception.
+          -> (failure -> STM m failure)
+          -- ^ return `True` when a failure should throw an exception
+          -> Mempool m txid tx
+          -> MempoolWriter txid tx failure Int m
+getWriter getTxId acquireCtx preProcess validateTx failureFilterFn (Mempool mempool) =
+  MempoolWriter {
+    txId = getTxId,
+
+    mempoolAddTxs = \txs -> do
+      ctx <- acquireCtx
+      let txs' = preProcess ctx <$> txs -- TODO could run in parallel
+      atomically $ do
+        MempoolSeq { mempoolSet, mempoolSeq } <- readTVar mempool
+        result <- sequence
+          [bimapM (fmap SubmitFail . failureFilterFn) (pure . (txid,)) validated
+          | (tx, tx') <- zip txs txs'
+          , let txid = getTxId tx
+                validated =
+                  validateTx tx' (txid `Set.member` mempoolSet)
+          ]
+        let (validIds, validTxs) = unzip . rights $ result
+            mempoolTxs' = MempoolSeq {
+               mempoolSet = Set.union mempoolSet (Set.fromList validIds),
+               mempoolSeq = Foldable.foldl' (Seq.|>) mempoolSeq validTxs
+            }
+        writeTVar mempool mempoolTxs'
+        return $ fromLeft SubmitSuccess <$> result
+  }