From 9b87f102781a718044bc3de71c2952453e0d60e2 Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 16:14:58 +0000
Subject: [PATCH 1/9] Rename ResolveSlot in TimeFromSlot

---
 .../src/Cardano/Wallet/Deposit/IO/Network/Type.hs        | 9 ++++-----
 .../src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs     | 4 ++--
 .../Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs    | 6 +++---
 .../src/Cardano/Wallet/Deposit/Time.hs                   | 8 ++++++--
 .../Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs  | 6 +++---
 5 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/Network/Type.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/Network/Type.hs
index bf094edc5b2..0aa20d75c13 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/Network/Type.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO/Network/Type.hs
@@ -11,7 +11,9 @@ import Prelude
 
 import Cardano.Wallet.Deposit.Read
     ( Slot
-    , WithOrigin
+    )
+import Cardano.Wallet.Deposit.Time
+    ( LookupTimeFromSlot
     )
 import Cardano.Wallet.Network
     ( ChainFollower (..)
@@ -26,9 +28,6 @@ import Control.Tracer
 import Data.List.NonEmpty
     ( NonEmpty
     )
-import Data.Map.Strict
-    ( Map
-    )
 import Data.Set
     ( Set
     )
@@ -65,7 +64,7 @@ data NetworkEnv m block = NetworkEnv
         -- ^ Get the current 'TimeInterpreter' from the Cardano node.
     , slotsToUTCTimes
         :: Set Slot
-        -> m (Map Slot (WithOrigin UTCTime))
+        -> m LookupTimeFromSlot
     -- ^ Try to convert a set of slots to their UTCTimes counterparts
     , utcTimeToSlot
         :: UTCTime
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
index 12bb3e0f097..62656b7c40f 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
@@ -6,7 +6,7 @@ module Cardano.Wallet.Deposit.Pure.API.TxHistory
     , ByTime
     , DownTime
     , ResolveAddress
-    , ResolveSlot
+    , LookupTimeFromSlot
     , TxHistory (..)
     , firstJust
     , transfers
@@ -84,4 +84,4 @@ instance Monoid TxHistory where
     mempty = TxHistory mempty mempty
 
 type ResolveAddress = Address -> Maybe Customer
-type ResolveSlot = Slot -> Maybe DownTime
+type LookupTimeFromSlot = Slot -> Maybe DownTime
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
index 306b43bd0f7..c6cb496696b 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
@@ -19,8 +19,8 @@ import Cardano.Wallet.Deposit.Pure
     ( ValueTransfer (..)
     )
 import Cardano.Wallet.Deposit.Pure.API.TxHistory
-    ( ResolveAddress
-    , ResolveSlot
+    ( LookupTimeFromSlot
+    , ResolveAddress
     , TxHistory (..)
     )
 import Cardano.Wallet.Deposit.Read
@@ -70,7 +70,7 @@ mockTxHistory
     -- ^ Current time.
     -> ResolveAddress
     -- ^ Compute a customer from an address.
-    -> ResolveSlot
+    -> LookupTimeFromSlot
     -- ^ Compute a time from a slot.
     -> [Address]
     -- ^ List of addresses to use.
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs
index 87d97d53a84..be2019a1d54 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs
@@ -19,6 +19,7 @@ module Cardano.Wallet.Deposit.Time
     , toTimeTranslation
 
     -- * wishlist
+    , LookupTimeFromSlot
     , unsafeUTCTimeOfSlot
     , unsafeSlotsToUTCTimes
     , unsafeSlotOfUTCTime
@@ -92,9 +93,12 @@ mockSlottingParameters = SlottingParameters
 {-----------------------------------------------------------------------------
     TimeInterpreter
 ------------------------------------------------------------------------------}
-unsafeSlotsToUTCTimes :: Set.Set Slot -> Map.Map Slot (WithOrigin UTCTime)
+
+type LookupTimeFromSlot = Slot -> Maybe (WithOrigin UTCTime)
+
+unsafeSlotsToUTCTimes :: Set.Set Slot -> LookupTimeFromSlot
 unsafeSlotsToUTCTimes slots =
-    Map.fromList $ do
+    flip Map.lookup $ Map.fromList $ do
         slot <- Set.toList slots
         time <- maybeToList $ unsafeUTCTimeOfSlot slot
         pure (slot, time)
diff --git a/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs b/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs
index 0e89cc1f056..d76cc53d3ce 100644
--- a/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs
+++ b/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs
@@ -6,8 +6,8 @@ where
 import Prelude
 
 import Cardano.Wallet.Deposit.Pure.API.TxHistory
-    ( ResolveAddress
-    , ResolveSlot
+    ( LookupTimeFromSlot
+    , ResolveAddress
     , TxHistory
     )
 import Cardano.Wallet.Deposit.Pure.API.TxHistory.Mock
@@ -76,7 +76,7 @@ getMockDepositsByTimeWithCount nDeposits = do
 
 getCachedMockDeposits
     :: ResolveAddress
-    -> ResolveSlot
+    -> LookupTimeFromSlot
     -> Int
     -> [Address]
     -> IO TxHistory

From 8812ac76542d20408d452add275bf19e9dbfb2af Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 10:56:07 +0000
Subject: [PATCH 2/9] Add roll forward support in TxHistory

---
 .../Wallet/Deposit/Pure/API/TxHistory.hs      | 60 ++++++++++++-
 .../Wallet/Deposit/Pure/API/TxHistory/Mock.hs |  5 +-
 .../Cardano/Wallet/Deposit/Pure/Balance.hs    | 86 ++++++++++++++++---
 .../UI/Deposit/Handlers/Deposits/Mock.hs      |  5 +-
 4 files changed, 136 insertions(+), 20 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
index 62656b7c40f..5ba566c52d1 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE RankNTypes #-}
 
 module Cardano.Wallet.Deposit.Pure.API.TxHistory
     ( ByCustomer
@@ -10,6 +12,7 @@ module Cardano.Wallet.Deposit.Pure.API.TxHistory
     , TxHistory (..)
     , firstJust
     , transfers
+    , rollForward
     )
 where
 
@@ -19,25 +22,35 @@ import Cardano.Wallet.Deposit.Map
     ( F
     , Map (..)
     , W
+    , singletonFinger
+    , singletonMap
     )
-
 import Cardano.Wallet.Deposit.Pure.Address
     ( Customer
     )
+import Cardano.Wallet.Deposit.Pure.Balance
+    ( ValueTransferMap
+    )
 import Cardano.Wallet.Deposit.Pure.UTxO.ValueTransfer
     ( ValueTransfer
     )
 import Cardano.Wallet.Deposit.Read
     ( Address
+    , WithOrigin (..)
+    )
+import Cardano.Wallet.Deposit.Time
+    ( LookupTimeFromSlot
     )
 import Cardano.Wallet.Read
     ( Slot
     , TxId
-    , WithOrigin
     )
 import Data.Foldable
     ( Foldable (..)
     )
+import Data.Maybe
+    ( maybeToList
+    )
 import Data.Monoid
     ( First (..)
     )
@@ -48,10 +61,13 @@ import Data.Time
     ( UTCTime
     )
 
+import qualified Data.Map.Monoidal.Strict as MonoidalMap
+
 firstJust :: a -> First a
 firstJust = First . Just
 
-transfers :: Foldable (Map xs) => Map xs ValueTransfer -> ValueTransfer
+transfers
+    :: Foldable (Map xs) => Map xs ValueTransfer -> ValueTransfer
 transfers = fold
 
 type DownTime = Down (WithOrigin UTCTime)
@@ -84,4 +100,40 @@ instance Monoid TxHistory where
     mempty = TxHistory mempty mempty
 
 type ResolveAddress = Address -> Maybe Customer
-type LookupTimeFromSlot = Slot -> Maybe DownTime
+
+rollForward
+    :: ValueTransferMap
+    -> ResolveAddress
+    -> LookupTimeFromSlot
+    -> Slot
+    -> TxHistory
+    -> TxHistory
+rollForward valueTransferMap resolveAddress timeFromSlot slot =
+    (<> txHistory')
+  where
+    txHistory' =
+        blockToTxHistory valueTransferMap resolveAddress timeFromSlot slot
+
+blockToTxHistory
+    :: ValueTransferMap
+    -> ResolveAddress
+    -> LookupTimeFromSlot
+    -> Slot
+    -> TxHistory
+blockToTxHistory valueTransferMap resolveAddress timeFromSlot slot =
+    fold $ do
+        time <- fmap Down $ maybeToList $ timeFromSlot slot
+        (address, valueTransferByTxId) <- MonoidalMap.toList valueTransferMap
+        (txId, valueTransfer) <- MonoidalMap.toList valueTransferByTxId
+        customer <- maybeToList $ resolveAddress address
+        let byTime =
+                singletonFinger () time
+                    $ singletonMap (First $ Just slot) customer
+                    $ singletonMap (First $ Just address) txId
+                    $ Value valueTransfer
+        let byCustomer =
+                singletonMap () customer
+                    $ singletonFinger (First $ Just address) time
+                    $ singletonMap (First $ Just slot) txId
+                    $ Value valueTransfer
+        pure $ TxHistory{byCustomer, byTime}
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
index c6cb496696b..ee72a4cbcc6 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
@@ -46,6 +46,9 @@ import Data.Maybe
 import Data.Monoid
     ( First (..)
     )
+import Data.Ord
+    ( Down (..)
+    )
 import Data.Time
     ( UTCTime (..)
     )
@@ -103,7 +106,7 @@ mockTxHistory now solveAddress solveSlot addresses ns =
                             Just c -> c
                             Nothing -> error "fakeDepositsCreate: address not found"
                     let time = case solveSlot slot of
-                            Just t -> t
+                            Just t -> Down t
                             Nothing -> error "fakeDepositsCreate: slot not found"
                         singletonByTime =
                             singletonMap () time
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs
index 00012f05d7e..b11e134a2ad 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs
@@ -1,10 +1,13 @@
 {-# LANGUAGE BangPatterns #-}
+
 -- | Wallet balance.
 module Cardano.Wallet.Deposit.Pure.Balance
     ( balance
     , availableUTxO
     , IsOurs
     , applyBlock
+    , ValueTransferMap
+    , getTxValueTransfer
     ) where
 
 import Prelude
@@ -15,14 +18,33 @@ import Cardano.Wallet.Deposit.Pure.UTxO.DeltaUTxO
 import Cardano.Wallet.Deposit.Pure.UTxO.Tx
     ( IsOurs
     , applyTx
+    , resolveInputs
+    , valueTransferFromResolvedTx
     )
 import Cardano.Wallet.Deposit.Pure.UTxO.UTxO
     ( UTxO
     , balance
     , excluding
     )
+import Cardano.Wallet.Deposit.Pure.UTxO.ValueTransfer
+    ( ValueTransfer
+    )
+import Cardano.Wallet.Deposit.Read
+    ( Address
+    , TxId
+    )
+import Cardano.Wallet.Read
+    ( getTxId, IsEra, Block
+    )
+import Control.Monad
+    ( guard
+    )
 import Data.Foldable
-    ( foldMap'
+    ( Foldable (..)
+    , foldMap'
+    )
+import Data.Map.Monoidal.Strict
+    ( MonoidalMap
     )
 import Data.Set
     ( Set
@@ -31,10 +53,13 @@ import Data.Set
 import qualified Cardano.Wallet.Deposit.Pure.UTxO.DeltaUTxO as DeltaUTxO
 import qualified Cardano.Wallet.Deposit.Write as Write
 import qualified Cardano.Wallet.Read as Read
+import qualified Data.Map.Monoidal.Strict as MonoidalMap
+import qualified Data.Map.Strict as Map
 
 {-----------------------------------------------------------------------------
     Wallet Balance
 ------------------------------------------------------------------------------}
+
 -- | Available = excluding pending transactions
 availableUTxO :: UTxO -> Set Write.Tx -> UTxO
 availableUTxO u pending =
@@ -48,31 +73,70 @@ availableUTxO u pending =
     getUsedTxIn :: Read.Tx Read.Conway -> Set Read.TxIn
     getUsedTxIn tx =
         Read.getInputs tx
-        <> Read.getCollateralInputs tx
+            <> Read.getCollateralInputs tx
 
 {-----------------------------------------------------------------------------
     Applying Blocks
 ------------------------------------------------------------------------------}
+
+-- | Get the value transfer of a transaction.
+getTxValueTransfer
+    :: IsEra era
+    => IsOurs Address
+    -> UTxO
+    -> Read.Tx era
+    -> ValueTransferMap
+getTxValueTransfer isOurs u tx = fold $ do
+    (addr, value) <-
+        Map.assocs
+            . valueTransferFromResolvedTx
+            . resolveInputs u
+            $ tx
+    guard $ isOurs addr
+    pure
+        $ MonoidalMap.singleton addr
+        $ MonoidalMap.singleton
+            (getTxId tx)
+            value
+
+-- | A summary of all value transfers in a block.
+type ValueTransferMap = MonoidalMap Address (MonoidalMap TxId ValueTransfer)
+
 -- | Apply a 'Block' to the 'UTxO'.
 --
 -- Returns both a delta and the new value.
 applyBlock
-    :: Read.IsEra era
-    => IsOurs Read.CompactAddr -> Read.Block era -> UTxO -> (DeltaUTxO, UTxO)
+    :: IsEra era
+    => IsOurs Read.CompactAddr
+    -> Block era
+    -> UTxO
+    -> (DeltaUTxO, UTxO, ValueTransferMap)
 applyBlock isOurs block u0 =
-    (DeltaUTxO.concat $ reverse dus, u1)
+    (DeltaUTxO.concat $ reverse dus, u1, totalValueTransfer)
   where
-    (dus, u1) =
-        mapAccumL' (applyTx isOurs) u0
+    (dus, (u1, totalValueTransfer)) =
+        mapAccumL' applyTx' (u0, mempty)
             $ Read.getEraTransactions block
+    applyTx' tx (u, total) =
+        let
+            total' =
+                let value = getTxValueTransfer isOurs u tx
+                    contrib
+                        | null value = mempty
+                        | otherwise = value
+                in  total <> contrib
+            (ds, u') = applyTx isOurs tx u
+        in
+            (ds, (u', total'))
 
 {-----------------------------------------------------------------------------
     Helpers
 ------------------------------------------------------------------------------}
+
 -- | Strict variant of 'mapAccumL'.
-mapAccumL' :: (a -> s -> (o,s)) -> s -> [a] -> ([o],s)
+mapAccumL' :: (a -> s -> (o, s)) -> s -> [a] -> ([o], s)
 mapAccumL' f = go []
   where
-    go os !s0 []     = (reverse os, s0)
-    go os !s0 (x:xs) = case f x s0 of
-        (!o,!s1) -> go (o:os) s1 xs
+    go os !s0 [] = (reverse os, s0)
+    go os !s0 (x : xs) = case f x s0 of
+        (!o, !s1) -> go (o : os) s1 xs
diff --git a/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs b/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs
index d76cc53d3ce..91f5c53edff 100644
--- a/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs
+++ b/lib/ui/src/Cardano/Wallet/UI/Deposit/Handlers/Deposits/Mock.hs
@@ -34,9 +34,6 @@ import Control.Concurrent.STM
 import Control.Monad.IO.Class
     ( MonadIO (..)
     )
-import Data.Ord
-    ( Down (..)
-    )
 import Data.Time
     ( UTCTime (..)
     , diffUTCTime
@@ -66,7 +63,7 @@ getMockDepositsByTimeWithCount
 getMockDepositsByTimeWithCount nDeposits = do
     addresses <- fmap snd <$> listCustomers
     solveAddress <- addressToCustomer
-    let solveSlot = fmap Down <$> unsafeUTCTimeOfSlot
+    let solveSlot = unsafeUTCTimeOfSlot
     liftIO
         $ getCachedMockDeposits
             solveAddress

From 608635f03d886fd181c30f15a4f2ced56c762616 Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 16:11:45 +0000
Subject: [PATCH 3/9] Include TxHistory roll forward in deposit wallet state
 roll forward

---
 .../src/Cardano/Wallet/Deposit/IO.hs          | 35 ++++++++++--
 .../src/Cardano/Wallet/Deposit/Pure.hs        | 53 +++++++++++++++----
 .../src/Cardano/Wallet/Deposit/Time.hs        |  1 +
 .../unit/Cardano/Wallet/Deposit/PureSpec.hs   | 19 +++++--
 4 files changed, 87 insertions(+), 21 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs
index 9ed07f8c392..b0af90bc82b 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs
@@ -46,12 +46,16 @@ import Cardano.Crypto.Wallet
 import Cardano.Wallet.Address.BIP32
     ( BIP32Path
     )
+import Cardano.Wallet.Deposit.IO.Network.Type
+    ( NetworkEnv (slotsToUTCTimes)
+    )
 import Cardano.Wallet.Deposit.Pure
     ( Customer
     , ValueTransfer
     , WalletPublicIdentity (..)
     , WalletState
     , Word31
+    , getEraSlotOfBlock
     )
 import Cardano.Wallet.Deposit.Pure.API.TxHistory
     ( ByCustomer
@@ -65,6 +69,9 @@ import Cardano.Wallet.Deposit.Read
 import Cardano.Wallet.Network.Checkpoints.Policy
     ( defaultPolicy
     )
+import Cardano.Wallet.Read
+    ( applyEraFun
+    )
 import Control.Tracer
     ( Tracer
     , contramap
@@ -72,6 +79,9 @@ import Control.Tracer
 import Data.Bifunctor
     ( first
     )
+import Data.Foldable
+    ( Foldable (..)
+    )
 import Data.List.NonEmpty
     ( NonEmpty
     )
@@ -92,6 +102,7 @@ import qualified Data.Delta as Delta
     ( Replace (..)
     )
 import qualified Data.Delta.Update as Delta
+import qualified Data.Set as Set
 import qualified Data.Store as Store
 
 {-----------------------------------------------------------------------------
@@ -113,8 +124,7 @@ data WalletBootEnv m = WalletBootEnv
 type WalletStore = Store.UpdateStore IO Wallet.DeltaWalletState
 
 -- | The full environment needed to run a wallet.
-data WalletEnv m
-    = WalletEnv
+data WalletEnv m = WalletEnv
     { bootEnv :: WalletBootEnv m
     -- ^ The boot environment.
     , store :: WalletStore
@@ -269,11 +279,26 @@ getAllDeposits w i =
     Wallet.getAllDeposits i <$> readWalletState w
 
 rollForward
-    :: WalletInstance -> NonEmpty (Read.EraValue Read.Block) -> tip -> IO ()
-rollForward w blocks _nodeTip =
+    :: WalletInstance
+    -> NonEmpty (Read.EraValue Read.Block)
+    -> tip
+    -> IO ()
+rollForward w blocks _nodeTip = do
+    let blockSlots =
+            Set.fromList
+                $ applyEraFun getEraSlotOfBlock
+                    <$> toList blocks
+    timeFromSlot <-
+        flip slotsToUTCTimes blockSlots
+            $ networkEnv
+            $ bootEnv
+            $ env w
     onWalletState w
         $ Delta.update
-        $ Delta.Replace . Wallet.rollForwardMany blocks
+        $ Delta.Replace
+            . Wallet.rollForwardMany
+                timeFromSlot
+                blocks
 
 rollBackward
     :: WalletInstance -> Read.ChainPoint -> IO Read.ChainPoint
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs
index 1799b42bed9..3c3d7af25d5 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs
@@ -36,6 +36,7 @@ module Cardano.Wallet.Deposit.Pure
     , ValueTransfer (..)
     , getTxHistoryByCustomer
     , getTxHistoryByTime
+    , getEraSlotOfBlock
 
       -- ** Writing to the blockchain
     , createPayment
@@ -82,6 +83,9 @@ import Cardano.Wallet.Deposit.Pure.API.TxHistory
     , DownTime
     , TxHistory (..)
     )
+import Cardano.Wallet.Deposit.Pure.Balance
+    ( ValueTransferMap
+    )
 import Cardano.Wallet.Deposit.Pure.UTxO.UTxOHistory
     ( UTxOHistory
     )
@@ -90,8 +94,14 @@ import Cardano.Wallet.Deposit.Pure.UTxO.ValueTransfer
     )
 import Cardano.Wallet.Deposit.Read
     ( Address
+    , Slot
     , TxId
-    , WithOrigin
+    , WithOrigin (..)
+    , getEraBHeader
+    , getEraSlotNo
+    )
+import Cardano.Wallet.Deposit.Time
+    ( LookupTimeFromSlot
     )
 import Data.Foldable
     ( fold
@@ -120,6 +130,7 @@ import Data.Word.Odd
     )
 
 import qualified Cardano.Wallet.Deposit.Pure.Address as Address
+import qualified Cardano.Wallet.Deposit.Pure.API.TxHistory as TxHistory
 import qualified Cardano.Wallet.Deposit.Pure.Balance as Balance
 import qualified Cardano.Wallet.Deposit.Pure.RollbackWindow as Rollback
 import qualified Cardano.Wallet.Deposit.Pure.Submissions as Sbm
@@ -230,31 +241,51 @@ getWalletTip :: WalletState -> Read.ChainPoint
 getWalletTip = walletTip
 
 rollForwardMany
-    :: NonEmpty (Read.EraValue Read.Block) -> WalletState -> WalletState
-rollForwardMany blocks w = foldl' (flip rollForwardOne) w blocks
+    :: LookupTimeFromSlot
+    -> NonEmpty (Read.EraValue Read.Block)
+    -> WalletState
+    -> WalletState
+rollForwardMany timeFromSlot blocks w =
+    foldl' (flip $ rollForwardOne timeFromSlot) w blocks
 
 rollForwardOne
-    :: Read.EraValue Read.Block -> WalletState -> WalletState
-rollForwardOne (Read.EraValue block) w =
+    :: LookupTimeFromSlot
+    -> Read.EraValue Read.Block
+    -> WalletState
+    -> WalletState
+rollForwardOne timeFromSlot (Read.EraValue block) w =
     w
         { walletTip = Read.getChainPoint block
-        , utxoHistory = rollForwardUTxO isOurs block (utxoHistory w)
+        , utxoHistory = utxoHistory'
         , submissions = Delta.apply (Sbm.rollForward block) (submissions w)
+        , txHistory =
+            TxHistory.rollForward
+                valueTransfers
+                (`addressToCustomer` w)
+                timeFromSlot
+                (getEraSlotOfBlock block)
+                (txHistory w)
         }
   where
+    (utxoHistory', valueTransfers) =
+        rollForwardUTxO isOurs block (utxoHistory w)
     isOurs :: Address -> Bool
     isOurs = Address.isOurs (addresses w)
 
+getEraSlotOfBlock :: Read.IsEra era => Read.Block era -> Slot
+getEraSlotOfBlock = At . getEraSlotNo . getEraBHeader
+
 rollForwardUTxO
     :: Read.IsEra era
     => (Address -> Bool)
     -> Read.Block era
     -> UTxOHistory
-    -> UTxOHistory
+    -> (UTxOHistory, ValueTransferMap)
 rollForwardUTxO isOurs block u =
-    UTxOHistory.rollForward slot deltaUTxO u
+    (UTxOHistory.rollForward slot deltaUTxO u, valueTransfers)
   where
-    (deltaUTxO, _) = Balance.applyBlock isOurs block (UTxOHistory.getUTxO u)
+    (deltaUTxO, _, valueTransfers) =
+            Balance.applyBlock isOurs block (UTxOHistory.getUTxO u)
     slot = Read.getEraSlotNo $ Read.getEraBHeader block
 
 rollBackward
@@ -282,8 +313,8 @@ rollBackward targetPoint w =
     -- any other point than the target point (or genesis).
     actualPoint =
         if (targetSlot `Rollback.member` UTxOHistory.getRollbackWindow h)
-            -- FIXME: Add test for rollback window of `submissions`
-            then targetPoint
+            then -- FIXME: Add test for rollback window of `submissions`
+                targetPoint
             else Read.GenesisPoint
 
 availableBalance :: WalletState -> Read.Value
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs
index be2019a1d54..f3888c7f71d 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Time.hs
@@ -25,6 +25,7 @@ module Cardano.Wallet.Deposit.Time
     , unsafeSlotOfUTCTime
     , systemStartMainnet
     , originTime
+
     ) where
 
 import Prelude
diff --git a/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs b/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs
index 369703d9e4b..85b990854fc 100644
--- a/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs
+++ b/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs
@@ -17,6 +17,12 @@ import Cardano.Crypto.Wallet
     , generate
     , toXPub
     )
+import Cardano.Wallet.Deposit.Pure.API.TxHistory
+    ( LookupTimeFromSlot
+    )
+import Cardano.Wallet.Deposit.Time
+    ( unsafeUTCTimeOfSlot
+    )
 import Test.Hspec
     ( Spec
     , describe
@@ -37,6 +43,9 @@ import qualified Data.ByteString.Char8 as B8
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
 
+timeFromSlot :: LookupTimeFromSlot
+timeFromSlot = unsafeUTCTimeOfSlot
+
 spec :: Spec
 spec = do
     describe "UTxO availableBalance" $ do
@@ -59,11 +68,11 @@ prop_availableBalance_rollForward_twice =
     tx1 = payFromFaucet [(addr1, Write.mkAda 1)]
     block1 = Read.mockNextBlock Read.GenesisPoint [tx1]
     chainPoint1 = Read.getChainPoint block1
-    w1 = Wallet.rollForwardOne (Read.EraValue block1) w0
+    w1 = Wallet.rollForwardOne  timeFromSlot (Read.EraValue block1) w0
 
     tx2 = payFromFaucet [(addr2, Write.mkAda 2)]
     block2 = Read.mockNextBlock chainPoint1 [tx2]
-    w2 = Wallet.rollForwardOne (Read.EraValue block2) w1
+    w2 = Wallet.rollForwardOne timeFromSlot (Read.EraValue block2) w1
 
 prop_availableBalance_rollForward_rollBackward :: Property
 prop_availableBalance_rollForward_rollBackward =
@@ -89,17 +98,17 @@ prop_availableBalance_rollForward_rollBackward =
 
     tx1 = payFromFaucet [(addr1, Write.mkAda 1)]
     block1 = Read.mockNextBlock chainPoint0 [tx1]
-    w1 = Wallet.rollForwardOne (Read.EraValue block1) w0
+    w1 = Wallet.rollForwardOne timeFromSlot (Read.EraValue block1) w0
     chainPoint1 = Read.getChainPoint block1
 
     tx2 = payFromFaucet [(addr2, Write.mkAda 2)]
     block2 = Read.mockNextBlock chainPoint1 [tx2]
     chainPoint2 = Read.getChainPoint block2
-    w2 = Wallet.rollForwardOne (Read.EraValue block2) w1
+    w2 = Wallet.rollForwardOne timeFromSlot (Read.EraValue block2) w1
 
     tx3 = spendOneTxOut (Wallet.availableUTxO w2)
     block3 = Read.mockNextBlock chainPoint2 [tx3]
-    w3 = Wallet.rollForwardOne (Read.EraValue block3) w2
+    w3 = Wallet.rollForwardOne timeFromSlot (Read.EraValue block3) w2
 
 emptyWalletWith17Addresses :: Wallet.WalletState
 emptyWalletWith17Addresses =

From d1df83ad5ec4fecbcb6ed88c6041c3fc9fc3a4b0 Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 14:43:54 +0000
Subject: [PATCH 4/9] Add map shallow modifiers for the depoist map datatype

---
 .../src/Cardano/Wallet/Deposit/Map.hs         | 50 ++++++++++++++++---
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs
index 46eb36d2102..1e5f3e1c478 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs
@@ -39,7 +39,13 @@ module Cardano.Wallet.Deposit.Map
       -- * Construction
     , singletonMap
     , singletonFinger
+
+      -- * Conversion
     , toFinger
+
+      -- * Modification
+    , onMap
+    , onFinger
     )
 where
 
@@ -82,12 +88,20 @@ data F (w :: Type) (k :: Type)
 
 -- | A nested monoidal map. Every nesting can also be patched with a monoid 'w'.
 data Map :: [Type] -> Type -> Type where
-    Value :: v -> Map '[] v
-    -- ^ A leaf node with a value.
-    Map :: w -> k ^^^ Map ks v -> Map (W w k ': ks) v
-    -- ^ A node with a patch 'w' and a nested monoidal map.
-    Finger :: w -> TimedSeq k (Map ks v) -> Map (F w k ': ks) v
-    -- ^ A node with a patch 'w' and a nested finger tree of maps.
+    Value
+        :: v
+        -> Map '[] v
+        -- ^ A leaf node with a value.
+    Map
+        :: w
+        -> k ^^^ Map ks v
+        -> Map (W w k ': ks) v
+        -- ^ A node with a patch 'w' and a nested monoidal map.
+    Finger
+        :: w
+        -> TimedSeq k (Map ks v)
+        -> Map (F w k ': ks) v
+        -- ^ A node with a patch 'w' and a nested finger tree of maps.
 
 deriving instance Show v => Show (Map '[] v)
 
@@ -114,6 +128,13 @@ deriving instance
     )
     => Show (Map (F w k ': ks) v)
 
+deriving instance
+    ( Eq w
+    , Eq k
+    , Eq (Map ks v)
+    )
+    => Eq (Map (F w k ': ks) v)
+
 instance Functor (Map '[]) where
     fmap f (Value v) = Value (f v)
 
@@ -241,7 +262,8 @@ singletonFinger
 singletonFinger w k m =
     Finger w $ FingerTree.singleton (Timed (Last (Just k)) m)
 
-toFinger :: Monoid (Map ks a) => Map (W w k : ks) a -> Map (F w k : ks) a
+toFinger
+    :: Monoid (Map ks a) => Map (W w k : ks) a -> Map (F w k : ks) a
 toFinger (Map w m) = Finger w $ FingerTree.fromList $ do
     (k, v) <- MonoidalMap.toList m
     pure $ Timed (Last (Just k)) v
@@ -262,3 +284,17 @@ lookupFinger k1 k2 (Finger w m) = do
     case extractInterval k1 k2 m of
         Timed (Last Nothing) _ -> Nothing
         Timed _ m' -> Just (w, m')
+
+-- | Apply a function to the nested monoidal map keeping the patch.
+onMap
+    :: Map (W w k : ks) a
+    -> (MonoidalMap k (Map ks a) -> MonoidalMap k (Map ks a))
+    -> Map (W w k : ks) a
+onMap (Map w m) f = Map w $ f m
+
+-- | Apply a function to the nested finger tree keeping the patch.
+onFinger
+    :: Map (F w k : ks) a
+    -> (TimedSeq k (Map ks a) -> TimedSeq k (Map ks a))
+    -> Map (F w k : ks) a
+onFinger (Finger w m) f = Finger w $ f m

From 1b0d5dabdea99249ce30060c3add0b0a16d8bd76 Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 14:45:49 +0000
Subject: [PATCH 5/9] Add dropAfter splitting function to for TimeSeq values

---
 .../src/Cardano/Wallet/Deposit/Map/Timed.hs   |  5 +++++
 .../Cardano/Wallet/Deposit/Map/TimedSpec.hs   | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs
index 128bf787ee1..424d0a646b5 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs
@@ -11,6 +11,7 @@ module Cardano.Wallet.Deposit.Map.Timed
     , extractInterval
     , minKey
     , maxKey
+    , dropAfter
     )
 where
 
@@ -183,3 +184,7 @@ extractInterval t0 t1 tseq =
     measure
         $ takeUntil (\q -> time q > Last (Just t1))
         $ dropUntil (\q -> time q >= Last (Just t0)) tseq
+
+-- | Extract all elements from a tseq that are before the given time.
+dropAfter :: (Ord t, Monoid a) => t -> TimedSeq t a -> TimedSeq t a
+dropAfter t = takeUntil (\q -> time q > Last (Just t))
diff --git a/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/Map/TimedSpec.hs b/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/Map/TimedSpec.hs
index 3918bab588e..2916c0f8f88 100644
--- a/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/Map/TimedSpec.hs
+++ b/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/Map/TimedSpec.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE OverloadedLists #-}
 {-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE TypeApplications #-}
 
 module Cardano.Wallet.Deposit.Map.TimedSpec where
 
@@ -10,6 +11,7 @@ import Prelude
 import Cardano.Wallet.Deposit.Map.Timed
     ( Timed (..)
     , TimedSeq
+    , dropAfter
     , maxKey
     , minKey
     , takeAfter
@@ -358,3 +360,20 @@ spec = do
                     [ [t4 <> t5 <> t6 <> t7]
                     , [t0 <> t1 <> t2 <> t3]
                     ]
+
+    describe "dropAfter function" $ do
+        it "works on empty" $ do
+            dropAfter (t "2021-01-01 00:00:00") (fromList @UTimed [])
+                `shouldBe` fromList []
+        it "drop a single" $ do
+            dropAfter (t "2021-01-01 00:00:00") (fromList @UTimed [t0])
+                `shouldBe` fromList [t0]
+        it "take one and drop the second, early cut" $ do
+            dropAfter (t "2021-01-01 00:00:00") (fromList @UTimed [t0, t1])
+                `shouldBe` fromList [t0]
+        it "take one and drop the second, late cut" $ do
+            dropAfter (t "2021-01-01 23:59:59") (fromList @UTimed [t0, t1])
+                `shouldBe` fromList [t0]
+        it "can take all" $ do
+            dropAfter (t "2021-01-02 00:00:00") (fromList @UTimed [t0, t1])
+                `shouldBe` fromList [t0, t1]

From 267105ff6f7350d9f4b4d516acc99732449c811d Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Sat, 2 Nov 2024 11:26:33 +0000
Subject: [PATCH 6/9] Add IsList instance to TimedSeq type

---
 .../src/Cardano/Wallet/Deposit/Map/Timed.hs         | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs
index 424d0a646b5..11b40997573 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map/Timed.hs
@@ -1,6 +1,8 @@
 {-# LANGUAGE DeriveFoldable #-}
 {-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE TypeFamilies #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Cardano.Wallet.Deposit.Map.Timed
@@ -40,6 +42,12 @@ import Data.Function
 import Data.Monoid
     ( Last (..)
     )
+import GHC.IsList
+    ( IsList (..)
+    )
+
+import qualified Data.FingerTree as FingerTree
+import qualified Data.Foldable as F
 
 -- | A value paired with a timestamp.
 data Timed t a = Timed
@@ -60,6 +68,11 @@ instance Monoid a => Measured (Timed t a) (Timed t a) where
 -- | A sequence of timed values with a monoidal annotation as itself
 type TimedSeq t a = FingerTree (Timed t a) (Timed t a)
 
+instance Monoid a => IsList (TimedSeq t a) where
+    type Item (TimedSeq t a) = Timed t a
+    fromList = FingerTree.fromList
+    toList = F.toList
+
 takeAfterElement
     :: (Monoid a, Ord q)
     => (t -> q)

From f9e53debd9d11b088e3cc42db903048dbf2b1fc8 Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 14:58:38 +0000
Subject: [PATCH 7/9] Add roll backward support in TxHistory

---
 .../Wallet/Deposit/Pure/API/TxHistory.hs      | 36 ++++++++++++++++---
 .../Wallet/Deposit/Pure/API/TxHistory/Mock.hs | 12 +++----
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
index 5ba566c52d1..d980b1706b0 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 
@@ -13,6 +14,7 @@ module Cardano.Wallet.Deposit.Pure.API.TxHistory
     , firstJust
     , transfers
     , rollForward
+    , rollBackward
     )
 where
 
@@ -22,9 +24,15 @@ import Cardano.Wallet.Deposit.Map
     ( F
     , Map (..)
     , W
+    , onFinger
+    , onMap
     , singletonFinger
     , singletonMap
     )
+import Cardano.Wallet.Deposit.Map.Timed
+    ( TimedSeq
+    , dropAfter
+    )
 import Cardano.Wallet.Deposit.Pure.Address
     ( Customer
     )
@@ -101,6 +109,9 @@ instance Monoid TxHistory where
 
 type ResolveAddress = Address -> Maybe Customer
 
+-- | Roll forward a block into the transaction history. This function
+-- relies on the TxHistory to be sorted by time both on the time and
+-- customer views.
 rollForward
     :: ValueTransferMap
     -> ResolveAddress
@@ -128,12 +139,29 @@ blockToTxHistory valueTransferMap resolveAddress timeFromSlot slot =
         customer <- maybeToList $ resolveAddress address
         let byTime =
                 singletonFinger () time
-                    $ singletonMap (First $ Just slot) customer
-                    $ singletonMap (First $ Just address) txId
+                    $ singletonMap (firstJust slot) customer
+                    $ singletonMap (firstJust address) txId
                     $ Value valueTransfer
         let byCustomer =
                 singletonMap () customer
-                    $ singletonFinger (First $ Just address) time
-                    $ singletonMap (First $ Just slot) txId
+                    $ singletonFinger (firstJust address) time
+                    $ singletonMap (firstJust slot) txId
                     $ Value valueTransfer
         pure $ TxHistory{byCustomer, byTime}
+
+-- | Roll backward the transaction history to a given slot. This function
+-- relies on the TxHistory to be sorted by time both on the time and
+-- customer views.
+rollBackward
+    :: LookupTimeFromSlot
+    -> Slot
+    -> TxHistory
+    -> TxHistory
+rollBackward timeFromSlot slot TxHistory{byCustomer, byTime} =
+    TxHistory
+        { byCustomer = onMap byCustomer $ fmap (`onFinger` takeToSlot)
+        , byTime = onFinger byTime takeToSlot
+        }
+  where
+    takeToSlot :: Monoid a => TimedSeq DownTime a -> TimedSeq DownTime a
+    takeToSlot x = maybe x ((`dropAfter` x) . Down) $ timeFromSlot slot
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
index ee72a4cbcc6..f93ddc76b69 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/API/TxHistory/Mock.hs
@@ -22,6 +22,7 @@ import Cardano.Wallet.Deposit.Pure.API.TxHistory
     ( LookupTimeFromSlot
     , ResolveAddress
     , TxHistory (..)
+    , firstJust
     )
 import Cardano.Wallet.Deposit.Read
     ( Address
@@ -43,9 +44,6 @@ import Control.Monad
 import Data.Maybe
     ( fromJust
     )
-import Data.Monoid
-    ( First (..)
-    )
 import Data.Ord
     ( Down (..)
     )
@@ -110,13 +108,13 @@ mockTxHistory now solveAddress solveSlot addresses ns =
                             Nothing -> error "fakeDepositsCreate: slot not found"
                         singletonByTime =
                             singletonMap () time
-                                $ singletonMap (First $ Just slot) customer
-                                $ singletonMap (First $ Just address) txId
+                                $ singletonMap (firstJust slot) customer
+                                $ singletonMap (firstJust address) txId
                                 $ Value value
                         singletonByCustomer =
                             singletonMap () customer
-                                $ singletonMap (First $ Just address) time
-                                $ singletonMap (First $ Just slot) txId
+                                $ singletonMap (firstJust address) time
+                                $ singletonMap (firstJust slot) txId
                                 $ Value value
                     pure (singletonByCustomer, singletonByTime)
         byCustomer' = Map w $ fmap toFinger f

From f417f15ccb0e7ca24baa3092670ad29f8350da2c Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Fri, 1 Nov 2024 15:11:05 +0000
Subject: [PATCH 8/9] Include TxHistory roll backward in deposit wallet state
 roll back

---
 .../src/Cardano/Wallet/Deposit/IO.hs          | 32 ++++++++++++-------
 .../src/Cardano/Wallet/Deposit/Pure.hs        | 10 ++++--
 .../unit/Cardano/Wallet/Deposit/PureSpec.hs   | 11 ++++---
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs
index b0af90bc82b..a30fac1ff3f 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/IO.hs
@@ -60,6 +60,7 @@ import Cardano.Wallet.Deposit.Pure
 import Cardano.Wallet.Deposit.Pure.API.TxHistory
     ( ByCustomer
     , ByTime
+    , LookupTimeFromSlot
     )
 import Cardano.Wallet.Deposit.Read
     ( Address
@@ -70,7 +71,8 @@ import Cardano.Wallet.Network.Checkpoints.Policy
     ( defaultPolicy
     )
 import Cardano.Wallet.Read
-    ( applyEraFun
+    ( Slot
+    , applyEraFun
     )
 import Control.Tracer
     ( Tracer
@@ -285,14 +287,9 @@ rollForward
     -> IO ()
 rollForward w blocks _nodeTip = do
     let blockSlots =
-            Set.fromList
-                $ applyEraFun getEraSlotOfBlock
-                    <$> toList blocks
-    timeFromSlot <-
-        flip slotsToUTCTimes blockSlots
-            $ networkEnv
-            $ bootEnv
-            $ env w
+            applyEraFun getEraSlotOfBlock
+                <$> toList blocks
+    timeFromSlot <- slotResolver w blockSlots
     onWalletState w
         $ Delta.update
         $ Delta.Replace
@@ -302,10 +299,23 @@ rollForward w blocks _nodeTip = do
 
 rollBackward
     :: WalletInstance -> Read.ChainPoint -> IO Read.ChainPoint
-rollBackward w point =
+rollBackward w point = do
+    timeFromSlot <- slotResolver w [Read.slotFromChainPoint point]
     onWalletState w
         $ Delta.updateWithResult
-        $ first Delta.Replace . Wallet.rollBackward point
+        $ first Delta.Replace . Wallet.rollBackward timeFromSlot point
+
+-- | Compute a slot resolver for the given slots.
+slotResolver
+    :: WalletInstance
+    -> [Slot]
+    -> IO LookupTimeFromSlot
+slotResolver w slots = do
+    let slotSet = Set.fromList slots
+    flip slotsToUTCTimes slotSet
+        $ networkEnv
+        $ bootEnv
+        $ env w
 
 {-----------------------------------------------------------------------------
     Operations
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs
index 3c3d7af25d5..026036ff822 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure.hs
@@ -289,16 +289,22 @@ rollForwardUTxO isOurs block u =
     slot = Read.getEraSlotNo $ Read.getEraBHeader block
 
 rollBackward
-    :: Read.ChainPoint
+    :: LookupTimeFromSlot
+    -> Read.ChainPoint
     -> WalletState
     -> (WalletState, Read.ChainPoint)
-rollBackward targetPoint w =
+rollBackward timeFromSlot targetPoint w =
     ( w
         { walletTip = actualPoint
         , utxoHistory =
             UTxOHistory.rollBackward actualSlot (utxoHistory w)
         , submissions =
             Delta.apply (Sbm.rollBackward actualSlot) (submissions w)
+        , txHistory =
+            TxHistory.rollBackward
+                timeFromSlot
+                (Read.slotFromChainPoint targetPoint)
+                (txHistory w)
         }
     , actualPoint
     )
diff --git a/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs b/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs
index 85b990854fc..2859130d2c1 100644
--- a/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs
+++ b/lib/customer-deposit-wallet/test/unit/Cardano/Wallet/Deposit/PureSpec.hs
@@ -68,7 +68,7 @@ prop_availableBalance_rollForward_twice =
     tx1 = payFromFaucet [(addr1, Write.mkAda 1)]
     block1 = Read.mockNextBlock Read.GenesisPoint [tx1]
     chainPoint1 = Read.getChainPoint block1
-    w1 = Wallet.rollForwardOne  timeFromSlot (Read.EraValue block1) w0
+    w1 = Wallet.rollForwardOne timeFromSlot (Read.EraValue block1) w0
 
     tx2 = payFromFaucet [(addr2, Write.mkAda 2)]
     block2 = Read.mockNextBlock chainPoint1 [tx2]
@@ -76,13 +76,16 @@ prop_availableBalance_rollForward_twice =
 
 prop_availableBalance_rollForward_rollBackward :: Property
 prop_availableBalance_rollForward_rollBackward =
-        Wallet.availableBalance (fst $ Wallet.rollBackward chainPoint0 w3)
+        Wallet.availableBalance
+            (fst $ Wallet.rollBackward timeFromSlot chainPoint0 w3)
             === Wallet.availableBalance w0
     .&&.
-        Wallet.availableBalance (fst $ Wallet.rollBackward chainPoint1 w3)
+        Wallet.availableBalance
+            (fst $ Wallet.rollBackward timeFromSlot chainPoint1 w3)
             === Wallet.availableBalance w1
     .&&.
-        Wallet.availableBalance (fst $ Wallet.rollBackward chainPoint2 w3)
+        Wallet.availableBalance
+            (fst $ Wallet.rollBackward timeFromSlot chainPoint2 w3)
             === Wallet.availableBalance w2
     .&&.
         Wallet.availableBalance w3

From ae7c59c0a902787ba431b5c285d1cbfb53210fb9 Mon Sep 17 00:00:00 2001
From: paolino <paolo.veronelli@gmail.com>
Date: Mon, 4 Nov 2024 17:07:54 +0000
Subject: [PATCH 9/9] Add a dsl to define deposit pure tests

---
 .../customer-deposit-wallet.cabal             |   8 +-
 .../src/Cardano/Wallet/Deposit/Map.hs         |   6 +
 .../Cardano/Wallet/Deposit/Pure/Balance.hs    |   4 +-
 .../src/Cardano/Wallet/Deposit/Testing/DSL.hs | 205 ++++++++++++++++
 .../Wallet/Deposit/Testing/DSL/ByTime.hs      | 218 ++++++++++++++++++
 .../Wallet/Deposit/Testing/DSL/Types.hs       |  24 ++
 .../src/Cardano/Wallet/Deposit/Write.hs       |  53 ++++-
 7 files changed, 503 insertions(+), 15 deletions(-)
 create mode 100644 lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL.hs
 create mode 100644 lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/ByTime.hs
 create mode 100644 lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/Types.hs

diff --git a/lib/customer-deposit-wallet/customer-deposit-wallet.cabal b/lib/customer-deposit-wallet/customer-deposit-wallet.cabal
index 8a923a85c93..81c65b6ef57 100644
--- a/lib/customer-deposit-wallet/customer-deposit-wallet.cabal
+++ b/lib/customer-deposit-wallet/customer-deposit-wallet.cabal
@@ -69,11 +69,12 @@ library
     , delta-types
     , fingertree
     , io-classes
-    , microlens
+    , lens
     , monoidal-containers
     , mtl
     , mwc-random
     , OddWord
+    , operational
     , random
     , text
     , time
@@ -94,6 +95,9 @@ library
     Cardano.Wallet.Deposit.Pure.Submissions
     Cardano.Wallet.Deposit.Pure.UTxO
     Cardano.Wallet.Deposit.Read
+    Cardano.Wallet.Deposit.Testing.DSL
+    Cardano.Wallet.Deposit.Testing.DSL.ByTime
+    Cardano.Wallet.Deposit.Testing.DSL.Types
     Cardano.Wallet.Deposit.Time
     Cardano.Wallet.Deposit.Write
 
@@ -113,7 +117,7 @@ test-suite scenario
     , contra-tracer
     , customer-deposit-wallet
     , delta-store
-    , hspec                      >=2.8.2
+    , hspec
 
   other-modules:
     Test.Scenario.Blockchain
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs
index 1e5f3e1c478..a7bb9c3612c 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Map.hs
@@ -46,6 +46,7 @@ module Cardano.Wallet.Deposit.Map
       -- * Modification
     , onMap
     , onFinger
+    , Peel
     )
 where
 
@@ -298,3 +299,8 @@ onFinger
     -> (TimedSeq k (Map ks a) -> TimedSeq k (Map ks a))
     -> Map (F w k : ks) a
 onFinger (Finger w m) f = Finger w $ f m
+
+type family Peel x where
+    Peel (Map (W w k : xs) v) =  Map xs v
+    Peel (Map (F w k : xs) v) =  Map xs v
+    Peel (Map '[] v) = v
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs
index b11e134a2ad..99a0e440f99 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Pure/Balance.hs
@@ -34,7 +34,9 @@ import Cardano.Wallet.Deposit.Read
     , TxId
     )
 import Cardano.Wallet.Read
-    ( getTxId, IsEra, Block
+    ( Block
+    , IsEra
+    , getTxId
     )
 import Control.Monad
     ( guard
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL.hs
new file mode 100644
index 00000000000..fa537619e1e
--- /dev/null
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL.hs
@@ -0,0 +1,205 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Cardano.Wallet.Deposit.Testing.DSL where
+
+import Prelude
+
+import Cardano.Wallet.Deposit.Pure
+    ( Customer
+    , WalletState
+    )
+import Cardano.Wallet.Deposit.Pure.API.TxHistory
+    ( ByTime
+    )
+import Cardano.Wallet.Deposit.Read
+    ( Address
+    )
+import Cardano.Wallet.Deposit.Testing.DSL.ByTime
+    ( ByTimeM
+    , at
+    , deposited
+    , forCustomer
+    , inTx
+    , newTimes
+    , withdrawn
+    )
+import Cardano.Wallet.Deposit.Testing.DSL.Types
+    ( BlockI (..)
+    , TxI (..)
+    , UnspentI (..)
+    )
+import Cardano.Wallet.Deposit.Write
+    ( TxBody
+    , addTxIn
+    , addTxOut
+    , emptyTxBody
+    , mkAda
+    , mkTx
+    , mkTxOut
+    )
+import Cardano.Wallet.Read
+    ( Slot
+    , WithOrigin (At, Origin)
+    , getTxId
+    , pattern TxIn
+    )
+import Control.Lens
+    ( _2
+    , zoom
+    )
+import Control.Monad
+    ( (>=>)
+    )
+import Control.Monad.Operational
+    ( ProgramT
+    , ProgramViewT (..)
+    , singleton
+    , viewT
+    )
+import Control.Monad.State
+    ( MonadState (..)
+    , MonadTrans (..)
+    , StateT
+    , execStateT
+    )
+import Data.Map
+    ( Map
+    )
+import Data.Time
+    ( UTCTime
+    )
+
+import qualified Data.Map as Map
+
+data Scenario p a where
+    ExistsTx :: Scenario p TxI
+    Deposit :: TxI -> Customer -> Int -> Scenario p UnspentI
+    Withdrawal :: TxI -> UnspentI -> Scenario p ()
+    CreateBlock :: Slot -> [TxI] -> Scenario p BlockI
+    RollForward :: [BlockI] -> Scenario p ()
+    RollBackward :: Slot -> Scenario p ()
+    HistoryByTime :: Scenario p ByTime
+    NewHistoryByTime :: ByTimeM m ByTime -> Scenario p ByTime
+    Assert :: p -> Scenario p ()
+
+type ScenarioP p m = ProgramT (Scenario p) m
+
+existsTx :: ScenarioP p m TxI
+existsTx = singleton ExistsTx
+
+deposit :: TxI -> Customer -> Int -> ScenarioP p m UnspentI
+deposit tx customer value = singleton (Deposit tx customer value)
+
+withdrawal :: TxI -> UnspentI -> ScenarioP p m ()
+withdrawal tx unspent = singleton (Withdrawal tx unspent)
+
+block :: Slot -> [TxI] -> ScenarioP p m BlockI
+block slot txs = singleton (CreateBlock slot txs)
+
+rollForward :: [BlockI] -> ScenarioP p m ()
+rollForward blocks = singleton (RollForward blocks)
+
+rollBackward :: Slot -> ScenarioP p m ()
+rollBackward slot = singleton (RollBackward slot)
+
+historyByTime :: ScenarioP p m ByTime
+historyByTime = singleton HistoryByTime
+
+newHistoryByTime :: ByTimeM m ByTime -> ScenarioP p m ByTime
+newHistoryByTime = singleton . NewHistoryByTime
+
+assert :: p -> ScenarioP p m ()
+assert = singleton . Assert
+
+example
+    :: Monad m
+    => (forall a . Eq a => a -> a -> p)
+    -> ProgramT (Scenario p) m ()
+example assertEq = do
+    h0 <- historyByTime
+    h0' <- newHistoryByTime $ pure mempty
+    assert $ assertEq h0 h0'
+    tx1 <- existsTx
+    unspent1 <- deposit tx1 1 100
+    _ <- deposit tx1 2 200
+    b1 <- block Origin [tx1]
+    tx2 <- existsTx
+    _ <- deposit tx2 1 200
+    withdrawal tx2 unspent1
+    b2 <- block (At 1) [tx2]
+    rollForward [b1, b2]
+    h1 <- historyByTime
+    h1' <- newHistoryByTime $ newTimes $ do
+        at Origin $ do
+            forCustomer 1 $ do
+                inTx tx1 $ deposited 100
+            forCustomer 2 $ do
+                inTx tx1 $ deposited 200
+        at (At 1) $ do
+            forCustomer 1 $ do
+                inTx tx2 $ do
+                    deposited 200
+                    withdrawn 100
+    assert $ assertEq h1 h1'
+    rollBackward (At 1)
+    rollForward [b2]
+    h2 <- historyByTime
+    assert $ assertEq h1 h2
+
+newTxId :: Monad m => StateT (Map TxI TxBody) m TxI
+newTxId = do
+    txs <- get
+    let TxI z = fst $ Map.findMax txs
+        txId = TxI (z + 1)
+    put $ Map.insert txId emptyTxBody txs
+    return txId
+
+interpret
+    :: Monad m
+    => (Customer -> Address)
+    -> (Slot -> WithOrigin UTCTime)
+    -> ScenarioP
+        p
+        (StateT (WalletState, Map TxI TxBody) m)
+        ()
+    -> StateT WalletState m ()
+interpret customerAddresses _slotTimes p = do
+    walletState <- get
+    (walletState', _) <- lift $ execStateT (go p) (walletState, mempty)
+    put walletState'
+  where
+    go = viewT >=> eval
+    -- eval :: Monad m => ProgramViewT (ScenarioP p m) m a -> m a
+    eval (Return x) = return x
+    eval (ExistsTx :>>= k) = do
+        txId <- zoom _2 newTxId
+        go $ k txId
+    eval (Deposit tx customer value :>>= k) = do
+        let v = mkAda $ fromIntegral value
+            txOut = mkTxOut (customerAddresses customer) v
+        (w, txs) <- get
+        let txBody = txs Map.! tx
+            (txBody', ix) = addTxOut txOut txBody
+        put (w, Map.insert tx txBody' txs)
+
+        go $ k $ UnspentI (tx, ix)
+    eval (Withdrawal tx (UnspentI (tx', ix)) :>>= k) = do
+        (w, txs) <- get
+        let txInBody = txs Map.! tx'
+            txId = getTxId $ mkTx txInBody
+        let txBody = txs Map.! tx
+            txIn = TxIn (txId) ix
+            txBody' = addTxIn txIn txBody
+        put (w, Map.insert tx txBody' txs)
+        go $ k ()
+    -- eval (CreateBlock slot txs :>>= k) = do
+    --     (w, txs) <- get
+    --     let txs' = Map.restrictKeys txs (Map.fromList txs)
+    --     put (w, txs')
+    --     go $ k $ BlockI slot
+    eval _ = error "Not implemented"
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/ByTime.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/ByTime.hs
new file mode 100644
index 00000000000..9beba479f48
--- /dev/null
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/ByTime.hs
@@ -0,0 +1,218 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Cardano.Wallet.Deposit.Testing.DSL.ByTime where
+
+import Prelude
+
+import Cardano.Wallet.Deposit.Map
+    ( Map (Map, Value)
+    , W
+    , toFinger
+    )
+import Cardano.Wallet.Deposit.Pure
+    ( Customer
+    , ValueTransfer (received, spent)
+    )
+import Cardano.Wallet.Deposit.Pure.API.TxHistory
+    ( ByTime
+    , DownTime
+    , firstJust
+    )
+import Cardano.Wallet.Deposit.Read
+    ( Address
+    )
+import Cardano.Wallet.Deposit.Testing.DSL.Types
+    ( TxI
+    )
+import Cardano.Wallet.Deposit.Write
+    ( mkAda
+    )
+import Cardano.Wallet.Read
+    ( Slot
+    , TxId
+    , WithOrigin (..)
+    )
+import Control.Monad.Reader
+    ( MonadReader (..)
+    , ReaderT (..)
+    , asks
+    , withReaderT
+    )
+import Control.Monad.State
+    ( StateT
+    , execStateT
+    , modify'
+    )
+import Control.Monad.Trans
+    ( MonadTrans (..)
+    )
+import Data.Map.Monoidal.Strict
+    ( MonoidalMap
+    )
+import Data.Monoid
+    ( First
+    )
+import Data.Ord
+    ( Down (..)
+    )
+import Data.Time
+    ( UTCTime
+    )
+
+import qualified Data.Map.Monoidal.Strict as MonoidalMap
+
+-- -------------------------------------------------------------------------------
+-- -- AtTime
+-- -------------------------------------------------------------------------------
+
+type ByTimeM =
+    ReaderT
+        ( (TxI -> TxId, Customer -> Address)
+        , Slot -> WithOrigin UTCTime
+        )
+
+at
+    :: Monad m
+    => Slot
+    -> StateT
+        (MonoidalMap Customer (Map '[W (First Address) TxId] ValueTransfer))
+        (ReaderT (TxI -> TxId, Customer -> Address) m)
+        ()
+    -> StateT
+        ( MonoidalMap
+            DownTime
+            (Map '[W (First Slot) Customer, W (First Address) TxId] ValueTransfer)
+        )
+        (ByTimeM m)
+        ()
+at t v = do
+    timeOf <- asks snd
+    txs <- lift $ withReaderT fst $ newCustomers t v
+    modify' $ MonoidalMap.insert (Down $ timeOf t) txs
+
+newTimes
+    :: Monad m
+    => StateT
+        ( MonoidalMap
+            DownTime
+            (Map '[W (First Slot) Customer, W (First Address) TxId] ValueTransfer)
+        )
+        (ByTimeM m)
+        ()
+    -> ByTimeM m ByTime
+newTimes v = toFinger . Map () <$> execStateT v mempty
+
+-- -------------------------------------------------------------------------------
+-- -- Customer
+-- -------------------------------------------------------------------------------
+
+forCustomer
+    :: Monad m
+    => Customer
+    -> StateT
+        (MonoidalMap TxId (Map '[] ValueTransfer))
+        (ReaderT (TxI -> TxId) m)
+        ()
+    -> StateT
+        (MonoidalMap Customer (Map '[W (First Address) TxId] ValueTransfer))
+        (ReaderT (TxI -> TxId, Customer -> Address) m)
+        ()
+forCustomer c v = do
+    addrOf <- asks snd
+    txs <- lift $ withReaderT fst $ newTxIds (addrOf c) v
+    modify' $ MonoidalMap.insert c txs
+
+newCustomers
+    :: Monad m
+    => Slot
+    -> StateT
+        (MonoidalMap Customer (Map '[W (First Address) TxId] ValueTransfer))
+        (ReaderT (TxI -> TxId, Customer -> Address) m)
+        ()
+    -> ReaderT
+        (TxI -> TxId, Customer -> Address)
+        m
+        (Map '[W (First Slot) Customer, W (First Address) TxId] ValueTransfer)
+newCustomers slot v = Map (firstJust slot) <$> execStateT v mempty
+
+-------------------------------------------------------------------------------
+-- Tx
+-------------------------------------------------------------------------------
+
+inTx
+    :: Monad m
+    => TxI
+    -> StateT ValueTransfer m ()
+    -> StateT
+        (MonoidalMap TxId (Map '[] ValueTransfer))
+        (ReaderT (TxI -> TxId) m)
+        ()
+inTx tx v = do
+    w <- lift . lift $ newValueTransferP v
+    txIdOf <- ask
+    modify' $ MonoidalMap.insert (txIdOf tx) w
+
+newTxIds
+    :: Monad m
+    => Address
+    -> StateT
+        (MonoidalMap TxId (Map '[] ValueTransfer))
+        (ReaderT (TxI -> TxId) m)
+        ()
+    -> ReaderT (TxI -> TxId) m (Map '[W (First Address) TxId] ValueTransfer)
+newTxIds addr v = Map (firstJust addr) <$> execStateT v mempty
+
+-------------------------------------------------------------------------------
+-- Value transfer
+-------------------------------------------------------------------------------
+
+deposited :: Monad m => Int -> StateT ValueTransfer m ()
+deposited n = modify' $ \s -> s{received = mkAda $ fromIntegral n}
+
+withdrawn :: Monad m => Int -> StateT ValueTransfer m ()
+withdrawn n = modify' $ \s -> s{spent = mkAda $ fromIntegral n}
+
+newValueTransferP
+    :: forall m
+     . Monad m
+    => StateT ValueTransfer m ()
+    -> m (Map '[] ValueTransfer)
+newValueTransferP v = Value <$> execStateT v mempty
+
+-------------------------------------------------------------------------------
+-- Example
+-------------------------------------------------------------------------------
+-- >>> ex (TxI 1) (TxI 2)
+-- No instance for `Show
+--                    (ReaderT
+--                       ((TxI -> TxId, Customer -> Address), WithOrigin UTCTime -> Slot)
+--                       m0_aCM2Y[tau:0]
+--                       ByTime)'
+--   arising from a use of `evalPrint'
+-- In a stmt of an interactive GHCi command: evalPrint it_aCM0w
+ex
+    :: Monad m
+    => TxI
+    -> TxI
+    -> ( ByTimeM
+            m
+            ByTime
+       )
+ex tx1 tx2 = newTimes $ do
+    at Origin $ do
+        forCustomer 1 $ do
+            inTx tx1 $ do
+                deposited 1
+        forCustomer 2 $ do
+            inTx tx1 $ do
+                deposited 1
+                withdrawn 2
+    at (At 1) $ do
+        forCustomer 1 $ do
+            inTx tx2 $ do
+                withdrawn 1
+            inTx tx2 $ do
+                deposited 1
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/Types.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/Types.hs
new file mode 100644
index 00000000000..da92b7ce845
--- /dev/null
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Testing/DSL/Types.hs
@@ -0,0 +1,24 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Cardano.Wallet.Deposit.Testing.DSL.Types where
+
+import Prelude
+
+import Cardano.Wallet.Deposit.Read
+    ( Ix
+    )
+
+newtype TxI = TxI Int
+    deriving (Eq, Ord, Show)
+
+newtype UnspentI = UnspentI (TxI, Ix)
+    deriving (Eq, Ord, Show)
+
+newtype BlockI = BlockI Int
+    deriving (Eq, Ord, Show)
+
+newtype TimeI = TimeI Int
+    deriving (Eq, Ord, Show)
diff --git a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Write.hs b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Write.hs
index 77f1717a918..f6d6ed92e17 100644
--- a/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Write.hs
+++ b/lib/customer-deposit-wallet/src/Cardano/Wallet/Deposit/Write.hs
@@ -1,3 +1,7 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeApplications #-}
+
 -- | Indirection module that re-exports types
 -- used for writing transactions to the blockchain,
 -- in the most recent and the next future eras.
@@ -6,9 +10,7 @@
 module Cardano.Wallet.Deposit.Write
     ( -- * Basic types
       Address
-
     , Value
-
     , TxId
     , Tx
     , mkTx
@@ -16,7 +18,7 @@ module Cardano.Wallet.Deposit.Write
     , TxIn
     , TxOut
 
-    -- * Transaction balancing
+      -- * Transaction balancing
     , Write.IsRecentEra
     , Write.Conway
     , L.PParams
@@ -32,12 +34,16 @@ module Cardano.Wallet.Deposit.Write
     , Write.ErrBalanceTx (..)
     , Write.balanceTx
 
-    -- ** Time interpreter
+      -- ** Time interpreter
     , Write.TimeTranslation
-    -- * Helper functions
+
+      -- * Helper functions
     , mkAda
     , mkTxOut
     , toConwayTx
+    , emptyTxBody
+    , addTxOut
+    , addTxIn
     ) where
 
 import Prelude
@@ -56,6 +62,13 @@ import Cardano.Wallet.Deposit.Read
 import Cardano.Wallet.Read.Tx
     ( toConwayOutput
     )
+import Control.Lens
+    ( Lens'
+    , lens
+    , makeLenses
+    , (&)
+    , (.~)
+    )
 import Data.Map
     ( Map
     )
@@ -70,10 +83,6 @@ import Data.Sequence.Strict
 import Data.Set
     ( Set
     )
-import Lens.Micro
-    ( (&)
-    , (.~)
-    )
 
 import qualified Cardano.Ledger.Api as L
 import qualified Cardano.Ledger.Api.Tx.In as L
@@ -96,6 +105,26 @@ data TxBody = TxBody
     }
     deriving (Show)
 
+txOutsL :: Lens' TxBody (Map Ix TxOut)
+txOutsL = lens txouts (\s a -> s { txouts = a })
+
+makeLenses ''TxBody
+
+nextIx :: TxBody -> Ix
+nextIx = maybe minBound (succ . fst) . Map.lookupMax . txouts
+
+addTxOut :: TxOut -> TxBody -> (TxBody, Ix)
+addTxOut txout txbody = (txBody', txIx)
+    where
+        txBody' = txbody & txOutsL .~ Map.insert txIx txout (txouts txbody)
+        txIx = nextIx txbody
+
+addTxIn :: TxIn -> TxBody -> TxBody
+addTxIn txin txbody = txbody { spendInputs = Set.insert txin (spendInputs txbody) }
+
+emptyTxBody :: TxBody
+emptyTxBody = TxBody mempty mempty mempty Nothing
+
 mkAda :: Integer -> Value
 mkAda = Read.injectCoin . Read.CoinC
 
@@ -116,9 +145,9 @@ mkTx txbody = Read.Tx $ L.mkBasicTx txBody
                 .~ Set.map toLedgerTxIn (collInputs txbody)
                 )
             & (L.outputsTxBodyL .~ toLedgerTxOuts (txouts txbody))
-            & (L.collateralReturnTxBodyL
-                .~ toLedgerMaybeTxOut (collRet txbody)
-                )
+            & ( L.collateralReturnTxBodyL
+                    .~ toLedgerMaybeTxOut (collRet txbody)
+              )
 
 toLedgerTxIn :: TxIn -> L.TxIn L.StandardCrypto
 toLedgerTxIn = id