Keep unknown items in local state

Changelog: Keep unknown items in local state

4 files changed, 133 insertions, 28 deletions

diff --git a/main/Test.hs b/main/Test.hs
index c6448b8..2155e09 100644
--- a/main/Test.hs
+++ b/main/Test.hs
@@ -16,6 +16,7 @@ import Data.Bool
 import Data.ByteString qualified as B
 import Data.ByteString.Char8 qualified as BC
 import Data.ByteString.Lazy qualified as BL
+import Data.ByteString.Lazy.Char8 qualified as BL
 import Data.Foldable
 import Data.Ord
 import Data.Text (Text)
@@ -244,6 +245,7 @@ type Command = CommandM ()
 commands :: [(Text, Command)]
 commands = map (T.pack *** id)
     [ ("store", cmdStore)
+    , ("load", cmdLoad)
     , ("stored-generation", cmdStoredGeneration)
     , ("stored-roots", cmdStoredRoots)
     , ("stored-set-add", cmdStoredSetAdd)
@@ -259,6 +261,9 @@ commands = map (T.pack *** id)
     , ("peer-drop", cmdPeerDrop)
     , ("peer-list", cmdPeerList)
     , ("test-message-send", cmdTestMessageSend)
+    , ("local-state-get", cmdLocalStateGet)
+    , ("local-state-replace", cmdLocalStateReplace)
+    , ("local-state-wait", cmdLocalStateWait)
     , ("shared-state-get", cmdSharedStateGet)
     , ("shared-state-wait", cmdSharedStateWait)
     , ("watch-local-identity", cmdWatchLocalIdentity)
@@ -299,6 +304,17 @@ cmdStore = do
     ref <- liftIO $ unsafeStoreRawBytes st $ BL.fromChunks [encodeUtf8 otype, BC.singleton ' ', BC.pack (show $ B.length cnt), BC.singleton '\n', cnt]
     cmdOut $ "store-done " ++ show (refDigest ref)
 
+cmdLoad :: Command
+cmdLoad = do
+    st <- asks tiStorage
+    [ tref ] <- asks tiParams
+    Just ref <- liftIO $ readRef st $ encodeUtf8 tref
+    header : content <- return $ BL.lines $ lazyLoadBytes ref
+    cmdOut $ "load-type " <> T.unpack (decodeUtf8 $ BL.toStrict header)
+    forM_ content $ \line -> do
+        cmdOut $ "load-line " <> T.unpack (decodeUtf8 $ BL.toStrict line)
+    cmdOut "load-done"
+
 cmdStoredGeneration :: Command
 cmdStoredGeneration = do
     st <- asks tiStorage
@@ -422,6 +438,7 @@ cmdCreateIdentity = do
         storeHead st $ LocalState
             { lsIdentity = idExtData identity
             , lsShared = shared
+            , lsOther = []
             }
     initTestHead h
 
@@ -511,13 +528,25 @@ cmdTestMessageSend = do
     sendManyToPeer peer $ map (TestMessage . wrappedLoad) refs
     cmdOut "test-message-send done"
 
-cmdSharedStateGet :: Command
-cmdSharedStateGet = do
+cmdLocalStateGet :: Command
+cmdLocalStateGet = do
     h <- getHead
-    cmdOut $ unwords $ "shared-state-get" : map (show . refDigest . storedRef) (lsShared $ headObject h)
+    cmdOut $ unwords $ "local-state-get" : map (show . refDigest . storedRef) [ headStoredObject h ]
 
-cmdSharedStateWait :: Command
-cmdSharedStateWait = do
+cmdLocalStateReplace :: Command
+cmdLocalStateReplace = do
+    st <- asks tiStorage
+    [ told, tnew ] <- asks tiParams
+    Just rold <- liftIO $ readRef st $ encodeUtf8 told
+    Just rnew <- liftIO $ readRef st $ encodeUtf8 tnew
+    ok <- updateLocalHead @LocalState $ \ls -> do
+        if storedRef ls == rold
+          then return ( wrappedLoad rnew, True )
+          else return ( ls, False )
+    cmdOut $ if ok then "local-state-replace-done" else "local-state-replace-failed"
+
+localStateWaitHelper :: Storable a => String -> (Head LocalState -> [ Stored a ]) -> Command
+localStateWaitHelper label sel = do
     st <- asks tiStorage
     out <- asks tiOutput
     h <- getOrLoadHead
@@ -525,15 +554,26 @@ cmdSharedStateWait = do
 
     liftIO $ do
         mvar <- newEmptyMVar
-        w <- watchHeadWith h (lsShared . headObject) $ \cur -> do
+        w <- watchHeadWith h sel $ \cur -> do
             mbobjs <- mapM (readRef st . encodeUtf8) trefs
             case map wrappedLoad <$> sequence mbobjs of
                 Just objs | filterAncestors (cur ++ objs) == cur -> do
-                    outLine out $ unwords $ "shared-state-wait" : map T.unpack trefs
+                    outLine out $ unwords $ label : map T.unpack trefs
                     void $ forkIO $ unwatchHead =<< takeMVar mvar
                 _ -> return ()
         putMVar mvar w
 
+cmdLocalStateWait :: Command
+cmdLocalStateWait = localStateWaitHelper "local-state-wait" ((: []) . headStoredObject)
+
+cmdSharedStateGet :: Command
+cmdSharedStateGet = do
+    h <- getHead
+    cmdOut $ unwords $ "shared-state-get" : map (show . refDigest . storedRef) (lsShared $ headObject h)
+
+cmdSharedStateWait :: Command
+cmdSharedStateWait = localStateWaitHelper "shared-state-wait" (lsShared . headObject)
+
 cmdWatchLocalIdentity :: Command
 cmdWatchLocalIdentity = do
     h <- getOrLoadHead
diff --git a/src/Erebos/State.hs b/src/Erebos/State.hs
index 324127a..3012064 100644
--- a/src/Erebos/State.hs
+++ b/src/Erebos/State.hs
@@ -22,13 +22,15 @@ module Erebos.State (
 import Control.Monad.Except
 import Control.Monad.Reader
 
+import Data.ByteString (ByteString)
+import Data.ByteString.Char8 qualified as BC
 import Data.Foldable
 import Data.Maybe
-import qualified Data.Text as T
-import qualified Data.Text.IO as T
+import Data.Text qualified as T
+import Data.Text.IO qualified as T
 import Data.Typeable
 import Data.UUID (UUID)
-import qualified Data.UUID as U
+import Data.UUID qualified as U
 
 import System.IO
 
@@ -40,6 +42,7 @@ import Erebos.Storage.Merge
 data LocalState = LocalState
     { lsIdentity :: Stored (Signed ExtendedIdentityData)
     , lsShared :: [Stored SharedState]
+    , lsOther :: [ ( ByteString, RecItem ) ]
     }
 
 data SharedState = SharedState
@@ -58,13 +61,16 @@ class Mergeable a => SharedType a where
     sharedTypeID :: proxy a -> SharedTypeID
 
 instance Storable LocalState where
-    store' st = storeRec $ do
-        storeRef "id" $ lsIdentity st
-        mapM_ (storeRef "shared") $ lsShared st
+    store' LocalState {..} = storeRec $ do
+        storeRef "id" lsIdentity
+        mapM_ (storeRef "shared") lsShared
+        storeRecItems lsOther
 
-    load' = loadRec $ LocalState
-        <$> loadRef "id"
-        <*> loadRefs "shared"
+    load' = loadRec $ do
+        lsIdentity <- loadRef "id"
+        lsShared <- loadRefs "shared"
+        lsOther <- filter ((`notElem` [ BC.pack "id", BC.pack "shared" ]) . fst) <$> loadRecItems
+        return LocalState {..}
 
 instance HeadType LocalState where
     headTypeID _ = mkHeadTypeID "1d7491a9-7bcb-4eaa-8f13-c8c4c4087e4e"
@@ -123,7 +129,8 @@ loadLocalStateHead st = loadHeads st >>= \case
             }
         storeHead st $ LocalState
             { lsIdentity = idExtData identity
-            , lsShared = [shared]
+            , lsShared = [ shared ]
+            , lsOther = []
             }
 
 localIdentity :: LocalState -> UnifiedIdentity
diff --git a/src/Erebos/Storage.hs b/src/Erebos/Storage.hs
index 2e6653a..2e60f4e 100644
--- a/src/Erebos/Storage.hs
+++ b/src/Erebos/Storage.hs
@@ -38,6 +38,7 @@ module Erebos.Storage (
     storeEmpty, storeInt, storeNum, storeText, storeBinary, storeDate, storeUUID, storeRef, storeRawRef,
     storeMbEmpty, storeMbInt, storeMbNum, storeMbText, storeMbBinary, storeMbDate, storeMbUUID, storeMbRef, storeMbRawRef,
     storeZRef,
+    storeRecItems,
 
     Load, LoadRec,
     evalLoad,
@@ -210,24 +211,28 @@ copyRef' st ref'@(Ref _ dgst) = refFromDigest st dgst >>= \case Just ref -> retu
                       mbobj <- sequence $ copyObject' st <$> mbobj'
                       sequence $ unsafeStoreObject st <$> join mbobj
 
+copyRecItem' :: forall c c'. (StorageCompleteness c, StorageCompleteness c') => Storage' c' -> RecItem' c -> IO (c (RecItem' c'))
+copyRecItem' st = \case
+    RecEmpty -> return $ return $ RecEmpty
+    RecInt x -> return $ return $ RecInt x
+    RecNum x -> return $ return $ RecNum x
+    RecText x -> return $ return $ RecText x
+    RecBinary x -> return $ return $ RecBinary x
+    RecDate x -> return $ return $ RecDate x
+    RecUUID x -> return $ return $ RecUUID x
+    RecRef x -> fmap RecRef <$> copyRef' st x
+
 copyObject' :: forall c c'. (StorageCompleteness c, StorageCompleteness c') => Storage' c' -> Object' c -> IO (c (Object' c'))
 copyObject' _ (Blob bs) = return $ return $ Blob bs
-copyObject' st (Rec rs) = fmap Rec . sequence <$> mapM copyItem rs
-    where copyItem :: (ByteString, RecItem' c) -> IO (c (ByteString, RecItem' c'))
-          copyItem (n, item) = fmap (n,) <$> case item of
-              RecEmpty -> return $ return $ RecEmpty
-              RecInt x -> return $ return $ RecInt x
-              RecNum x -> return $ return $ RecNum x
-              RecText x -> return $ return $ RecText x
-              RecBinary x -> return $ return $ RecBinary x
-              RecDate x -> return $ return $ RecDate x
-              RecUUID x -> return $ return $ RecUUID x
-              RecRef x -> fmap RecRef <$> copyRef' st x
+copyObject' st (Rec rs) = fmap Rec . sequence <$> mapM (\( n, item ) -> fmap ( n, ) <$> copyRecItem' st item) rs
 copyObject' _ ZeroObject = return $ return ZeroObject
 
 copyRef :: forall c c' m. (StorageCompleteness c, StorageCompleteness c', MonadIO m) => Storage' c' -> Ref' c -> m (LoadResult c (Ref' c'))
 copyRef st ref' = liftIO $ returnLoadResult <$> copyRef' st ref'
 
+copyRecItem :: forall c c' m. (StorageCompleteness c, StorageCompleteness c', MonadIO m) => Storage' c' -> RecItem' c -> m (LoadResult c (RecItem' c'))
+copyRecItem st item' = liftIO $ returnLoadResult <$> copyRecItem' st item'
+
 copyObject :: forall c c'. (StorageCompleteness c, StorageCompleteness c') => Storage' c' -> Object' c -> IO (LoadResult c (Object' c'))
 copyObject st obj' = returnLoadResult <$> copyObject' st obj'
 
@@ -790,6 +795,12 @@ storeZRef name x = StoreRecM $ do
         return $ if isZeroRef ref then []
                                   else [(BC.pack name, RecRef ref)]
 
+storeRecItems :: StorageCompleteness c => [ ( ByteString, RecItem ) ] -> StoreRec c
+storeRecItems items = StoreRecM $ do
+    st <- ask
+    tell $ flip map items $ \( name, value ) -> do
+        value' <- copyRecItem st value
+        return [ ( name, value' ) ]
 
 loadBlob :: (ByteString -> a) -> Load a
 loadBlob f = loadCurrentObject >>= \case
diff --git a/test/storage.test b/test/storage.test
index 0369807..3d898b1 100644
--- a/test/storage.test
+++ b/test/storage.test
@@ -427,3 +427,50 @@ test SharedStateWatcher:
 		send "shared-state-wait $s2" to p2
 		expect /shared-state-wait $s1/ from p2
 		expect /shared-state-wait $s2/ from p2
+
+
+test LocalStateKeepUnknown:
+	let refpat = /blake2#[0-9a-f]*/
+
+	spawn as p
+	with p:
+		send "create-identity Device"
+		send "watch-local-identity"
+		expect /local-identity Device/
+
+		send "local-state-get"
+		expect /local-state-get ($refpat)/ capture s1
+		send "load $s1"
+
+		expect /load-type rec [0-9]*/
+		expect /load-line id:r ($refpat)/ capture id1
+		local:
+			expect /load-(.*)/ capture done
+			guard (done == "done")
+
+		send:
+			"store rec"
+			"id:r $id1"
+			"TEST:i 123"
+			""
+		expect /store-done ($refpat)/ capture s2
+		send "local-state-replace $s1 $s2"
+		expect /local-state-replace-done/
+
+		send "local-state-get"
+		expect /local-state-get $s2/
+
+		send "update-local-identity Device2"
+		expect /local-identity Device2/
+
+		send "local-state-get"
+		expect /local-state-get ($refpat)/ capture s3
+		send "load $s3"
+
+		expect /load-type rec [0-9]*/
+		expect /load-line id:r ($refpat)/ capture id2
+		guard (id1 /= id2)
+		expect /load-line TEST:i 123/
+		local:
+			expect /load-(.*)/ capture done
+			guard (done == "done")