Generic collecting of objects through references

2 files changed, 20 insertions, 9 deletions

diff --git a/src/Network.hs b/src/Network.hs
index 198efb0..f294835 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -17,7 +17,6 @@ import Network.Socket
 import Network.Socket.ByteString (recvFrom, sendTo)
 
 import Identity
-import PubKey
 import Storage
 
 
@@ -108,10 +107,7 @@ peerDiscovery bhost sidentity = do
         putStrLn $ "Got announce: " ++ show ref ++ " from " ++ show peer
         when (ref /= storedRef sidentity) $ void $ sendTo sock (BL.toStrict $ BL.concat
             [ serializeObject $ transportToObject $ IdentityRequest ref (storedRef sidentity)
-            , lazyLoadBytes $ storedRef sidentity
-            , lazyLoadBytes $ storedRef $ signedData $ fromStored sidentity
-            , lazyLoadBytes $ storedRef $ idKeyIdentity $ fromStored $ signedData $ fromStored sidentity
-            , lazyLoadBytes $ storedRef $ signedSignature $ fromStored sidentity
+            , BL.concat $ map (lazyLoadBytes . storedRef) $ collectStoredObjects $ wrappedLoad $ storedRef sidentity
             ]) peer
 
     packet _ _ peer (IdentityRequest ref from) [] = do
@@ -126,10 +122,7 @@ peerDiscovery bhost sidentity = do
                    writeChan chan $ Peer (wrappedLoad from) (DatagramAddress peer)
                    void $ sendTo sock (BL.toStrict $ BL.concat
                        [ serializeObject $ transportToObject $ IdentityResponse (storedRef sidentity)
-                       , lazyLoadBytes $ storedRef sidentity
-                       , lazyLoadBytes $ storedRef $ signedData $ fromStored sidentity
-                       , lazyLoadBytes $ storedRef $ idKeyIdentity $ fromStored $ signedData $ fromStored sidentity
-                       , lazyLoadBytes $ storedRef $ signedSignature $ fromStored sidentity
+                       , BL.concat $ map (lazyLoadBytes . storedRef) $ collectStoredObjects $ wrappedLoad $ storedRef sidentity
                        ]) peer
            else putStrLn $ "Mismatched content"
 
diff --git a/src/Storage.hs b/src/Storage.hs
index c31230e..b3c2619 100644
--- a/src/Storage.hs
+++ b/src/Storage.hs
@@ -8,6 +8,7 @@ module Storage (
     Object(..), RecItem(..),
     serializeObject, deserializeObject, deserializeObjects,
     storeRawBytes, lazyLoadBytes,
+    collectObjects, collectStoredObjects,
 
     Head,
     headName, headRef, headObject,
@@ -73,6 +74,8 @@ import Data.Map (Map)
 import qualified Data.Map as M
 import Data.Maybe
 import Data.Ratio
+import Data.Set (Set)
+import qualified Data.Set as S
 import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Text.Encoding
@@ -263,6 +266,21 @@ deserializeObjects st bytes = do (obj, rest) <- deserializeObject st bytes
                                  (obj:) <$> deserializeObjects st rest
 
 
+collectObjects :: Object -> [Object]
+collectObjects obj = obj : map fromStored (fst $ collectOtherStored S.empty obj)
+
+collectStoredObjects :: Stored Object -> [Stored Object]
+collectStoredObjects obj = obj : (fst $ collectOtherStored S.empty $ fromStored obj)
+
+collectOtherStored :: Set Ref -> Object -> ([Stored Object], Set Ref)
+collectOtherStored seen (Rec items) = foldr helper ([], seen) $ map snd items
+    where helper (RecRef r) (xs, s) | r `S.notMember` s = let o = wrappedLoad r
+                                                              (xs', s') = collectOtherStored (S.insert r s) $ fromStored o
+                                                           in ((o : xs') ++ xs, s')
+          helper _          (xs, s) = (xs, s)
+collectOtherStored seen _ = ([], seen)
+
+
 data Head = Head String Ref
     deriving (Show)