6 files changed, 32 insertions, 20 deletions
diff --git a/src/Erebos/Chatroom.hs b/src/Erebos/Chatroom.hs
index c535845..c050381 100644
--- a/src/Erebos/Chatroom.hs
+++ b/src/Erebos/Chatroom.hs
@@ -371,12 +371,11 @@ updateChatroomByStateData lookupData newName newDesc = findAndUpdateChatroomStat
 
 listChatrooms :: MonadHead LocalState m => m [ChatroomState]
 listChatrooms = filter (not . roomStateDeleted) .
-    fromSetBy (comparing $ roomName <=< roomStateRoom) .
-    lookupSharedValue . lsShared . fromStored <$> getLocalHead
+    fromSetBy (comparing $ roomName <=< roomStateRoom) <$> lookupSharedValueM
 
 findChatroom :: MonadHead LocalState m => (ChatroomState -> Bool) -> m (Maybe ChatroomState)
 findChatroom p = do
-    list <- map snd . chatroomSetToList . lookupSharedValue . lsShared . fromStored <$> getLocalHead
+    list <- map snd . chatroomSetToList <$> lookupSharedValueM
     return $ find p list
 
 findChatroomByRoomData :: MonadHead LocalState m => Stored (Signed ChatroomData) -> m (Maybe ChatroomState)
@@ -453,7 +452,7 @@ data ChatroomSetChange = AddedChatroom ChatroomState
 watchChatrooms :: MonadIO m => Head LocalState -> (Set ChatroomState -> Maybe [ChatroomSetChange] -> IO ()) -> m WatchedHead
 watchChatrooms h f = liftIO $ do
     lastVar <- newIORef Nothing
-    watchHeadWith h (lookupSharedValue . lsShared . headObject) $ \cur -> do
+    watchHeadWith h lookupSharedValueH $ \cur -> do
         let curList = chatroomSetToList cur
         mbLast <- readIORef lastVar
         writeIORef lastVar $ Just curList
@@ -537,7 +536,7 @@ instance Service ChatroomService where
         svcModify $ \s -> s { psSendRoomUpdates = True }
 
         when (not previouslyUpdated) $ do
-            syncChatroomsToPeer . lookupSharedValue . lsShared . fromStored =<< getLocalHead
+            syncChatroomsToPeer =<< lookupSharedValueM
 
         when chatRoomQuery $ do
             rooms <- listChatrooms
diff --git a/src/Erebos/Conversation.hs b/src/Erebos/Conversation.hs
index e46f4a5..844b5dd 100644
--- a/src/Erebos/Conversation.hs
+++ b/src/Erebos/Conversation.hs
@@ -115,7 +115,7 @@ isSameConversation _ _ = False
 directMessageConversation :: MonadHead LocalState m => ComposedIdentity -> m Conversation
 directMessageConversation peer = do
     createOrUpdateDirectMessagePeer peer
-    (find (sameIdentity peer . msgPeer) . dmThreadList . lookupSharedValue . lsShared . fromStored <$> getLocalHead) >>= \case
+    (find (sameIdentity peer . msgPeer) . dmThreadList <$> lookupSharedValueM) >>= \case
         Just thread -> return $ DirectMessageConversation thread
         Nothing -> return $ DirectMessageConversation $ dmEmptyThread peer
 
@@ -135,7 +135,7 @@ reloadConversation cur@(ChatroomConversation rstate) =
     fromMaybe cur <$> chatroomConversation rstate
 
 lookupConversations :: MonadHead LocalState m => m [ Conversation ]
-lookupConversations = map DirectMessageConversation . dmThreadList . lookupSharedValue . lsShared . fromStored <$> getLocalHead
+lookupConversations = map DirectMessageConversation . dmThreadList <$> lookupSharedValueM
 
 lookupConversationByRef :: MonadHead LocalState m => RefDigest -> m (Maybe Conversation)
 lookupConversationByRef dgst = find ((dgst ==) . conversationReference) <$> lookupConversations
diff --git a/src/Erebos/DirectMessage.hs b/src/Erebos/DirectMessage.hs
index 3b298e7..ed2c2b4 100644
--- a/src/Erebos/DirectMessage.hs
+++ b/src/Erebos/DirectMessage.hs
@@ -120,8 +120,8 @@ instance Service DirectMessage where
         let msg = fromStored smsg
         powner <- asks $ finalOwner . svcPeerIdentity
         erb <- svcGetLocal
-        let DirectMessageThreads prev _ = lookupSharedValue $ lsShared $ fromStored erb
-            sent = concat $ propertyValue $ findMsgProperty powner msSent prev
+        DirectMessageThreads prev _ <- lookupSharedValueM
+        let sent = concat $ propertyValue $ findMsgProperty powner msSent prev
             received = concat $ propertyValue $ findMsgProperty powner msReceived prev
             received' = filterAncestors $ smsg : received
         if powner `sameIdentity` msgFrom msg ||
@@ -146,7 +146,7 @@ instance Service DirectMessage where
            else join $ asks $ dmOwnerMismatch . svcAttributes
 
     serviceNewPeer = do
-        syncDirectMessageToPeer . lookupSharedValue . lsShared . fromStored =<< svcGetLocal
+        syncDirectMessageToPeer =<< lookupSharedValueM
 
     serviceUpdatedPeer = do
         updateDirectMessagePeer . finalOwner =<< asks svcPeerIdentity
@@ -424,7 +424,7 @@ messageThreadFor pthread mss =
 watchDirectMessageThreads :: Head LocalState -> (DirectMessageThread -> DirectMessageThread -> IO ()) -> IO WatchedHead
 watchDirectMessageThreads h callback = do
     prevVar <- newMVar Nothing
-    watchHeadWith h (lookupSharedValue . lsShared . headObject) $ \(DirectMessageThreads sms _) -> do
+    watchHeadWith h lookupSharedValueH $ \(DirectMessageThreads sms _) -> do
         modifyMVar_ prevVar $ \case
             Just ( prev, prevPeers ) -> do
                 let addPeer (p : ps) p'
diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index 71929e3..6c86e83 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -667,7 +667,7 @@ instance Service DiscoveryService where
 
         searchForOwner <- asks (discoverySearchForOwner . svcAttributes) >>= \case
             True -> do
-                (lookupSharedValue . lsShared . fromStored <$> getLocalHead) >>= \case
+                lookupSharedValueM >>= \case
                     Just (self :: ComposedIdentity) -> do
                         return $ S.fromList $ map (refDigest . storedRef) $ idDataF self
                     Nothing -> do
diff --git a/src/Erebos/Invite.hs b/src/Erebos/Invite.hs
index c6e037d..5b35506 100644
--- a/src/Erebos/Invite.hs
+++ b/src/Erebos/Invite.hs
@@ -190,8 +190,7 @@ acceptInvite
     -> m ()
 acceptInvite from token = do
     prev <- find ((Just token ==) . acceptedInviteToken)
-        . fromSetBy (comparing acceptedInviteToken) . lookupSharedValue . lsShared . fromStored
-        <$> getLocalHead
+        . fromSetBy (comparing acceptedInviteToken) <$> lookupSharedValueM
     accepted <- mergeSorted @AcceptedInvite . (: []) <$> mstore AcceptedInviteData
         { aidPrev = maybe [] acceptedInviteData prev
         , aidToken = Just token
@@ -261,7 +260,7 @@ instance Service InviteService where
 
     serviceHandler = fromStored >>> \case
         AcceptInvite token -> do
-            invites <- fromSetBy (comparing inviteToken) . lookupSharedValue . lsShared . fromStored <$> getLocalHead
+            invites <- fromSetBy (comparing inviteToken) <$> lookupSharedValueM
             case find ((Just token ==) . inviteToken) invites of
                 Just invite
                     | Just name <- inviteContact invite
@@ -296,7 +295,7 @@ instance Service InviteService where
             asks (inviteHookReplyInvalid . svcAttributes) >>= ($ token)
             svcModify $ filter (/= token)
 
-            accepted <- fromSetBy (comparing acceptedInviteToken) . lookupSharedValue . lsShared . fromStored <$> getLocalHead
+            accepted <- fromSetBy (comparing acceptedInviteToken) <$> lookupSharedValueM
             case find ((Just token ==) . acceptedInviteToken) accepted of
                 Just invite -> do
                     aidata <- mstore AcceptedInviteData
@@ -316,7 +315,7 @@ instance Service InviteService where
               then do
                 svcSet $ filter (/= token) waitingTokens
 
-                accepted <- fromSetBy (comparing acceptedInviteToken) . lookupSharedValue . lsShared . fromStored <$> getLocalHead
+                accepted <- fromSetBy (comparing acceptedInviteToken) <$> lookupSharedValueM
                 case find ((Just token ==) . acceptedInviteToken) accepted of
                     Just invite -> do
                         aidata <- mstore AcceptedInviteData
@@ -343,7 +342,7 @@ instance Service InviteService where
             svcPrint $ "Received unknown invite packet"
 
     serviceNewPeer = do
-        invites <- fromSetBy (comparing acceptedInviteToken) . lookupSharedValue . lsShared . fromStored <$> getLocalHead
+        invites <- fromSetBy (comparing acceptedInviteToken) <$> lookupSharedValueM
         peerDigests <- asks $ identityOwnerDigests . svcPeerIdentity
         forM_ invites $ \case
             AcceptedInvite
diff --git a/src/Erebos/State.hs b/src/Erebos/State.hs
index 0f0adbe..8874085 100644
--- a/src/Erebos/State.hs
+++ b/src/Erebos/State.hs
@@ -10,6 +10,7 @@ module Erebos.State (
 
     updateLocalState, updateLocalState_,
     updateSharedState, updateSharedState_,
+    lookupSharedValueH, lookupSharedValueM,
     lookupSharedValue, makeSharedStateUpdate,
 
     localIdentity,
@@ -109,10 +110,12 @@ instance SharedType (Maybe ComposedIdentity) where
     sharedTypeID _ = mkSharedTypeID "0c6c1fe0-f2d7-4891-926b-c332449f7871"
 
 
-class (MonadIO m, MonadStorage m) => MonadHead a m where
+class (MonadIO m, MonadStorage m, HeadType a) => MonadHead a m where
     updateLocalHead :: (Stored a -> m (Stored a, b)) -> m b
     getLocalHead :: m (Stored a)
     getLocalHead = updateLocalHead $ \x -> return (x, x)
+    getLocalHeadCache :: proxy a -> m (HeadCacheType a)
+    getLocalHeadCache _ = headCacheInit @a <$> getLocalHead
 
 updateLocalHead_ :: MonadHead a m => (Stored a -> m (Stored a)) -> m ()
 updateLocalHead_ f = updateLocalHead (fmap (,()) . f)
@@ -121,6 +124,7 @@ instance (HeadType a, MonadIO m) => MonadHead a (ReaderT (Head a) m) where
     updateLocalHead f = do
         h <- ask
         snd <$> updateHead' h (\h' -> local (const h') (f $ headStoredObject h'))
+    getLocalHeadCache _ = asks headCache
 
 
 newtype LocalHeadT h m a = LocalHeadT { runLocalHeadT :: Storage -> Stored h -> m ( a, Stored h ) }
@@ -178,7 +182,7 @@ updateSharedState_ f = fmap fst <$> updateSharedState (fmap (,()) . f)
 updateSharedState :: forall a b m. (SharedType a, MonadHead LocalState m) => (a -> m (a, b)) -> Stored LocalState -> m (Stored LocalState, b)
 updateSharedState f = \ls -> do
     let shared = lsShared $ fromStored ls
-        val = lookupSharedValue shared
+    val <- lookupSharedValueM
     (val', x) <- f val
     (,x) <$> if toComponents val' == toComponents val
                 then return ls
@@ -203,6 +207,16 @@ lookupSharedValueObjects sid since = filterAncestors . helper
         | otherwise = helper $ ssPrev (fromStored x) ++ xs
     helper [] = []
 
+lookupSharedValueC :: forall a. SharedType a => HeadCacheType LocalState -> a
+lookupSharedValueC = mergeSorted . filterAncestors . map wrappedLoad . concatMap (ssValue . fromStored) .
+    fromMaybe [] . MS.lookup (sharedTypeID @a Proxy) . lscSharedCache
+
+lookupSharedValueH :: forall a. SharedType a => Head LocalState -> a
+lookupSharedValueH = lookupSharedValueC . headCache
+
+lookupSharedValueM :: forall a m. (SharedType a, MonadHead LocalState m) => m a
+lookupSharedValueM = lookupSharedValueC <$> getLocalHeadCache @LocalState Proxy
+
 lookupSharedValue :: forall a. SharedType a => [ Stored SharedState ] -> a
 lookupSharedValue = mergeSorted . filterAncestors . map wrappedLoad . concatMap (ssValue . fromStored) . lookupSharedValueObjects (sharedTypeID @a Proxy) []