Track local head cache in main/test commandsHEADmaster

2 files changed, 11 insertions, 0 deletions

diff --git a/main/Main.hs b/main/Main.hs
index 654c04b..5d5df63 100644
--- a/main/Main.hs
+++ b/main/Main.hs
@@ -51,6 +51,7 @@ import Erebos.Set
 import Erebos.State
 import Erebos.Storable
 import Erebos.Storage
+import Erebos.Storage.Head
 import Erebos.Storage.Merge
 import Erebos.Sync
 import Erebos.TextFormat
@@ -635,6 +636,11 @@ instance MonadHead LocalState CommandM where
         modify $ \s -> s { csHead = h' }
         return x
 
+    getLocalHeadCache _ = do
+        h <- maybe (fail "failed to reload head") return =<< reloadHead =<< gets csHead
+        modify $ \s -> s { csHead = h }
+        return $ headCache h
+
 type Command = CommandM ()
 
 getSelectedPeer :: CommandM Peer
diff --git a/main/Test.hs b/main/Test.hs
index f9657c2..6ee75bd 100644
--- a/main/Test.hs
+++ b/main/Test.hs
@@ -289,6 +289,11 @@ instance MonadHead LocalState CommandM where
         modify $ \s -> s { tsHead = Just h' }
         return x
 
+    getLocalHeadCache _ = do
+        h <- maybe (fail "failed to reload head") return =<< reloadHead =<< getOrLoadHead
+        modify $ \s -> s { tsHead = Just h }
+        return $ headCache h
+
 type Command = CommandM ()
 
 commands :: [ ( Text, Command ) ]