Rewrite Mempool tests using Mock.Utxo

Using the `Mock.Utxo` instead of the overly simplistic `TestTx` we were using
allows for more complicated dependencies between transactions. This better
resembles the real transactions, at the cost of some complexity.

Moreover, using the new transaction type reveals a bug that was not visible
with the previous transaction type, i.e., #1565.

ouroboros-consensus/ouroboros-consensus.cabal

@@ -300,8 +300,6 @@ test-suite test-consensus
                     Test.Consensus.Ledger.Mock
                     Test.Consensus.LocalStateQueryServer
                     Test.Consensus.Mempool
-                    Test.Consensus.Mempool.TestBlock
-                    Test.Consensus.Mempool.TestTx
                     Test.Consensus.Node
                     Test.Consensus.Protocol.PBFT
                     Test.Consensus.ResourceRegistry

ouroboros-consensus/src/Ouroboros/Consensus/Ledger/Mock/Block.hs

@@ -249,21 +249,19 @@ instance (SimpleCrypto c, Typeable ext, SupportedBlock (SimpleBlock c ext))
       mustSucceed (Right st)  = st
   ledgerTipPoint (SimpleLedgerState st) = mockTip st
 
-updateSimpleLedgerState :: (Monad m, SimpleCrypto c, Typeable ext)
+updateSimpleLedgerState :: (SimpleCrypto c, Typeable ext)
                         => SimpleBlock c ext
                         -> LedgerState (SimpleBlock c ext)
-                        -> ExceptT (MockError (SimpleBlock c ext))
-                                   m
-                                   (LedgerState (SimpleBlock c ext))
+                        -> Except (MockError (SimpleBlock c ext))
+                                  (LedgerState (SimpleBlock c ext))
 updateSimpleLedgerState b (SimpleLedgerState st) =
     SimpleLedgerState <$> updateMockState b st
 
-updateSimpleUTxO :: (Monad m, Mock.HasUtxo a)
+updateSimpleUTxO :: Mock.HasUtxo a
                  => a
                  -> TickedLedgerState (SimpleBlock c ext)
-                 -> ExceptT (MockError (SimpleBlock c ext))
-                            m
-                            (TickedLedgerState (SimpleBlock c ext))
+                 -> Except (MockError (SimpleBlock c ext))
+                           (TickedLedgerState (SimpleBlock c ext))
 updateSimpleUTxO b (TickedLedgerState slot (SimpleLedgerState st)) =
     TickedLedgerState slot . SimpleLedgerState <$> updateMockUTxO b st
 
@@ -279,7 +277,7 @@ instance (SimpleCrypto c, Typeable ext, SupportedBlock (SimpleBlock c ext))
   data GenTx (SimpleBlock c ext) = SimpleGenTx
     { simpleGenTx   :: !Mock.Tx
     , simpleGenTxId :: !Mock.TxId
-    } deriving stock    (Generic)
+    } deriving stock    (Generic, Eq, Ord)
       deriving anyclass (Serialise)
 
   txSize = fromIntegral . Lazy.length . serialise

ouroboros-consensus/src/Ouroboros/Consensus/Ledger/Mock/Stake.hs

@@ -47,7 +47,7 @@ newtype StakeDist = StakeDist { stakeDistToMap :: Map CoreNodeId Rational }
 stakeWithDefault :: Rational -> CoreNodeId -> StakeDist -> Rational
 stakeWithDefault d n = Map.findWithDefault d n . stakeDistToMap
 
-relativeStakes :: Map StakeHolder Int -> StakeDist
+relativeStakes :: Map StakeHolder Amount -> StakeDist
 relativeStakes m = StakeDist $
    let totalStake    = fromIntegral $ sum $ Map.elems m
    in  Map.fromList [ (nid, fromIntegral stake / totalStake)
@@ -58,10 +58,10 @@ relativeStakes m = StakeDist $
 --
 -- The 'Nothing' value holds the total stake of all addresses that don't
 -- get mapped to a NodeId.
-totalStakes :: Map Addr NodeId -> Utxo -> Map StakeHolder Int
+totalStakes :: Map Addr NodeId -> Utxo -> Map StakeHolder Amount
 totalStakes addrDist = foldl f Map.empty
  where
-   f :: Map StakeHolder Int -> TxOut -> Map StakeHolder Int
+   f :: Map StakeHolder Amount -> TxOut -> Map StakeHolder Amount
    f m (a, stake) = case Map.lookup a addrDist of
        Just (CoreId nid) -> Map.insertWith (+) (StakeCore nid)    stake m
        _                 -> Map.insertWith (+) StakeEverybodyElse stake m

ouroboros-consensus/src/Ouroboros/Consensus/Ledger/Mock/State.hs

@@ -45,42 +45,42 @@ data MockState blk = MockState {
 deriving instance Serialise (HeaderHash blk) => Serialise (MockState blk)
 
 data MockError blk =
-    MockInvalidInputs InvalidInputs
+    MockUtxoError UtxoError
   | MockInvalidHash (ChainHash blk) (ChainHash blk)
   deriving (Generic, NoUnexpectedThunks)
 
 deriving instance StandardHash blk => Show (MockError blk)
 deriving instance StandardHash blk => Eq   (MockError blk)
 deriving instance Serialise (HeaderHash blk) => Serialise (MockError blk)
 
-updateMockState :: ( Monad m
-                  , GetHeader blk
-                  , HasHeader (Header blk)
-                  , StandardHash blk
-                  , HasUtxo blk
-                  )
+updateMockState :: ( GetHeader blk
+                   , HasHeader (Header blk)
+                   , StandardHash blk
+                   , HasUtxo blk
+                   )
                 => blk
                 -> MockState blk
-                -> ExceptT (MockError blk) m (MockState blk)
+                -> Except (MockError blk) (MockState blk)
 updateMockState b st = do
     st' <- updateMockTip (getHeader b) st
     updateMockUTxO b st'
 
-updateMockTip :: (Monad m, HasHeader (Header blk), StandardHash blk)
+updateMockTip :: (HasHeader (Header blk), StandardHash blk)
               => Header blk
               -> MockState blk
-              -> ExceptT (MockError blk) m (MockState blk)
-updateMockTip hdr (MockState u c t) = ExceptT $ return $
-    if headerPrevHash hdr == pointHash t
-      then Right $ MockState u c (headerPoint hdr)
-      else Left  $ MockInvalidHash (headerPrevHash hdr) (pointHash t)
+              -> Except (MockError blk) (MockState blk)
+updateMockTip hdr (MockState u c t)
+    | headerPrevHash hdr == pointHash t
+    = return $ MockState u c (headerPoint hdr)
+    | otherwise
+    = throwError $ MockInvalidHash (headerPrevHash hdr) (pointHash t)
 
-updateMockUTxO :: (Monad m, HasUtxo a)
+updateMockUTxO :: HasUtxo a
                => a
                -> MockState blk
-               -> ExceptT (MockError blk) m (MockState blk)
+               -> Except (MockError blk) (MockState blk)
 updateMockUTxO b (MockState u c t) = do
-    u' <- withExceptT MockInvalidInputs $ updateUtxo b u
+    u' <- withExcept MockUtxoError $ updateUtxo b u
     return $ MockState u' (c `Set.union` confirmed b) t
 
 {-------------------------------------------------------------------------------

ouroboros-consensus/src/Ouroboros/Consensus/Ledger/Mock/UTxO.hs

@@ -1,11 +1,9 @@
-{-# LANGUAGE DeriveAnyClass             #-}
-{-# LANGUAGE DeriveGeneric              #-}
-{-# LANGUAGE DerivingStrategies         #-}
-{-# LANGUAGE DerivingVia                #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE PatternSynonyms            #-}
-{-# LANGUAGE ScopedTypeVariables        #-}
-{-# LANGUAGE UndecidableInstances       #-}
+{-# LANGUAGE DeriveAnyClass       #-}
+{-# LANGUAGE DeriveGeneric        #-}
+{-# LANGUAGE DerivingVia          #-}
+{-# LANGUAGE PatternSynonyms      #-}
+{-# LANGUAGE ScopedTypeVariables  #-}
+{-# LANGUAGE UndecidableInstances #-}
 
 module Ouroboros.Consensus.Ledger.Mock.UTxO (
     -- * Basic definitions
@@ -14,11 +12,12 @@ module Ouroboros.Consensus.Ledger.Mock.UTxO (
   , TxIn
   , TxOut
   , Addr
+  , Amount
+  , Ix
   , Utxo
     -- * Computing UTxO
-  , InvalidInputs(..)
+  , UtxoError(..)
   , HasUtxo(..)
-  , utxo
     -- * Genesis
   , genesisTx
   , genesisUtxo
@@ -27,7 +26,8 @@ module Ouroboros.Consensus.Ledger.Mock.UTxO (
 import           Codec.Serialise (Serialise (..))
 import           Control.DeepSeq (NFData (..), force)
 import           Control.Monad.Except
-import           Data.Either (fromRight)
+import           Control.Monad.State
+import           Data.Functor (($>))
 import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import           Data.Set (Set)
@@ -67,28 +67,33 @@ instance ToCBOR Tx where
 instance Condense Tx where
   condense (Tx ins outs) = condense (ins, outs)
 
-type TxId  = Hash ShortHash Tx
-type TxIn  = (TxId, Int)
-type TxOut = (Addr, Int)
-type Utxo  = Map TxIn TxOut
+type Ix     = Word
+type Amount = Word
+type TxId   = Hash ShortHash Tx
+type TxIn   = (TxId, Ix)
+type TxOut  = (Addr, Amount)
+type Utxo   = Map TxIn TxOut
 
 {-------------------------------------------------------------------------------
   Computing UTxO
 -------------------------------------------------------------------------------}
 
-newtype InvalidInputs = InvalidInputs (Set TxIn)
-  deriving stock    (Generic)
-  deriving newtype  (Eq, Show, Condense)
+data UtxoError
+  = MissingInput TxIn
+  | InputOutputMismatch
+      Amount  -- ^ Input
+      Amount  -- ^ Output
+  deriving stock    (Eq, Show, Generic)
   deriving anyclass (Serialise, NoUnexpectedThunks)
 
+instance Condense UtxoError where
+  condense = show
+
 class HasUtxo a where
   txIns      :: a -> Set TxIn
   txOuts     :: a -> Utxo
   confirmed  :: a -> Set TxId
-  updateUtxo :: Monad m => a -> Utxo -> ExceptT InvalidInputs m Utxo
-
-utxo :: (Monad m, HasUtxo a) => a -> ExceptT InvalidInputs m Utxo
-utxo a = updateUtxo a Map.empty
+  updateUtxo :: a -> Utxo -> Except UtxoError Utxo
 
 {-------------------------------------------------------------------------------
   HasUtxo instances
@@ -97,17 +102,28 @@ utxo a = updateUtxo a Map.empty
 instance HasUtxo Tx where
   txIns     (Tx ins _outs) = ins
   txOuts tx@(Tx _ins outs) =
-      Map.fromList $ map aux (zip [0..] outs)
+      Map.fromList $ zipWith aux [0..] outs
     where
-      aux :: (Int, TxOut) -> (TxIn, TxOut)
-      aux (ix, out) = ((hash tx, ix), out)
+      aux :: Ix -> TxOut -> (TxIn, TxOut)
+      aux ix out = ((hash tx, ix), out)
 
   confirmed       = Set.singleton . hash
-  updateUtxo tx u =
-      let notInUtxo = txIns tx Set.\\ (Map.keysSet u)
-      in case Set.null notInUtxo of
-           True  -> return $ (u `Map.union` txOuts tx) `Map.withoutKeys` txIns tx
-           False -> throwError $ InvalidInputs notInUtxo
+  updateUtxo tx = execStateT $ do
+    -- Remove all inputs from the Utxo and calculate the sum of all the input
+    -- amounts
+    inputAmount <- fmap sum $ forM (Set.toList (txIns tx)) $ \txIn -> do
+      u <- get
+      case Map.updateLookupWithKey (\_ _ -> Nothing) txIn u of
+        (Nothing,              _)  -> throwError $ MissingInput txIn
+        (Just (_addr, amount), u') -> put u' $> amount
+
+    -- Check that the sum of the inputs is equal to the sum of the outputs
+    let outputAmount = sum $ map snd $ Map.elems $ txOuts tx
+    when (inputAmount /= outputAmount) $
+      throwError $ InputOutputMismatch inputAmount outputAmount
+
+    -- Add the outputs to the Utxo
+    modify (`Map.union` txOuts tx)
 
 instance HasUtxo a => HasUtxo [a] where
   txIns      = foldr (Set.union . txIns)     Set.empty
@@ -130,6 +146,4 @@ genesisTx :: AddrDist -> Tx
 genesisTx addrDist = Tx mempty [(addr, 1000) | addr <- Map.keys addrDist]
 
 genesisUtxo :: AddrDist -> Utxo
-genesisUtxo addrDist =
-    fromRight (error "genesisLedger: invalid genesis tx") $
-      runExcept (utxo (genesisTx addrDist))
+genesisUtxo addrDist = txOuts (genesisTx addrDist)

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

@@ -419,7 +419,7 @@ implSyncWithLedger mpEnv@MempoolEnv{mpEnvTracer, mpEnvStateVar} = do
       mempoolSize <- getMempoolSize mpEnv
       snapshot    <- implGetSnapshot mpEnv
       return (map fst (vrInvalid vr), mempoolSize, snapshot)
-    unless (null removed) $ do
+    unless (null removed) $
       traceWith mpEnvTracer $ TraceMempoolRemoveTxs removed mempoolSize
     return snapshot