{-# LANGUAGE DerivingStrategies  #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE NamedFieldPuns      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}

-- | The module should be imported qualified.
--
module Ouroboros.Network.TxSubmission.Mempool.Simple
  ( Mempool (..)
  , MempoolSeq (..)
  , empty
  , new
  , read
  , getReader
  , getWriter
  ) where

import Prelude hiding (read, seq)

import Control.Concurrent.Class.MonadSTM.Strict
import Control.Monad (when)
import Control.Monad.Class.MonadThrow

import Data.Bifunctor (bimap)
import Data.Either (partitionEithers)
import Data.Foldable (toList)
import Data.Foldable qualified as Foldable
import Data.Function (on)
import Data.List (find, nubBy)
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 Data.Typeable (Typeable)

import Ouroboros.Network.SizeInBytes
import Ouroboros.Network.TxSubmission.Inbound.V2.Types
import Ouroboros.Network.TxSubmission.Mempool.Reader


data MempoolSeq txid tx = MempoolSeq {
    forall txid tx. MempoolSeq txid tx -> Set txid
mempoolSet :: !(Set txid),
    -- ^ cached set of `txid`s in the mempool
    forall txid tx. MempoolSeq txid tx -> Seq tx
mempoolSeq :: !(Seq tx)
    -- ^ sequence of all `tx`s
  }

-- | A simple in-memory mempool implementation.
--
newtype Mempool m txid tx = Mempool (StrictTVar m (MempoolSeq txid tx))


empty :: MonadSTM m => m (Mempool m txid tx)
empty :: forall (m :: * -> *) txid tx. MonadSTM m => m (Mempool m txid tx)
empty = StrictTVar m (MempoolSeq txid tx) -> Mempool m txid tx
forall (m :: * -> *) txid tx.
StrictTVar m (MempoolSeq txid tx) -> Mempool m txid tx
Mempool (StrictTVar m (MempoolSeq txid tx) -> Mempool m txid tx)
-> m (StrictTVar m (MempoolSeq txid tx)) -> m (Mempool m txid tx)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MempoolSeq txid tx -> m (StrictTVar m (MempoolSeq txid tx))
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
newTVarIO (Set txid -> Seq tx -> MempoolSeq txid tx
forall txid tx. Set txid -> Seq tx -> MempoolSeq txid tx
MempoolSeq Set txid
forall a. Set a
Set.empty Seq tx
forall a. Seq a
Seq.empty)


new :: ( MonadSTM m
       , Ord txid
       )
    => (tx -> txid)
    -> [tx]
    -> m (Mempool m txid tx)
new :: forall (m :: * -> *) txid tx.
(MonadSTM m, Ord txid) =>
(tx -> txid) -> [tx] -> m (Mempool m txid tx)
new tx -> txid
getTxId [tx]
txs =
    (StrictTVar m (MempoolSeq txid tx) -> Mempool m txid tx)
-> m (StrictTVar m (MempoolSeq txid tx)) -> m (Mempool m txid tx)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StrictTVar m (MempoolSeq txid tx) -> Mempool m txid tx
forall (m :: * -> *) txid tx.
StrictTVar m (MempoolSeq txid tx) -> Mempool m txid tx
Mempool
  (m (StrictTVar m (MempoolSeq txid tx)) -> m (Mempool m txid tx))
-> (MempoolSeq txid tx -> m (StrictTVar m (MempoolSeq txid tx)))
-> MempoolSeq txid tx
-> m (Mempool m txid tx)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MempoolSeq txid tx -> m (StrictTVar m (MempoolSeq txid tx))
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
newTVarIO
  (MempoolSeq txid tx -> m (Mempool m txid tx))
-> MempoolSeq txid tx -> m (Mempool m txid tx)
forall a b. (a -> b) -> a -> b
$ MempoolSeq { mempoolSet :: Set txid
mempoolSet = [txid] -> Set txid
forall a. Ord a => [a] -> Set a
Set.fromList (tx -> txid
getTxId (tx -> txid) -> [tx] -> [txid]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [tx]
txs),
                 mempoolSeq :: Seq tx
mempoolSeq = [tx] -> Seq tx
forall a. [a] -> Seq a
Seq.fromList [tx]
txs
               }


read :: MonadSTM m => Mempool m txid tx -> m [tx]
read :: forall (m :: * -> *) txid tx.
MonadSTM m =>
Mempool m txid tx -> m [tx]
read (Mempool StrictTVar m (MempoolSeq txid tx)
mempool) = Seq tx -> [tx]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Seq tx -> [tx])
-> (MempoolSeq txid tx -> Seq tx) -> MempoolSeq txid tx -> [tx]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MempoolSeq txid tx -> Seq tx
forall txid tx. MempoolSeq txid tx -> Seq tx
mempoolSeq (MempoolSeq txid tx -> [tx]) -> m (MempoolSeq txid tx) -> m [tx]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StrictTVar m (MempoolSeq txid tx) -> m (MempoolSeq txid tx)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> m a
readTVarIO StrictTVar m (MempoolSeq txid tx)
mempool


getReader :: forall tx txid m.
             ( MonadSTM m
             , Eq txid
             )
          => (tx -> txid)
          -> (tx -> SizeInBytes)
          -> Mempool m txid tx
          -> TxSubmissionMempoolReader txid tx Int m
getReader :: forall tx txid (m :: * -> *).
(MonadSTM m, Eq txid) =>
(tx -> txid)
-> (tx -> SizeInBytes)
-> Mempool m txid tx
-> TxSubmissionMempoolReader txid tx Int m
getReader tx -> txid
getTxId tx -> SizeInBytes
getTxSize (Mempool StrictTVar m (MempoolSeq txid tx)
mempool) =
    -- Using `0`-based index.  `mempoolZeroIdx = -1` so that
    -- `mempoolTxIdsAfter mempoolZeroIdx` returns all txs.
    TxSubmissionMempoolReader { STM m (MempoolSnapshot txid tx Int)
mempoolGetSnapshot :: STM m (MempoolSnapshot txid tx Int)
mempoolGetSnapshot :: STM m (MempoolSnapshot txid tx Int)
mempoolGetSnapshot,
                                mempoolZeroIdx :: Int
mempoolZeroIdx = -Int
1
                              }
  where
    mempoolGetSnapshot :: STM m (MempoolSnapshot txid tx Int)
    mempoolGetSnapshot :: STM m (MempoolSnapshot txid tx Int)
mempoolGetSnapshot = Seq tx -> MempoolSnapshot txid tx Int
getSnapshot (Seq tx -> MempoolSnapshot txid tx Int)
-> (MempoolSeq txid tx -> Seq tx)
-> MempoolSeq txid tx
-> MempoolSnapshot txid tx Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MempoolSeq txid tx -> Seq tx
forall txid tx. MempoolSeq txid tx -> Seq tx
mempoolSeq (MempoolSeq txid tx -> MempoolSnapshot txid tx Int)
-> STM m (MempoolSeq txid tx)
-> STM m (MempoolSnapshot txid tx Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StrictTVar m (MempoolSeq txid tx) -> STM m (MempoolSeq txid tx)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m (MempoolSeq txid tx)
mempool

    getSnapshot :: Seq tx
                -> MempoolSnapshot txid tx Int
    getSnapshot :: Seq tx -> MempoolSnapshot txid tx Int
getSnapshot Seq tx
seq =
      MempoolSnapshot {
          mempoolTxIdsAfter :: Int -> [(txid, Int, SizeInBytes)]
mempoolTxIdsAfter = \Int
idx -> (Int -> tx -> (txid, Int, SizeInBytes))
-> [Int] -> [tx] -> [(txid, Int, SizeInBytes)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> tx -> (txid, Int, SizeInBytes)
f [Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1..]
                                                (Seq tx -> [tx]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Seq tx -> [tx]) -> Seq tx -> [tx]
forall a b. (a -> b) -> a -> b
$ Int -> Seq tx -> Seq tx
forall a. Int -> Seq a -> Seq a
Seq.drop (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Seq tx
seq),
          mempoolLookupTx :: Int -> Maybe tx
mempoolLookupTx   = \Int
idx -> Int -> Seq tx -> Maybe tx
forall a. Int -> Seq a -> Maybe a
Seq.lookup Int
idx Seq tx
seq,
          mempoolHasTx :: txid -> Bool
mempoolHasTx      = \txid
txid -> Maybe tx -> Bool
forall a. Maybe a -> Bool
isJust (Maybe tx -> Bool) -> Maybe tx -> Bool
forall a b. (a -> b) -> a -> b
$ (tx -> Bool) -> Seq tx -> Maybe tx
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\tx
tx -> tx -> txid
getTxId tx
tx txid -> txid -> Bool
forall a. Eq a => a -> a -> Bool
== txid
txid) Seq tx
seq
       }

    f :: Int -> tx -> (txid, Int, SizeInBytes)
    f :: Int -> tx -> (txid, Int, SizeInBytes)
f Int
idx tx
tx = (tx -> txid
getTxId tx
tx, Int
idx, tx -> SizeInBytes
getTxSize tx
tx)


data InvalidTxsError where
    InvalidTxsError :: forall txid failure.
                       ( Typeable txid
                       , Typeable failure
                       , Show txid
                       , Show failure
                       )
                    => [(txid, failure)]
                    -> InvalidTxsError

deriving instance Show InvalidTxsError
instance Exception InvalidTxsError


-- | A simple mempool writer.
--
getWriter :: forall tx txid ctx failure m.
             ( MonadSTM m
             , MonadThrow m
             , Ord txid
             , Typeable txid
             , Typeable failure
             , Show txid
             , Show failure
             )
          => (tx -> txid)
          -- ^ get txid of a tx
          -> m ctx
          -- ^ monadic validation ctx
          -> (ctx -> tx -> Either failure ())
          -- ^ validate a tx, any failing `tx` throws an exception.
          -> (failure -> Bool)
          -- ^ return `True` when a failure should throw an exception
          -> Mempool m txid tx
          -> TxSubmissionMempoolWriter txid tx Int m
getWriter :: forall tx txid ctx failure (m :: * -> *).
(MonadSTM m, MonadThrow m, Ord txid, Typeable txid,
 Typeable failure, Show txid, Show failure) =>
(tx -> txid)
-> m ctx
-> (ctx -> tx -> Either failure ())
-> (failure -> Bool)
-> Mempool m txid tx
-> TxSubmissionMempoolWriter txid tx Int m
getWriter tx -> txid
getTxId m ctx
getValidationCtx ctx -> tx -> Either failure ()
validateTx failure -> Bool
failureFilterFn (Mempool StrictTVar m (MempoolSeq txid tx)
mempool) =
    TxSubmissionMempoolWriter {
        txId :: tx -> txid
txId = tx -> txid
getTxId,

        mempoolAddTxs :: [tx] -> m [txid]
mempoolAddTxs = \[tx]
txs -> do
          ctx <- m ctx
getValidationCtx
          (invalidTxIds, validTxs) <- atomically $ do
            MempoolSeq { mempoolSet, mempoolSeq } <- readTVar mempool
            let (invalidTxIds, validTxs) =
                    bimap (filter (failureFilterFn . snd))
                          (nubBy (on (==) getTxId))
                  . partitionEithers
                  . map (\tx
tx -> case ctx -> tx -> Either failure ()
validateTx ctx
ctx tx
tx of
                                 Left  failure
e -> (txid, failure) -> Either (txid, failure) tx
forall a b. a -> Either a b
Left (tx -> txid
getTxId tx
tx, failure
e)
                                 Right ()
_ -> tx -> Either (txid, failure) tx
forall a b. b -> Either a b
Right tx
tx
                        )
                  . filter (\tx
tx -> tx -> txid
getTxId tx
tx txid -> Set txid -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set txid
mempoolSet)
                  $ txs
                mempoolTxs' = MempoolSeq {
                    mempoolSet :: Set txid
mempoolSet = (Set txid -> tx -> Set txid) -> Set txid -> [tx] -> Set txid
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Set txid
s tx
tx -> tx -> txid
getTxId tx
tx txid -> Set txid -> Set txid
forall a. Ord a => a -> Set a -> Set a
`Set.insert` Set txid
s)
                                                 Set txid
mempoolSet
                                                 [tx]
validTxs,
                    mempoolSeq :: Seq tx
mempoolSeq = (Seq tx -> tx -> Seq tx) -> Seq tx -> [tx] -> Seq tx
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' Seq tx -> tx -> Seq tx
forall a. Seq a -> a -> Seq a
(Seq.|>) Seq tx
mempoolSeq [tx]
validTxs
                  }
            writeTVar mempool mempoolTxs'
            return (invalidTxIds, map getTxId validTxs)
          when (not (null invalidTxIds)) $
            throwIO (InvalidTxsError invalidTxIds)
          return validTxs
      }