3 files changed, 22 insertions, 6 deletions
diff --git a/src/Erebos/State.hs b/src/Erebos/State.hs
index fa58118..0f0adbe 100644
--- a/src/Erebos/State.hs
+++ b/src/Erebos/State.hs
@@ -120,7 +120,7 @@ updateLocalHead_ f = updateLocalHead (fmap (,()) . f)
 instance (HeadType a, MonadIO m) => MonadHead a (ReaderT (Head a) m) where
     updateLocalHead f = do
         h <- ask
-        snd <$> updateHead h f
+        snd <$> updateHead' h (\h' -> local (const h') (f $ headStoredObject h'))
 
 
 newtype LocalHeadT h m a = LocalHeadT { runLocalHeadT :: Storage -> Stored h -> m ( a, Stored h ) }
diff --git a/src/Erebos/Storage.hs b/src/Erebos/Storage.hs
index d9559bc..f6ee200 100644
--- a/src/Erebos/Storage.hs
+++ b/src/Erebos/Storage.hs
@@ -14,7 +14,7 @@ module Erebos.Storage (
     HeadID, HeadTypeID,
     headId, headStorage, headRef, headObject, headStoredObject,
     loadHeads, loadHead, reloadHead,
-    storeHead, replaceHead, updateHead, updateHead_,
+    storeHead, replaceHead, updateHead, updateHead', updateHead_,
 
     WatchedHead,
     watchHead, watchHeadWith, unwatchHead,
diff --git a/src/Erebos/Storage/Head.hs b/src/Erebos/Storage/Head.hs
index 5c17ceb..343669a 100644
--- a/src/Erebos/Storage/Head.hs
+++ b/src/Erebos/Storage/Head.hs
@@ -15,7 +15,7 @@ module Erebos.Storage.Head (
 
     -- * Loading and storing heads
     loadHeads, loadHead, reloadHead,
-    storeHead, replaceHead, updateHead, updateHead_,
+    storeHead, replaceHead, updateHead, updateHead', updateHead_,
     loadHeadRaw, storeHeadRaw, replaceHeadRaw,
 
     -- * Watching heads
@@ -233,12 +233,28 @@ updateHead
         -- ^ First element contains either the new head as @`Just' h@, or
         -- `Nothing' in case the head no longer exists in storage. Second
         -- element is the value from last call to the update function.
-updateHead h f = do
-    (o, x) <- f $ headStoredObject h
+updateHead h f = updateHead' h (f . headStoredObject)
+
+-- | Update existing existing `Head' of type @a@ in the storage, using a given
+-- function. The update function may be called multiple times in case the head
+-- content changes concurrently during evaluation.
+updateHead'
+    :: (HeadType a, MonadIO m)
+    => Head a  -- ^ Existing head to be updated
+    -> (Head a -> m ( Stored a, b ))
+        -- ^ Function that gets current value of the head and returns updated
+        -- value, along with a custom extra value to be returned from
+        -- `updateHead' call. The function may be called multiple times.
+    -> m ( Maybe (Head a), b )
+        -- ^ First element contains either the new head as @`Just' h@, or
+        -- `Nothing' in case the head no longer exists in storage. Second
+        -- element is the value from last call to the update function.
+updateHead' h f = do
+    (o, x) <- f h
     replaceHead h o >>= \case
         Right h' -> return (Just h', x)
         Left Nothing -> return (Nothing, x)
-        Left (Just h') -> updateHead h' f
+        Left (Just h') -> updateHead' h' f
 
 -- | Update existing existing `Head' of type @a@ in the storage, using a given
 -- function. The update function may be called multiple times in case the head