Generalize local state helper functions

6 files changed, 41 insertions, 27 deletions

diff --git a/src/Attach.hs b/src/Attach.hs
index 89ed4bb..e028718 100644
--- a/src/Attach.hs
+++ b/src/Attach.hs
@@ -88,15 +88,14 @@ attachAccept printMsg h peer = do
             liftIO $ printMsg $ "Confirmed peer, but verification of received identity failed"
             return (Nothing, NoPairing)
         OurRequestConfirm (Just (AttachIdentity _ keys (Just identity))) -> do
-            liftIO $ do
-                printMsg $ "Accepted updated identity"
-                updateLocalState_ h $ finalizeAttach identity keys
+            liftIO $ printMsg $ "Accepted updated identity"
+            flip runReaderT h $ updateLocalState_ $ finalizeAttach identity keys
             return (Nothing, PairingDone)
         OurRequestReady -> throwError $ "alredy accepted, waiting for peer"
         PeerRequest {} -> throwError $ "waiting for peer"
         PeerRequestConfirm -> do
             liftIO $ printMsg $ "Accepted new attached device, seding updated identity"
-            owner <- liftIO $ mergeSharedIdentity h
+            owner <- runReaderT mergeSharedIdentity h
             PeerIdentityFull pid <- peerIdentity peer
             Just secret <- liftIO $ loadKey $ idKeyIdentity owner
             liftIO $ do
diff --git a/src/Contact.hs b/src/Contact.hs
index 9accc4d..e0f1a74 100644
--- a/src/Contact.hs
+++ b/src/Contact.hs
@@ -123,17 +123,15 @@ contactAccept printMsg h peer = do
             return (Nothing, OurRequestReady)
         OurRequestConfirm (Just ContactAccepted) -> do
             PeerIdentityFull pid <- peerIdentity peer
-            liftIO $ do
-                printMsg $ "Contact accepted"
-                updateLocalState_ h $ finalizeContact pid
+            liftIO $ printMsg $ "Contact accepted"
+            flip runReaderT h $ updateLocalState_ $ finalizeContact pid
             return (Nothing, PairingDone)
         OurRequestReady -> throwError $ "alredy accepted, waiting for peer"
         PeerRequest {} -> throwError $ "waiting for peer"
         PeerRequestConfirm -> do
             PeerIdentityFull pid <- peerIdentity peer
-            liftIO $ do
-                printMsg $ "Contact accepted"
-                updateLocalState_ h $ finalizeContact pid
+            liftIO $ printMsg $ "Contact accepted"
+            flip runReaderT h $ updateLocalState_ $ finalizeContact pid
             return (Just $ PairingAccept ContactAccepted, PairingDone)
         PairingDone -> throwError $ "alredy done"
         PairingFailed -> throwError $ "alredy failed"
diff --git a/src/Main.hs b/src/Main.hs
index 9dcbae9..1515648 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -99,7 +99,7 @@ main = do
             Nothing -> error "ref does not exist"
             Just ref -> print $ storedGeneration (wrappedLoad ref :: Stored Object)
 
-        ["update-identity"] -> updateSharedIdentity =<< loadLocalStateHead st
+        ["update-identity"] -> runReaderT updateSharedIdentity =<< loadLocalStateHead st
 
         ("update-identity" : srefs) -> do
             sequence <$> mapM (readRef st . BC.pack) srefs >>= \case
@@ -290,8 +290,7 @@ cmdHistory = void $ do
 
 cmdUpdateIdentity :: Command
 cmdUpdateIdentity = void $ do
-    ehead <- asks ciHead
-    liftIO $ updateSharedIdentity ehead
+    runReaderT updateSharedIdentity =<< asks ciHead
 
 cmdAttach :: Command
 cmdAttach = join $ attachToOwner
diff --git a/src/Message.hs b/src/Message.hs
index 192ab9d..1dadc29 100644
--- a/src/Message.hs
+++ b/src/Message.hs
@@ -121,7 +121,7 @@ sendDirectMessage h peer text = do
         self = headLocalIdentity h
         powner = finalOwner pid
 
-    smsg <- liftIO $ updateSharedState h $ \prev -> do
+    smsg <- flip runReaderT h $ updateSharedState $ \prev -> do
         let sent = findMsgProperty powner msSent prev
             received = findMsgProperty powner msReceived prev
 
diff --git a/src/Service.hs b/src/Service.hs
index 0942159..1d506aa 100644
--- a/src/Service.hs
+++ b/src/Service.hs
@@ -110,6 +110,12 @@ data ServiceHandlerState s = ServiceHandlerState
 newtype ServiceHandler s a = ServiceHandler (ReaderT (ServiceInput s) (WriterT [ServiceReply s] (StateT (ServiceHandlerState s) (ExceptT String IO))) a)
     deriving (Functor, Applicative, Monad, MonadReader (ServiceInput s), MonadWriter [ServiceReply s], MonadState (ServiceHandlerState s), MonadError String, MonadIO)
 
+instance MonadHead LocalState (ServiceHandler s) where
+    updateLocalHead f = do
+        (ls, x) <- liftIO . f =<< gets svcLocal
+        modify $ \s -> s { svcLocal = ls }
+        return x
+
 runServiceHandler :: Service s => Head LocalState -> ServiceInput s -> ServiceState s -> ServiceGlobalState s -> ServiceHandler s () -> IO ([ServiceReply s], (ServiceState s, ServiceGlobalState s))
 runServiceHandler h input svc global shandler = do
     let sstate = ServiceHandlerState { svcValue = svc, svcGlobal = global, svcLocal = headStoredObject h }
diff --git a/src/State.hs b/src/State.hs
index 55c55e1..a715f8a 100644
--- a/src/State.hs
+++ b/src/State.hs
@@ -2,6 +2,7 @@ module State (
     LocalState(..),
     SharedState, SharedType(..),
     SharedTypeID, mkSharedTypeID,
+    MonadHead(..),
 
     loadLocalStateHead,
     updateLocalState, updateLocalState_,
@@ -16,6 +17,8 @@ module State (
     interactiveIdentityUpdate,
 ) where
 
+import Control.Monad.Reader
+
 import Data.Foldable
 import Data.Maybe
 import qualified Data.Text as T
@@ -78,6 +81,15 @@ instance SharedType (Signed IdentityData) where
     sharedTypeID _ = mkSharedTypeID "0c6c1fe0-f2d7-4891-926b-c332449f7871"
 
 
+class MonadHead a m where
+    updateLocalHead :: (Stored a -> IO (Stored a, b)) -> m b
+
+instance (HeadType a, MonadIO m) => MonadHead a (ReaderT (Head a) m) where
+    updateLocalHead f = do
+        h <- ask
+        liftIO $ snd <$> updateHead h f
+
+
 loadLocalStateHead :: Storage -> IO (Head LocalState)
 loadLocalStateHead st = loadHeads st >>= \case
     (h:_) -> return h
@@ -114,20 +126,20 @@ headLocalIdentity h =
             (validateIdentity $ lsIdentity ls)
 
 
-updateLocalState_ :: Head LocalState -> (Stored LocalState -> IO (Stored LocalState)) -> IO ()
-updateLocalState_ h f = updateLocalState h (fmap (,()) . f)
+updateLocalState_ :: MonadHead LocalState m => (Stored LocalState -> IO (Stored LocalState)) -> m ()
+updateLocalState_ f = updateLocalState (fmap (,()) . f)
 
-updateLocalState :: Head LocalState -> (Stored LocalState -> IO (Stored LocalState, a)) -> IO a
-updateLocalState h f = snd <$> updateHead h f
+updateLocalState :: MonadHead LocalState m => (Stored LocalState -> IO (Stored LocalState, a)) -> m a
+updateLocalState = updateLocalHead
 
-updateSharedState_ :: SharedType a => Head LocalState -> ([Stored a] -> IO ([Stored a])) -> IO ()
-updateSharedState_ h f = updateSharedState h (fmap (,()) . f)
+updateSharedState_ :: (SharedType a, MonadHead LocalState m) => ([Stored a] -> IO ([Stored a])) -> m ()
+updateSharedState_ f = updateSharedState (fmap (,()) . f)
 
-updateSharedState :: forall a b. SharedType a => Head LocalState -> ([Stored a] -> IO ([Stored a], b)) -> IO b
-updateSharedState h f = updateLocalState h $ \ls -> do
+updateSharedState :: forall a b m. (SharedType a, MonadHead LocalState m) => ([Stored a] -> IO ([Stored a], b)) -> m b
+updateSharedState f = updateLocalHead $ \ls -> do
     let shared = lsShared $ fromStored ls
         val = lookupSharedValue shared
-        st = refStorage $ headRef h
+        st = storedStorage ls
     (val', x) <- f val
     (,x) <$> if val' == val
                 then return ls
@@ -148,14 +160,14 @@ makeSharedStateUpdate st val prev = wrappedStore st SharedState
     }
 
 
-mergeSharedIdentity :: Head LocalState -> IO UnifiedIdentity
-mergeSharedIdentity = flip updateSharedState $ \sdata -> do
+mergeSharedIdentity :: MonadHead LocalState m => m UnifiedIdentity
+mergeSharedIdentity = updateSharedState $ \sdata -> do
     let Just cidentity = validateIdentityF sdata
     identity <- mergeIdentity cidentity
     return ([idData identity], identity)
 
-updateSharedIdentity :: Head LocalState -> IO ()
-updateSharedIdentity = flip updateSharedState_ $ \sdata -> do
+updateSharedIdentity :: MonadHead LocalState m => m ()
+updateSharedIdentity = updateSharedState_ $ \sdata -> do
     let Just identity = validateIdentityF sdata
     (:[]) . idData <$> interactiveIdentityUpdate identity