1 files changed, 37 insertions, 6 deletions
diff --git a/src/Erebos/State.hs b/src/Erebos/State.hs
index 076a8c0..06e5c54 100644
--- a/src/Erebos/State.hs
+++ b/src/Erebos/State.hs
@@ -6,6 +6,7 @@ module Erebos.State (
     MonadStorage(..),
     MonadHead(..),
     updateLocalHead_,
+    LocalHeadT(..),
 
     updateLocalState, updateLocalState_,
     updateSharedState, updateSharedState_,
@@ -17,9 +18,11 @@ module Erebos.State (
     mergeSharedIdentity,
 ) where
 
+import Control.Monad
 import Control.Monad.Except
 import Control.Monad.Reader
 
+import Data.Bifunctor
 import Data.ByteString (ByteString)
 import Data.ByteString.Char8 qualified as BC
 import Data.Typeable
@@ -66,7 +69,7 @@ instance Storable LocalState where
         lsPrev <- loadMbRawWeak "PREV"
         lsIdentity <- loadRef "id"
         lsShared <- loadRefs "shared"
-        lsOther <- filter ((`notElem` [ BC.pack "id", BC.pack "shared" ]) . fst) <$> loadRecItems
+        lsOther <- filter ((`notElem` [ BC.pack "PREV", BC.pack "id", BC.pack "shared" ]) . fst) <$> loadRecItems
         return LocalState {..}
 
 instance HeadType LocalState where
@@ -101,6 +104,35 @@ instance (HeadType a, MonadIO m) => MonadHead a (ReaderT (Head a) m) where
         snd <$> updateHead h f
 
 
+newtype LocalHeadT h m a = LocalHeadT { runLocalHeadT :: Storage -> Stored h -> m ( a, Stored h ) }
+
+instance Functor m => Functor (LocalHeadT h m) where
+    fmap f (LocalHeadT act) = LocalHeadT $ \st h -> first f <$> act st h
+
+instance Monad m => Applicative (LocalHeadT h m) where
+    pure x = LocalHeadT $ \_ h -> pure ( x, h )
+    (<*>) = ap
+
+instance Monad m => Monad (LocalHeadT h m) where
+    return = pure
+    LocalHeadT act >>= f = LocalHeadT $ \st h -> do
+        ( x, h' ) <- act st h
+        let (LocalHeadT act') = f x
+        act' st h'
+
+instance MonadIO m => MonadIO (LocalHeadT h m) where
+    liftIO act = LocalHeadT $ \_ h -> ( , h ) <$> liftIO act
+
+instance MonadIO m => MonadStorage (LocalHeadT h m) where
+    getStorage = LocalHeadT $ \st h -> return ( st, h )
+
+instance (HeadType h, MonadIO m) => MonadHead h (LocalHeadT h m) where
+    updateLocalHead f = LocalHeadT $ \st h -> do
+        let LocalHeadT act = f h
+        ( ( h', x ), _ ) <- act st h
+        return ( x, h' )
+
+
 localIdentity :: LocalState -> UnifiedIdentity
 localIdentity ls = maybe (error "failed to verify local identity")
     (updateOwners $ maybe [] idExtDataF $ lookupSharedValue $ lsShared ls)
@@ -128,12 +160,11 @@ updateSharedState :: forall a b m. (SharedType a, MonadHead LocalState m) => (a
 updateSharedState f = \ls -> do
     let shared = lsShared $ fromStored ls
         val = lookupSharedValue shared
-    st <- getStorage
     (val', x) <- f val
     (,x) <$> if toComponents val' == toComponents val
                 then return ls
-                else do shared' <- makeSharedStateUpdate st val' shared
-                        wrappedStore st (fromStored ls) { lsShared = [shared'] }
+                else do shared' <- makeSharedStateUpdate val' shared
+                        mstore (fromStored ls) { lsShared = [shared'] }
 
 lookupSharedValue :: forall a. SharedType a => [Stored SharedState] -> a
 lookupSharedValue = mergeSorted . filterAncestors . map wrappedLoad . concatMap (ssValue . fromStored) . filterAncestors . helper
@@ -141,8 +172,8 @@ lookupSharedValue = mergeSorted . filterAncestors . map wrappedLoad . concatMap
                         | otherwise = helper $ ssPrev (fromStored x) ++ xs
           helper [] = []
 
-makeSharedStateUpdate :: forall a m. MonadIO m => SharedType a => Storage -> a -> [Stored SharedState] -> m (Stored SharedState)
-makeSharedStateUpdate st val prev = liftIO $ wrappedStore st SharedState
+makeSharedStateUpdate :: forall a m. (SharedType a, MonadStorage m) => a -> [ Stored SharedState ] -> m (Stored SharedState)
+makeSharedStateUpdate val prev = mstore SharedState
     { ssPrev = prev
     , ssType = Just $ sharedTypeID @a Proxy
     , ssValue = storedRef <$> toComponents val