From 622d76d77b9892cfb4bbaa65bdc58b1dbe6fe36e Mon Sep 17 00:00:00 2001
From: Franco Testagrossa <franco.testagrossa@gmail.com>
Date: Tue, 4 Jun 2024 18:34:36 +0200
Subject: [PATCH] Minor dry on balanceUTxOInHead

---
 .../test/Hydra/Chain/Direct/TxTraceSpec.hs    | 36 ++++++++-----------
 1 file changed, 15 insertions(+), 21 deletions(-)

diff --git a/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs b/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs
index d5f9249fe99..5bd610022ac 100644
--- a/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs
+++ b/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs
@@ -296,6 +296,20 @@ data TxResult = TxResult
 initialAmount :: Natural
 initialAmount = 10
 
+balanceUTxOInHead :: Ord k => Map k Natural -> Map k Natural -> Map k Natural
+balanceUTxOInHead currentUtxoInHead someUTxOToDecrement =
+  let
+    currentUtxoInHead' = fmap naturalToInteger currentUtxoInHead
+    someUTxOToDecrement' = fmap (negate . naturalToInteger) someUTxOToDecrement
+   in
+    Map.map naturalFromInteger
+      . Map.filter (> 0)
+      . Map.map sum
+      $ Map.unionWith
+        (++)
+        (Map.map (: []) currentUtxoInHead')
+        (Map.map (: []) someUTxOToDecrement')
+
 instance StateModel Model where
   data Action Model a where
     Decrement :: {actor :: Actor, snapshot :: ModelSnapshot} -> Action Model TxResult
@@ -394,19 +408,6 @@ instance StateModel Model where
       subset <- sublistOf (Map.toList model)
       return $ Map.fromList subset
 
-    balanceUTxOInHead currentUtxoInHead someUTxOToDecrement =
-      let
-        currentUtxoInHead' = fmap naturalToInteger currentUtxoInHead
-        someUTxOToDecrement' = fmap (negate . naturalToInteger) someUTxOToDecrement
-       in
-        Map.map naturalFromInteger
-          . Map.filter (> 0)
-          . Map.map sum
-          $ Map.unionWith
-            (++)
-            (Map.map (: []) currentUtxoInHead')
-            (Map.map (: []) someUTxOToDecrement')
-
     reduceValues :: ModelUTxO -> Gen ModelUTxO
     reduceValues = Map.traverseWithKey reduceValue
      where
@@ -485,14 +486,7 @@ instance StateModel Model where
         m
           { headState = Open
           , latestSnapshot = snapshotNumber snapshot
-          , utxoInHead =
-              -- XXX: remove the balance
-              let currentUTxOInHead = Map.map naturalToInteger $ utxoInHead m
-                  utxoToDecommit = Map.map (negate . naturalToInteger) $ decommitUTxO snapshot
-                  balancedUTxOInHead =
-                    Map.map sum $
-                      Map.unionWith (++) (Map.map (: []) currentUTxOInHead) (Map.map (: []) utxoToDecommit)
-               in Map.map naturalFromInteger balancedUTxOInHead
+          , utxoInHead = balanceUTxOInHead (utxoInHead m) (decommitUTxO snapshot)
           }
       Close{snapshot} ->
         m