Storage: typed heads

8 files changed, 264 insertions, 211 deletions

diff --git a/src/Attach.hs b/src/Attach.hs
index 761da0f..95f0a67 100644
--- a/src/Attach.hs
+++ b/src/Attach.hs
@@ -111,8 +111,7 @@ instance Service AttachService where
             verifyAttachedIdentity sdata >>= \case
                 Just identity -> do
                     svcPrint $ "Accepted updated identity"
-                    st <- storedStorage <$> svcGetLocal
-                    finalizeAttach st identity keys
+                    svcSetLocal =<< finalizeAttach identity keys =<< svcGetLocal
                 Nothing -> do
                     svcPrint $ "Failed to verify new identity"
                     svcSet $ AttachFailed
@@ -151,9 +150,10 @@ attachToOwner _ self peer = do
         NoAttach -> return (Just $ AttachRequest (nonceDigest self pid nonce BA.empty), OurRequest nonce)
         _ -> throwError "alredy in progress"
 
-attachAccept :: (MonadIO m, MonadError String m) => (String -> IO ()) -> UnifiedIdentity -> Peer -> m ()
-attachAccept printMsg self peer = do
-    let st = storedStorage $ idData self
+attachAccept :: (MonadIO m, MonadError String m) => (String -> IO ()) -> Head LocalState -> Peer -> m ()
+attachAccept printMsg h peer = do
+    let st = refStorage $ headRef h
+        self = headLocalIdentity h
     sendToPeerWith self peer $ \case
         NoAttach -> throwError $ "none in progress"
         OurRequest {} -> throwError $ "waiting for peer"
@@ -161,14 +161,15 @@ attachAccept printMsg self peer = do
             liftIO $ printMsg $ "Confirmed peer, waiting for updated identity"
             return (Nothing, OurRequestReady)
         OurRequestConfirm (Just (identity, keys)) -> do
-            liftIO $ printMsg $ "Accepted updated identity"
-            finalizeAttach st identity keys
+            liftIO $ do
+                printMsg $ "Accepted updated identity"
+                updateLocalState_ h $ finalizeAttach identity keys
             return (Nothing, AttachDone)
         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 st
+            owner <- liftIO $ mergeSharedIdentity h
             PeerIdentityFull pid <- return $ peerIdentity peer
             Just secret <- liftIO $ loadKey $ idKeyIdentity owner
             liftIO $ do
@@ -209,15 +210,15 @@ verifyAttachedIdentity sdata = do
         return identity
 
 
-finalizeAttach :: MonadIO m => Storage -> UnifiedIdentity -> [ScrubbedBytes] -> m ()
-finalizeAttach st identity skeys = do
-    liftIO $ updateLocalState_ st $ \slocal -> do
-        let owner = finalOwner identity
-        pkeys <- mapM (copyStored st) [ idKeyIdentity owner, idKeyMessage owner ]
-        mapM_ storeKey $ catMaybes [ keyFromData sec pub | sec <- skeys, pub <- pkeys ]
+finalizeAttach :: MonadIO m => UnifiedIdentity -> [ScrubbedBytes] -> Stored LocalState -> m (Stored LocalState)
+finalizeAttach identity skeys slocal = liftIO $ do
+    let owner = finalOwner identity
+        st = storedStorage slocal
+    pkeys <- mapM (copyStored st) [ idKeyIdentity owner, idKeyMessage owner ]
+    mapM_ storeKey $ catMaybes [ keyFromData sec pub | sec <- skeys, pub <- pkeys ]
 
-        shared <- makeSharedStateUpdate st (idDataF owner) (lsShared $ fromStored slocal)
-        wrappedStore st (fromStored slocal)
-            { lsIdentity = idData identity
-            , lsShared = [ shared ]
+    shared <- makeSharedStateUpdate st (idDataF owner) (lsShared $ fromStored slocal)
+    wrappedStore st (fromStored slocal)
+        { lsIdentity = idData identity
+        , lsShared = [ shared ]
             }
diff --git a/src/Main.hs b/src/Main.hs
index 34c2b3b..c961f4f 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -74,7 +74,7 @@ main = do
             Nothing -> error "ref does not exist"
             Just ref -> print $ storedGeneration (wrappedLoad ref :: Stored Object)
 
-        ["update-identity"] -> updateSharedIdentity st
+        ["update-identity"] -> updateSharedIdentity =<< loadLocalStateHead st
 
         ("update-identity" : srefs) -> do
             sequence <$> mapM (readRef st . BC.pack) srefs >>= \case
@@ -89,15 +89,14 @@ main = do
 
 interactiveLoop :: Storage -> String -> IO ()
 interactiveLoop st bhost = runInputT defaultSettings $ do
-    origIdentity <- liftIO $ loadLocalIdentity st
-    outputStrLn $ T.unpack $ displayIdentity origIdentity
+    erebosHead <- liftIO $ loadLocalStateHead st
+    outputStrLn $ T.unpack $ displayIdentity $ headLocalIdentity erebosHead
 
     haveTerminalUI >>= \case True -> return ()
                              False -> error "Requires terminal"
     extPrint <- getExternalPrint
     let extPrintLn str = extPrint $ str ++ "\n";
     server <- liftIO $ do
-        erebosHead <- loadLocalStateHead st
         startServer erebosHead extPrintLn bhost
             [ SomeService @AttachService Proxy
             , SomeService @SyncService Proxy
@@ -139,9 +138,12 @@ interactiveLoop st bhost = runInputT defaultSettings $ do
                                        then (cmdSetPeer $ read scmd, args)
                                        else (fromMaybe (cmdUnknown scmd) $ lookup scmd commands, args)
                     _        -> (cmdSend, input)
-            curIdentity <- liftIO $ loadLocalIdentity st
+            h <- liftIO (reloadHead erebosHead) >>= \case
+                Just h  -> return h
+                Nothing -> do lift $ lift $ extPrint "current head deleted"
+                              mzero
             res <- liftIO $ runExceptT $ flip execStateT cstate $ runReaderT cmd CommandInput
-                { ciSelf = curIdentity
+                { ciHead = h
                 , ciServer = server
                 , ciLine = line
                 , ciPrint = extPrintLn
@@ -158,7 +160,7 @@ interactiveLoop st bhost = runInputT defaultSettings $ do
 
 
 data CommandInput = CommandInput
-    { ciSelf :: UnifiedIdentity
+    { ciHead :: Head LocalState
     , ciServer :: Server
     , ciLine :: String
     , ciPrint :: String -> IO ()
@@ -215,41 +217,38 @@ cmdSetPeer n | n < 1     = liftIO $ putStrLn "Invalid peer index"
 
 cmdSend :: Command
 cmdSend = void $ do
-    self <- asks ciSelf
+    ehead <- asks ciHead
     Just peer <- gets csPeer
     text <- asks ciLine
-    smsg <- sendDirectMessage self peer $ T.pack text
+    smsg <- sendDirectMessage ehead peer $ T.pack text
     tzone <- liftIO $ getCurrentTimeZone
     liftIO $ putStrLn $ formatMessage tzone $ fromStored smsg
 
 cmdHistory :: Command
 cmdHistory = void $ do
-    self <- asks ciSelf
-    let st = storedStorage $ idData self
+    ehead <- asks ciHead
     Just peer <- gets csPeer
     PeerIdentityFull pid <- return $ peerIdentity peer
     let powner = finalOwner pid
 
-    Just erebosHead <- liftIO $ loadHead st "erebos"
-    let erebos = wrappedLoad (headRef erebosHead)
     Just thread <- return $ find (sameIdentity powner . msgPeer) $
-        messageThreadView $ lookupSharedValue $ lsShared $ fromStored erebos
+        messageThreadView $ lookupSharedValue $ lsShared $ headObject ehead
     tzone <- liftIO $ getCurrentTimeZone
     liftIO $ mapM_ (putStrLn . formatMessage tzone) $ reverse $ take 50 $ threadToList thread
 
 cmdUpdateIdentity :: Command
 cmdUpdateIdentity = void $ do
-    st <- asks $ storedStorage . idData . ciSelf
-    liftIO $ updateSharedIdentity st
+    ehead <- asks ciHead
+    liftIO $ updateSharedIdentity ehead
 
 cmdAttach :: Command
 cmdAttach = join $ attachToOwner
     <$> asks ciPrint
-    <*> asks ciSelf
+    <*> asks (headLocalIdentity . ciHead)
     <*> (maybe (throwError "no peer selected") return =<< gets csPeer)
 
 cmdAttachAccept :: Command
 cmdAttachAccept = join $ attachAccept
     <$> asks ciPrint
-    <*> asks ciSelf
+    <*> asks ciHead
     <*> (maybe (throwError "no peer selected") return =<< gets csPeer)
diff --git a/src/Message.hs b/src/Message.hs
index 0039d7e..874e375 100644
--- a/src/Message.hs
+++ b/src/Message.hs
@@ -113,14 +113,15 @@ findMsgProperty pid sel mss = concat $ flip findProperty mss $ \x -> do
     return $ sel x
 
 
-sendDirectMessage :: (MonadIO m, MonadError String m) => UnifiedIdentity -> Peer -> Text -> m (Stored DirectMessage)
-sendDirectMessage self peer text = do
+sendDirectMessage :: (MonadIO m, MonadError String m) => Head LocalState -> Peer -> Text -> m (Stored DirectMessage)
+sendDirectMessage h peer text = do
     pid <- case peerIdentity peer of PeerIdentityFull pid -> return pid
                                      _ -> throwError "incomplete peer identity"
-    let st = storedStorage $ idData self
+    let st = refStorage $ headRef h
+        self = headLocalIdentity h
         powner = finalOwner pid
 
-    smsg <- liftIO $ updateSharedState st $ \prev -> do
+    smsg <- liftIO $ updateSharedState h $ \prev -> do
         let sent = findMsgProperty powner msSent prev
             received = findMsgProperty powner msReceived prev
 
diff --git a/src/Network.hs b/src/Network.hs
index f07e7ce..5685627 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -181,7 +181,7 @@ receivedWaitingRef nref wr@(WaitingRef _ _ mvar) = do
     checkWaitingRef wr
 
 
-startServer :: Head -> (String -> IO ()) -> String -> [SomeService] -> IO Server
+startServer :: Head LocalState -> (String -> IO ()) -> String -> [SomeService] -> IO Server
 startServer origHead logd bhost services = do
     let storage = refStorage $ headRef origHead
     chanPeer <- newChan
@@ -271,7 +271,7 @@ startServer origHead logd bhost services = do
                                    forM_ objs $ storeObject $ peerInStorage peer
                                    identity <- readMVar midentity
                                    let svcs = map someServiceID services
-                                   handlePacket logd identity secure peer chanSvc svcs header >>= \case
+                                   handlePacket logd origHead identity secure peer chanSvc svcs header >>= \case
                                        Just peer' -> do
                                            writeChan chanPeer peer'
                                            return $ M.insert paddr peer' pvalue
@@ -307,13 +307,18 @@ startServer origHead logd bhost services = do
                                     { svcPeer = peerId
                                     , svcPrintOp = logd
                                     }
-                            (rsp, (s', gs')) <- handleServicePacket storage inp s gs (wrappedLoad ref :: Stored s)
-                            identity <- readMVar midentity
-                            runExceptT (sendToPeerList identity peer rsp) >>= \case
-                                Left err -> logd $ "failed to send response to peer: " ++ show err
-                                Right () -> return ()
-                            return (M.insert svc (SomeServiceState proxy s') svcs,
-                                M.insert svc (SomeServiceGlobalState proxy gs') global)
+                            reloadHead origHead >>= \case
+                                Nothing -> do
+                                    logd $ "current head deleted"
+                                    return (svcs, global)
+                                Just h -> do
+                                    (rsp, (s', gs')) <- handleServicePacket h inp s gs (wrappedLoad ref :: Stored s)
+                                    identity <- readMVar midentity
+                                    runExceptT (sendToPeerList identity peer rsp) >>= \case
+                                        Left err -> logd $ "failed to send response to peer: " ++ show err
+                                        Right () -> return ()
+                                    return (M.insert svc (SomeServiceState proxy s') svcs,
+                                        M.insert svc (SomeServiceGlobalState proxy gs') global)
                         _ -> do
                             logd $ "unhandled service '" ++ show (toUUID svc) ++ "'"
                             return (svcs, global)
@@ -352,10 +357,10 @@ appendDistinct x (y:ys) | x == y    = y : ys
                         | otherwise = y : appendDistinct x ys
 appendDistinct x [] = [x]
 
-handlePacket :: (String -> IO ()) -> UnifiedIdentity -> Bool
+handlePacket :: (String -> IO ()) -> Head LocalState -> UnifiedIdentity -> Bool
     -> Peer -> Chan (Peer, ServiceID, Ref) -> [ServiceID]
     -> TransportHeader -> IO (Maybe Peer)
-handlePacket logd identity secure opeer chanSvc svcs (TransportHeader headers) = do
+handlePacket logd origHead identity secure opeer chanSvc svcs (TransportHeader headers) = do
     let sidentity = idData identity
         DatagramAddress paddr = peerAddress opeer
         plaintextRefs = map (refDigest . storedRef) $ concatMap (collectStoredObjects . wrappedLoad) $ concat
@@ -373,7 +378,7 @@ handlePacket logd identity secure opeer chanSvc svcs (TransportHeader headers) =
                 gets (peerChannel . phPeer) >>= \case
                     ChannelOurAccept acc ch | refDigest (storedRef acc) == refDigest ref -> do
                         updatePeer $ \p -> p { peerChannel = ChannelEstablished (fromStored ch) }
-                        finalizedChannel identity
+                        finalizedChannel origHead identity
                     _ -> return ()
 
             Rejected _ -> return ()
@@ -442,7 +447,7 @@ handlePacket logd identity secure opeer chanSvc svcs (TransportHeader headers) =
             TrChannelAccept accref -> do
                 let process = do
                         addHeader $ Acknowledged accref
-                        handleChannelAccept identity accref
+                        handleChannelAccept origHead identity accref
                 gets (peerChannel . phPeer) >>= \case
                     ChannelWait {} -> process
                     ChannelOurRequest {} -> process
@@ -550,8 +555,8 @@ handleChannelRequest identity reqref = do
         Nothing -> do
             updatePeer $ \p -> p { peerChannel = ChannelPeerRequest reqref }
 
-handleChannelAccept :: UnifiedIdentity -> PartialRef -> PacketHandler ()
-handleChannelAccept identity accref = do
+handleChannelAccept :: Head LocalState -> UnifiedIdentity -> PartialRef -> PacketHandler ()
+handleChannelAccept oh identity accref = do
     pst <- gets $ peerStorage . phPeer
     copyRef pst accref >>= \case
         Right acc -> do
@@ -570,12 +575,12 @@ handleChannelAccept identity accref = do
                 { peerIdentity = PeerIdentityFull pid
                 , peerChannel = ChannelEstablished $ fromStored ch
                 }
-            finalizedChannel identity
+            finalizedChannel oh identity
         Left dgst -> throwError $ "missing accept data " ++ BC.unpack (showRefDigest dgst)
 
 
-finalizedChannel :: UnifiedIdentity -> PacketHandler ()
-finalizedChannel self = do
+finalizedChannel :: Head LocalState -> UnifiedIdentity -> PacketHandler ()
+finalizedChannel oh self = do
     -- Identity update
     ist <- gets $ peerInStorage . phPeer
     addHeader $ AnnounceSelf $ partialRef ist $ storedRef $ idData $ self
@@ -585,8 +590,8 @@ finalizedChannel self = do
     gets phPeer >>= \case
         peer | PeerIdentityFull pid <- peerIdentity peer
              , finalOwner pid `sameIdentity` finalOwner self -> do
-                 shared <- lsShared . fromStored <$>
-                     liftIO (loadLocalState $ storedStorage $ idData self)
+                 Just h <- liftIO $ reloadHead oh
+                 let shared = lsShared $ headObject h
                  addHeader $ ServiceType $ serviceID @SyncService Proxy
                  mapM_ (addHeader . ServiceRef . partialRef ist . storedRef) shared
                  mapM_ (addBody . storedRef) shared
diff --git a/src/Service.hs b/src/Service.hs
index b5106ce..704bc67 100644
--- a/src/Service.hs
+++ b/src/Service.hs
@@ -91,11 +91,9 @@ data ServiceHandlerState s = ServiceHandlerState
 newtype ServiceHandler s a = ServiceHandler (ReaderT ServiceInput (WriterT [ServiceReply s] (StateT (ServiceHandlerState s) (ExceptT String IO))) a)
     deriving (Functor, Applicative, Monad, MonadReader ServiceInput, MonadWriter [ServiceReply s], MonadState (ServiceHandlerState s), MonadError String, MonadIO)
 
-handleServicePacket :: Service s => Storage -> ServiceInput -> ServiceState s -> ServiceGlobalState s -> Stored s -> IO ([ServiceReply s], (ServiceState s, ServiceGlobalState s))
-handleServicePacket st input svc global packet = do
-    herb <- loadLocalStateHead st
-    let erb = wrappedLoad $ headRef herb
-        sstate = ServiceHandlerState { svcValue = svc, svcGlobal = global, svcLocal = erb }
+handleServicePacket :: Service s => Head LocalState -> ServiceInput -> ServiceState s -> ServiceGlobalState s -> Stored s -> IO ([ServiceReply s], (ServiceState s, ServiceGlobalState s))
+handleServicePacket h input svc global packet = do
+    let sstate = ServiceHandlerState { svcValue = svc, svcGlobal = global, svcLocal = headStoredObject h }
         ServiceHandler handler = serviceHandler packet
     (runExceptT $ flip runStateT sstate $ execWriterT $ flip runReaderT input $ handler) >>= \case
         Left err -> do
@@ -103,9 +101,9 @@ handleServicePacket st input svc global packet = do
             return ([], (svc, global))
         Right (rsp, sstate')
             | svcLocal sstate' == svcLocal sstate -> return (rsp, (svcValue sstate', svcGlobal sstate'))
-            | otherwise -> replaceHead (svcLocal sstate') (Right herb) >>= \case
-                Left  _ -> handleServicePacket st input svc global packet
-                Right _ -> return (rsp, (svcValue sstate', svcGlobal sstate'))
+            | otherwise -> replaceHead h (svcLocal sstate') >>= \case
+                Left (Just h') -> handleServicePacket h' input svc global packet
+                _              -> return (rsp, (svcValue sstate', svcGlobal sstate'))
 
 svcGet :: ServiceHandler s (ServiceState s)
 svcGet = gets svcValue
diff --git a/src/State.hs b/src/State.hs
index 15ae7d2..8e9e320 100644
--- a/src/State.hs
+++ b/src/State.hs
@@ -3,21 +3,19 @@ module State (
     SharedState, SharedType(..),
     SharedTypeID, mkSharedTypeID,
 
-    loadLocalState, loadLocalStateHead,
+    loadLocalStateHead,
     updateLocalState, updateLocalState_,
 
     updateSharedState, updateSharedState_,
     lookupSharedValue, makeSharedStateUpdate,
 
-    loadLocalIdentity, headLocalIdentity,
+    headLocalIdentity,
 
     mergeSharedIdentity,
     updateSharedIdentity,
     interactiveIdentityUpdate,
 ) where
 
-import Control.Monad
-
 import Data.Foldable
 import Data.Maybe
 import qualified Data.Text as T
@@ -62,6 +60,9 @@ instance Storable LocalState where
         <$> loadRef "id"
         <*> loadRefs "shared"
 
+instance HeadType LocalState where
+    headTypeID _ = mkHeadTypeID "1d7491a9-7bcb-4eaa-8f13-c8c4c4087e4e"
+
 instance Storable SharedState where
     store' st = storeRec $ do
         mapM_ (storeRef "PREV") $ ssPrev st
@@ -77,80 +78,69 @@ instance SharedType (Signed IdentityData) where
     sharedTypeID _ = mkSharedTypeID "0c6c1fe0-f2d7-4891-926b-c332449f7871"
 
 
-loadLocalState :: Storage -> IO (Stored LocalState)
-loadLocalState = return . wrappedLoad . headRef <=< loadLocalStateHead
-
-loadLocalStateHead :: Storage -> IO Head
-loadLocalStateHead st = loadHeadDef st "erebos" $ do
-    putStr "Name: "
-    hFlush stdout
-    name <- T.getLine
-
-    putStr "Device: "
-    hFlush stdout
-    devName <- T.getLine
-
-    (owner, secret) <- if
-        | T.null name -> return (Nothing, Nothing)
-        | otherwise -> do
-            (secret, public) <- generateKeys st
-            (_secretMsg, publicMsg) <- generateKeys st
-
-            return . (, Just secret) . Just =<< wrappedStore st =<< sign secret =<<
-                wrappedStore st (emptyIdentityData public)
-                { iddName = Just name, iddKeyMessage = Just publicMsg }
-
-    (devSecret, devPublic) <- generateKeys st
-    (_devSecretMsg, devPublicMsg) <- generateKeys st
-
-    identity <- wrappedStore st =<< maybe return signAdd secret =<< sign devSecret =<< wrappedStore st (emptyIdentityData devPublic)
-        { iddName = if T.null devName then Nothing else Just devName
-        , iddOwner = owner
-        , iddKeyMessage = Just devPublicMsg
-        }
-
-    shared <- wrappedStore st $ SharedState
-        { ssPrev = []
-        , ssType = Just $ sharedTypeID @(Signed IdentityData) Proxy
-        , ssValue = [storedRef $ fromMaybe identity owner]
-        }
-    return $ LocalState
-        { lsIdentity = identity
-        , lsShared = [shared]
-        }
-
-loadLocalIdentity :: Storage -> IO UnifiedIdentity
-loadLocalIdentity = return . headLocalIdentity <=< loadLocalStateHead
-
-headLocalIdentity :: Head -> UnifiedIdentity
+loadLocalStateHead :: Storage -> IO (Head LocalState)
+loadLocalStateHead st = loadHeads st >>= \case
+    (h:_) -> return h
+    [] -> do
+        putStr "Name: "
+        hFlush stdout
+        name <- T.getLine
+
+        putStr "Device: "
+        hFlush stdout
+        devName <- T.getLine
+
+        (owner, secret) <- if
+            | T.null name -> return (Nothing, Nothing)
+            | otherwise -> do
+                (secret, public) <- generateKeys st
+                (_secretMsg, publicMsg) <- generateKeys st
+
+                return . (, Just secret) . Just =<< wrappedStore st =<< sign secret =<<
+                    wrappedStore st (emptyIdentityData public)
+                    { iddName = Just name, iddKeyMessage = Just publicMsg }
+
+        (devSecret, devPublic) <- generateKeys st
+        (_devSecretMsg, devPublicMsg) <- generateKeys st
+
+        identity <- wrappedStore st =<< maybe return signAdd secret =<< sign devSecret =<< wrappedStore st (emptyIdentityData devPublic)
+            { iddName = if T.null devName then Nothing else Just devName
+            , iddOwner = owner
+            , iddKeyMessage = Just devPublicMsg
+            }
+
+        shared <- wrappedStore st $ SharedState
+            { ssPrev = []
+            , ssType = Just $ sharedTypeID @(Signed IdentityData) Proxy
+            , ssValue = [storedRef $ fromMaybe identity owner]
+            }
+        storeHead st $ LocalState
+            { lsIdentity = identity
+            , lsShared = [shared]
+            }
+
+headLocalIdentity :: Head LocalState -> UnifiedIdentity
 headLocalIdentity h =
-    let ls = load $ headRef h
+    let ls = headObject h
      in maybe (error "failed to verify local identity")
             (updateOwners (lookupSharedValue $ lsShared ls))
             (validateIdentity $ lsIdentity ls)
 
 
-updateLocalState_ :: Storage -> (Stored LocalState -> IO (Stored LocalState)) -> IO ()
-updateLocalState_ st f = updateLocalState st (fmap (,()) . f)
-
-updateLocalState :: Storage -> (Stored LocalState -> IO (Stored LocalState, a)) -> IO a
-updateLocalState st f = do
-    Just erebosHead <- loadHead st "erebos"
-    let ls = wrappedLoad (headRef erebosHead)
-    (ls', x) <- f ls
-    when (ls' /= ls) $ do
-        Right _ <- replaceHead ls' (Right erebosHead)
-        return ()
-    return x
+updateLocalState_ :: Head LocalState -> (Stored LocalState -> IO (Stored LocalState)) -> IO ()
+updateLocalState_ h f = updateLocalState h (fmap (,()) . f)
 
+updateLocalState :: Head LocalState -> (Stored LocalState -> IO (Stored LocalState, a)) -> IO a
+updateLocalState h f = snd <$> updateHead h f
 
-updateSharedState_ :: SharedType a => Storage -> ([Stored a] -> IO ([Stored a])) -> IO ()
-updateSharedState_ st f = updateSharedState st (fmap (,()) . f)
+updateSharedState_ :: SharedType a => Head LocalState -> ([Stored a] -> IO ([Stored a])) -> IO ()
+updateSharedState_ h f = updateSharedState h (fmap (,()) . f)
 
-updateSharedState :: forall a b. SharedType a => Storage -> ([Stored a] -> IO ([Stored a], b)) -> IO b
-updateSharedState st f = updateLocalState st $ \ls -> do
+updateSharedState :: forall a b. SharedType a => Head LocalState -> ([Stored a] -> IO ([Stored a], b)) -> IO b
+updateSharedState h f = updateLocalState h $ \ls -> do
     let shared = lsShared $ fromStored ls
         val = lookupSharedValue shared
+        st = refStorage $ headRef h
     (val', x) <- f val
     (,x) <$> if val' == val
                 then return ls
@@ -171,14 +161,14 @@ makeSharedStateUpdate st val prev = wrappedStore st SharedState
     }
 
 
-mergeSharedIdentity :: Storage -> IO UnifiedIdentity
-mergeSharedIdentity st = updateSharedState st $ \sdata -> do
+mergeSharedIdentity :: Head LocalState -> IO UnifiedIdentity
+mergeSharedIdentity = flip updateSharedState $ \sdata -> do
     let Just cidentity = validateIdentityF sdata
     identity <- mergeIdentity cidentity
     return ([idData identity], identity)
 
-updateSharedIdentity :: Storage -> IO ()
-updateSharedIdentity st = updateSharedState_ st $ \sdata -> do
+updateSharedIdentity :: Head LocalState -> IO ()
+updateSharedIdentity = flip updateSharedState_ $ \sdata -> do
     let Just identity = validateIdentityF sdata
     (:[]) . idData <$> interactiveIdentityUpdate identity
 
diff --git a/src/Storage.hs b/src/Storage.hs
index 5a5d992..92a1e1f 100644
--- a/src/Storage.hs
+++ b/src/Storage.hs
@@ -16,9 +16,11 @@ module Storage (
     storeObject,
     collectObjects, collectStoredObjects,
 
-    Head,
-    headName, headRef, headObject,
-    loadHeads, loadHead, loadHeadDef, replaceHead,
+    Head, HeadType(..),
+    HeadTypeID, mkHeadTypeID,
+    headId, headRef, headObject, headStoredObject,
+    loadHeads, loadHead, reloadHead,
+    storeHead, replaceHead, updateHead, updateHead_,
     watchHead,
 
     Storable(..), ZeroStorable(..),
@@ -88,10 +90,13 @@ import Data.Time.Calendar
 import Data.Time.Clock
 import Data.Time.Format
 import Data.Time.LocalTime
+import Data.Typeable
 import Data.UUID (UUID)
 import qualified Data.UUID as U
+import qualified Data.UUID.V4 as U
 
 import System.Directory
+import System.FilePath
 import System.INotify
 import System.IO.Error
 import System.IO.Unsafe
@@ -106,7 +111,7 @@ openStorage :: FilePath -> IO Storage
 openStorage path = do
     createDirectoryIfMissing True $ path ++ "/objects"
     createDirectoryIfMissing True $ path ++ "/heads"
-    watchers <- newMVar (Nothing, [])
+    watchers <- newMVar ([], [])
     refgen <- newMVar =<< HT.new
     return $ Storage { stBacking = StorageDir path watchers, stParent = Nothing, stRefGeneration = refgen }
 
@@ -357,90 +362,137 @@ collectOtherStored seen _ = ([], seen)
 
 type Head = Head' Complete
 
-headName :: Head -> String
-headName (Head name _) = name
+headId :: Head a -> HeadID
+headId (Head uuid _) = uuid
 
-headRef :: Head -> Ref
-headRef (Head _ ref) = ref
+headRef :: Head a -> Ref
+headRef (Head _ sx) = storedRef sx
 
-headObject :: Storable a => Head -> a
-headObject = load . headRef
+headObject :: Head a -> a
+headObject (Head _ sx) = fromStored sx
 
+headStoredObject :: Head a -> Stored a
+headStoredObject (Head _ sx) = sx
+
+deriving instance StorableUUID HeadID
+deriving instance StorableUUID HeadTypeID
+
+mkHeadTypeID :: String -> HeadTypeID
+mkHeadTypeID = maybe (error "Invalid head type ID") HeadTypeID . U.fromString
+
+class Storable a => HeadType a where
+    headTypeID :: proxy a -> HeadTypeID
 
-loadHeads :: Storage -> IO [Head]
-loadHeads s@(Storage { stBacking = StorageDir { dirPath = spath }}) = do
-    let hpath = spath ++ "/heads/"
-    files <- filterM (doesFileExist . (hpath++)) =<< getDirectoryContents hpath
-    forM files $ \hname -> do
-        (h:_) <- BC.lines <$> B.readFile (hpath ++ "/" ++ hname)
-        Just ref <- readRef s h
-        return $ Head hname ref
-loadHeads Storage { stBacking = StorageMemory { memHeads = theads } } = readMVar theads
 
-loadHead :: Storage -> String -> IO (Maybe Head)
-loadHead s@(Storage { stBacking = StorageDir { dirPath = spath }}) hname = do
+headTypePath :: FilePath -> HeadTypeID -> FilePath
+headTypePath spath (HeadTypeID tid) = spath </> "heads" </> U.toString tid
+
+headPath :: FilePath -> HeadTypeID -> HeadID -> FilePath
+headPath spath tid (HeadID hid) = headTypePath spath tid </> U.toString hid
+
+loadHeads :: forall a. HeadType a => Storage -> IO [Head a]
+loadHeads s@(Storage { stBacking = StorageDir { dirPath = spath }}) = do
+    let hpath = headTypePath spath $ headTypeID @a Proxy
+
+    files <- filterM (doesFileExist . (hpath </>)) =<<
+        handleJust (\e -> guard (isDoesNotExistError e)) (const $ return [])
+        (getDirectoryContents hpath)
+    fmap catMaybes $ forM files $ \hname -> do
+        case U.fromString hname of
+             Just hid -> do
+                 (h:_) <- BC.lines <$> B.readFile (hpath </> hname)
+                 Just ref <- readRef s h
+                 return $ Just $ Head (HeadID hid) $ wrappedLoad ref
+             Nothing -> return Nothing
+loadHeads Storage { stBacking = StorageMemory { memHeads = theads } } = do
+    let toHead ((tid, hid), ref) | tid == headTypeID @a Proxy = Just $ Head hid $ wrappedLoad ref
+                                 | otherwise                  = Nothing
+    catMaybes . map toHead <$> readMVar theads
+
+loadHead :: forall a. HeadType a => Storage -> HeadID -> IO (Maybe (Head a))
+loadHead s@(Storage { stBacking = StorageDir { dirPath = spath }}) hid = do
     handleJust (guard . isDoesNotExistError) (const $ return Nothing) $ do
-        let hpath = spath ++ "/heads/"
-        (h:_) <- BC.lines <$> B.readFile (hpath ++ hname)
+        (h:_) <- BC.lines <$> B.readFile (headPath spath (headTypeID @a Proxy) hid)
         Just ref <- readRef s h
-        return $ Just $ Head hname ref
-loadHead Storage { stBacking = StorageMemory { memHeads = theads } } hname =
-    find ((==hname) . headName) <$> readMVar theads
-
-loadHeadDef :: Storable a => Storage -> String -> IO a -> IO Head
-loadHeadDef s hname gen = loadHead s hname >>= \case
-        Just h  -> return h
-        Nothing -> do obj <- gen
-                      Right h <- replaceHead obj (Left (s, hname))
-                      return h
-
-replaceHead :: Storable a => a -> Either (Storage, String) Head -> IO (Either (Maybe Head) Head)
-replaceHead obj prev = do
-    let (st, name) = either id (\(Head n (Ref s _)) -> (s, n)) prev
-    ref <- store st obj
+        return $ Just $ Head hid $ wrappedLoad ref
+loadHead Storage { stBacking = StorageMemory { memHeads = theads } } hid = do
+    fmap (Head hid . wrappedLoad) . lookup (headTypeID @a Proxy, hid) <$> readMVar theads
+
+reloadHead :: HeadType a => Head a -> IO (Maybe (Head a))
+reloadHead (Head hid (Stored (Ref st _) _)) = loadHead st hid
+
+storeHead :: forall a. HeadType a => Storage -> a -> IO (Head a)
+storeHead st obj = do
+    let tid = headTypeID @a Proxy
+    hid <- HeadID <$> U.nextRandom
+    stored <- wrappedStore st obj
+    case stBacking st of
+         StorageDir { dirPath = spath } -> do
+             Right () <- writeFileChecked (headPath spath tid hid) Nothing $
+                 showRef (storedRef stored) `B.append` BC.singleton '\n'
+             return ()
+         StorageMemory { memHeads = theads } -> do
+             modifyMVar_ theads $ return . (((tid, hid), storedRef stored) :)
+    return $ Head hid stored
+
+replaceHead :: forall a. HeadType a => Head a -> Stored a -> IO (Either (Maybe (Head a)) (Head a))
+replaceHead prev@(Head hid pobj) stored = do
+    let st = storedStorage pobj
+        tid = headTypeID @a Proxy
     case stBacking st of
          StorageDir { dirPath = spath } -> do
-             let filename = spath ++ "/heads/" ++ name
+             let filename = headPath spath tid hid
                  showRefL r = showRef r `B.append` BC.singleton '\n'
 
-             writeFileChecked filename (either (const Nothing) (Just . showRefL . headRef) prev) (showRefL ref) >>= \case
+             writeFileChecked filename (Just $ showRefL $ headRef prev) (showRefL $ storedRef stored) >>= \case
                  Left Nothing -> return $ Left Nothing
                  Left (Just bs) -> do Just oref <- readRef st $ BC.takeWhile (/='\n') bs
-                                      return $ Left $ Just $ Head name oref
-                 Right () -> return $ Right $ Head name ref
+                                      return $ Left $ Just $ Head hid $ wrappedLoad oref
+                 Right () -> return $ Right $ Head hid stored
 
          StorageMemory { memHeads = theads, memWatchers = twatch } -> do
              res <- modifyMVar theads $ \hs -> do
-                 ws <- map snd . filter ((==name) . fst) <$> readMVar twatch
-                 case (partition ((== name) . headName) hs, prev) of
-                      (([], _), Left _)  -> let h = Head name ref
-                                             in return (h:hs, Right (h, ws))
-                      (([], _), Right _) -> return (hs, Left Nothing)
-                      ((h:_, _), Left _) -> return (hs, Left (Just h))
-                      ((h:_, hs'), Right h') | headRef h == headRef h' -> let nh = Head name ref
-                                                                           in return (nh:hs', Right (nh, ws))
-                                             | otherwise -> return (hs, Left (Just h))
+                 ws <- map snd . filter ((==(tid, hid)) . fst) <$> readMVar twatch
+                 return $ case partition ((==(tid, hid)) . fst) hs of
+                     ([] , _  ) -> (hs, Left Nothing)
+                     ((_, r):_, hs') | r == storedRef pobj -> (((tid, hid), storedRef stored) : hs',
+                                                                  Right (Head hid stored, ws))
+                                     | otherwise -> (hs, Left $ Just $ Head hid $ wrappedLoad r)
              case res of
-                  Right (h, ws) -> mapM_ ($h) ws >> return (Right h)
+                  Right (h, ws) -> mapM_ ($ headRef h) ws >> return (Right h)
                   Left x -> return $ Left x
 
-watchHead :: Head -> (Head -> IO ()) -> IO ()
-watchHead (Head name (Ref st _)) cb = do
+updateHead :: HeadType a => Head a -> (Stored a -> IO (Stored a, b)) -> IO (Maybe (Head a), b)
+updateHead h f = do
+    (o, x) <- f $ headStoredObject h
+    replaceHead h o >>= \case
+        Right h' -> return (Just h', x)
+        Left Nothing -> return (Nothing, x)
+        Left (Just h') -> updateHead h' f
+
+updateHead_ :: HeadType a => Head a -> (Stored a -> IO (Stored a)) -> IO (Maybe (Head a))
+updateHead_ h = fmap fst . updateHead h . (fmap (,()) .)
+
+watchHead :: forall a. HeadType a => Head a -> (Head a -> IO ()) -> IO ()
+watchHead (Head hid (Stored (Ref st _) _)) cb = do
+    let cb' = cb . Head hid . wrappedLoad
+        tid = headTypeID @a Proxy
     case stBacking st of
-         StorageDir { dirPath = spath, dirWatchers = mvar } -> modifyMVar_ mvar $ \(mbi, watchers) -> do
-             inotify <- (\f -> maybe f return mbi) $ do
-                 inotify <- initINotify
-                 void $ addWatch inotify [Move] (BC.pack $ spath ++ "/heads") $ \case
-                     MovedIn { filePath = fpath } -> do
-                         let cname = BC.unpack fpath
-                         loadHead st cname >>= \case
-                             Just h -> mapM_ ($h) . map snd . filter ((== cname) . fst) . snd =<< readMVar mvar
-                             Nothing -> return ()
-                     _ -> return ()
-                 return inotify
-             return (Just inotify, (name, cb) : watchers)
-
-         StorageMemory { memWatchers = mvar } -> modifyMVar_ mvar $ return . ((name, cb) :)
+         StorageDir { dirPath = spath, dirWatchers = mvar } -> modifyMVar_ mvar $ \(ilist, watchers) -> do
+             ilist' <- case lookup tid ilist of
+                 Just _ -> return ilist
+                 Nothing -> do
+                     inotify <- initINotify
+                     void $ addWatch inotify [Move] (BC.pack $ headTypePath spath tid) $ \case
+                         MovedIn { filePath = fpath } | Just ihid <- HeadID <$> U.fromASCIIBytes fpath -> do
+                             loadHead @a st ihid >>= \case
+                                 Just h -> mapM_ ($ headRef h) . map snd . filter ((== (tid, ihid)) . fst) . snd =<< readMVar mvar
+                                 Nothing -> return ()
+                         _ -> return ()
+                     return $ (tid, inotify) : ilist
+             return (ilist', ((tid, hid), cb') : watchers)
+
+         StorageMemory { memWatchers = mvar } -> modifyMVar_ mvar $ return . (((tid, hid), cb') :)
 
 
 class Storable a where
diff --git a/src/Storage/Internal.hs b/src/Storage/Internal.hs
index c70e8ae..e4e4f00 100644
--- a/src/Storage/Internal.hs
+++ b/src/Storage/Internal.hs
@@ -24,6 +24,7 @@ import qualified Data.HashTable.IO as HT
 import Data.List
 import Data.Map (Map)
 import qualified Data.Map as M
+import Data.UUID (UUID)
 
 import Foreign.Storable (peek)
 
@@ -57,12 +58,12 @@ showParentStorage Storage { stParent = Just st } = "@" ++ show st
 
 data StorageBacking c
          = StorageDir { dirPath :: FilePath
-                      , dirWatchers :: MVar (Maybe INotify, [(String, Head' c -> IO ())])
+                      , dirWatchers :: MVar ([(HeadTypeID, INotify)], [((HeadTypeID, HeadID), Ref' c -> IO ())])
                       }
-         | StorageMemory { memHeads :: MVar [Head' c]
+         | StorageMemory { memHeads :: MVar [((HeadTypeID, HeadID), Ref' c)]
                          , memObjs :: MVar (Map RefDigest BL.ByteString)
                          , memKeys :: MVar (Map RefDigest ScrubbedBytes)
-                         , memWatchers :: MVar [(String, Head' c -> IO ())]
+                         , memWatchers :: MVar [((HeadTypeID, HeadID), Ref' c -> IO ())]
                          }
     deriving (Eq)
 
@@ -111,9 +112,15 @@ hashToRefDigest = RefDigest . hashFinalize . hashUpdates hashInit . BL.toChunks
 newtype Generation = Generation Int
     deriving (Eq, Show)
 
-data Head' c = Head String (Ref' c)
+data Head' c a = Head HeadID (Stored' c a)
     deriving (Show)
 
+newtype HeadID = HeadID UUID
+    deriving (Eq, Ord, Show)
+
+newtype HeadTypeID = HeadTypeID UUID
+    deriving (Eq, Ord)
+
 data Stored' c a = Stored (Ref' c) a
     deriving (Show)