Keep the default chain state as part of the history.

This removes the need of passing the initial chain state around and
allows us to use the LocalChainState handle on the BehaviorSpec.

hydra-node/src/Hydra/Chain.hs

@@ -173,34 +173,37 @@ instance Arbitrary Lovelace where
 instance (IsTx tx, Arbitrary (ChainStateType tx)) => Arbitrary (PostTxError tx) where
   arbitrary = genericArbitrary
 
-newtype ChainStateHistory tx = UnsafeChainStateHistory (NonEmpty (ChainStateType tx))
+data ChainStateHistory tx = UnsafeChainStateHistory
+  { history :: NonEmpty (ChainStateType tx)
+  , defaultChainState :: ChainStateType tx
+  }
   deriving (Generic)
 
 currentState :: ChainStateHistory tx -> ChainStateType tx
-currentState (UnsafeChainStateHistory history) = head history
+currentState UnsafeChainStateHistory{history} = head history
 
 pushNewState :: ChainStateType tx -> ChainStateHistory tx -> ChainStateHistory tx
-pushNewState cs (UnsafeChainStateHistory history) = UnsafeChainStateHistory (cs <| history)
+pushNewState cs h@UnsafeChainStateHistory{history} = h{history = cs <| history}
 
 initHistory :: ChainStateType tx -> ChainStateHistory tx
-initHistory cs = UnsafeChainStateHistory (cs :| [])
+initHistory cs =
+  UnsafeChainStateHistory
+    { history = cs :| []
+    , defaultChainState = cs
+    }
 
 rollbackHistory ::
   IsChainState tx =>
-  ChainStateType tx ->
   ChainSlot ->
   ChainStateHistory tx ->
   ChainStateHistory tx
-rollbackHistory initialChainState rollbackChainSlot (UnsafeChainStateHistory history) =
-  UnsafeChainStateHistory $
-    fromMaybe (initialChainState :| []) $
-      let rolledBack =
-            dropWhile
-              (\cs -> chainStateSlot cs > rollbackChainSlot)
-              (toList history)
-       in if null rolledBack
-            then Nothing
-            else Just (fromList rolledBack)
+rollbackHistory rollbackChainSlot h@UnsafeChainStateHistory{history, defaultChainState} =
+  h{history = fromMaybe (defaultChainState :| []) (nonEmpty rolledBack)}
+ where
+  rolledBack =
+    dropWhile
+      (\cs -> chainStateSlot cs > rollbackChainSlot)
+      (toList history)
 
 deriving instance (Eq (ChainStateType tx)) => Eq (ChainStateHistory tx)
 deriving instance (Show (ChainStateType tx)) => Show (ChainStateHistory tx)

hydra-node/src/Hydra/Chain/Direct/Handlers.hs

@@ -34,6 +34,7 @@ import Hydra.Chain (
   ChainEvent (..),
   ChainStateHistory,
   ChainStateType,
+  IsChainState,
   PostChainTx (..),
   PostTxError (..),
   currentState,
@@ -53,7 +54,6 @@ import Hydra.Chain.Direct.State (
   contestationPeriod,
   fanout,
   getKnownUTxO,
-  initialChainState,
   initialize,
   observeSomeTx,
  )
@@ -80,9 +80,9 @@ data LocalChainState m tx = LocalChainState
 -- | Initialize a new local chain state with given 'ChainStateAt' (see also
 -- 'initialChainState').
 newLocalChainState ::
-  MonadSTM m =>
-  ChainStateHistory Tx ->
-  m (LocalChainState m Tx)
+  (MonadSTM m, IsChainState tx) =>
+  ChainStateHistory tx ->
+  m (LocalChainState m tx)
 newLocalChainState chainState = do
   tv <- newTVarIO chainState
   pure
@@ -101,7 +101,7 @@ newLocalChainState chainState = do
   rollback tv chainSlot = do
     rolledBack <-
       readTVar tv
-        <&> rollbackHistory initialChainState chainSlot
+        <&> rollbackHistory chainSlot
     writeTVar tv rolledBack
     pure (currentState rolledBack)
 

hydra-node/src/Hydra/HeadLogic.hs

@@ -945,7 +945,7 @@ recoverChainStateHistory initialChainState =
     HeadIsReadyToFanout -> history
     HeadFannedOut{chainState} -> pushNewState chainState history
     ChainRolledBack{chainState} ->
-      rollbackHistory initialChainState (chainStateSlot chainState) history
+      rollbackHistory (chainStateSlot chainState) history
     TickObserved{} -> history
 
 recoverState ::

hydra-node/test/Hydra/BehaviorSpec.hs

@@ -28,17 +28,16 @@ import Hydra.Cardano.Api (ChainPoint (..), SigningKey, SlotNo (SlotNo), Tx)
 import Hydra.Chain (
   Chain (..),
   ChainEvent (..),
-  ChainStateHistory (UnsafeChainStateHistory),
   ChainStateType,
   HeadId (HeadId),
   HeadParameters (..),
   IsChainState,
   OnChainTx (..),
   PostChainTx (..),
   chainStateSlot,
-  currentState,
   initHistory,
  )
+import Hydra.Chain.Direct.Handlers (getLatest, newLocalChainState, pushNew, rollback)
 import Hydra.Chain.Direct.State (ChainStateAt (..))
 import Hydra.ContestationPeriod (ContestationPeriod (UnsafeContestationPeriod), toNominalDiffTime)
 import Hydra.Crypto (HydraKey, aggregate, sign)
@@ -577,14 +576,15 @@ instance IsChainStateTest Tx where
 -- 'cancel'ed after use. Use 'withSimulatedChainAndNetwork' instead where
 -- possible.
 simulatedChainAndNetwork ::
+  forall m tx.
   (MonadTime m, MonadDelay m, MonadAsync m, IsChainStateTest tx) =>
   ChainStateType tx ->
   m (SimulatedChainNetwork tx m)
 simulatedChainAndNetwork initialChainState = do
   history <- newTVarIO []
   nodes <- newTVarIO []
-  chainStateVar <- newTVarIO (initHistory initialChainState)
-  tickThread <- async $ simulateTicks nodes chainStateVar
+  localChainState <- newLocalChainState (initHistory initialChainState)
+  tickThread <- async $ simulateTicks nodes localChainState
   pure $
     SimulatedChainNetwork
       { connectNode = \node -> do
@@ -595,35 +595,35 @@ simulatedChainAndNetwork initialChainState = do
                   Chain
                     { postTx = \tx -> do
                         now <- getCurrentTime
-                        createAndYieldEvent nodes history chainStateVar $ toOnChainTx now tx
+                        createAndYieldEvent nodes history localChainState $ toOnChainTx now tx
                     , draftCommitTx = \_ -> error "unexpected call to draftCommitTx"
                     , submitTx = \_ -> error "unexpected call to submitTx"
                     }
               , hn = createMockNetwork node nodes
               }
       , tickThread
-      , rollbackAndForward = rollbackAndForward nodes history chainStateVar
+      , rollbackAndForward = rollbackAndForward nodes history localChainState
       , simulateCommit = \(party, committed) ->
-          createAndYieldEvent nodes history chainStateVar $ OnCommitTx{party, committed}
+          createAndYieldEvent nodes history localChainState $ OnCommitTx{party, committed}
       }
  where
   -- seconds
   blockTime = 20
 
-  simulateTicks nodes chainStateVar = forever $ do
+  simulateTicks nodes localChainState = forever $ do
     threadDelay blockTime
     now <- getCurrentTime
     event <- atomically $ do
-      cs <- readTVar chainStateVar
-      pure $ Tick now (chainStateSlot (currentState cs))
+      cs <- getLatest localChainState
+      let chainSlot = chainStateSlot cs
+      pure $ Tick now chainSlot
     readTVarIO nodes >>= mapM_ (`handleChainEvent` event)
 
-  createAndYieldEvent nodes history chainStateVar tx = do
+  createAndYieldEvent nodes history localChainState tx = do
     chainEvent <- atomically $ do
-      cs <- currentState <$> readTVar chainStateVar
+      cs <- getLatest localChainState
       let cs' = advanceSlot cs
-      modifyTVar' chainStateVar $ \(UnsafeChainStateHistory prev) ->
-        UnsafeChainStateHistory (cs' :| tail prev)
+      pushNew localChainState cs'
       pure $
         Observation
           { observedTx = tx
@@ -638,23 +638,22 @@ simulatedChainAndNetwork initialChainState = do
     forM_ ns $ \n ->
       handleChainEvent n chainEvent
 
-  rollbackAndForward nodes history chainStateVar steps = do
+  rollbackAndForward nodes history localChainState steps = do
     -- Split the history after given steps
     (toReplay, kept) <- atomically $ do
       (toReplay, kept) <- splitAt (fromIntegral steps) <$> readTVar history
       writeTVar history kept
       pure (reverse toReplay, kept)
     -- Determine the new (last kept one) chainstate
-    let rolledBackChainStateHistory =
-          UnsafeChainStateHistory . fromList $
+    let chainSlot =
+          List.head $
             map
               ( \case
-                  Observation{newChainState} -> newChainState
+                  Observation{newChainState} -> chainStateSlot newChainState
                   _NoObservation -> error "unexpected non-observation ChainEvent"
               )
               kept
-        rolledBackChainState = currentState rolledBackChainStateHistory
-    atomically $ writeTVar chainStateVar rolledBackChainStateHistory
+    rolledBackChainState <- atomically $ rollback localChainState chainSlot
     -- Yield rollback events
     ns <- readTVarIO nodes
     forM_ ns $ \n -> handleChainEvent n Rollback{rolledBackChainState}