1 files changed, 44 insertions, 14 deletions

diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index d7c92e35..979dbb62 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -1,14 +1,19 @@
 module Network.BitTorrent.DHT.ContactInfo
        ( PeerStore
        , Network.BitTorrent.DHT.ContactInfo.lookup
-       , Network.BitTorrent.DHT.ContactInfo.insert
+       , Network.BitTorrent.DHT.ContactInfo.insertPeer
+       , knownSwarms
        ) where
 
+import Control.Applicative
 import Data.Default
 import Data.List as L
 import Data.Maybe
 import Data.HashMap.Strict as HM
 import Data.Serialize
+import Data.PSQueue as PSQ
+import Data.Time.Clock.POSIX
+import Data.ByteString (ByteString)
 
 import Data.Torrent
 import Network.BitTorrent.Address
@@ -107,21 +112,43 @@ import Network.BitTorrent.Address
 
 -- | Storage used to keep track a set of known peers in client,
 -- tracker or DHT sessions.
-newtype PeerStore ip = PeerStore (HashMap InfoHash [PeerAddr ip])
+newtype PeerStore ip = PeerStore (HashMap InfoHash (SwarmData ip))
+
+type Timestamp = POSIXTime
+
+data SwarmData ip = SwarmData
+    { peers :: PSQ (PeerAddr ip) Timestamp
+    , name :: Maybe ByteString
+    }
+
+knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
+knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
+
+swarmSinglton :: Ord ip => PeerAddr ip -> SwarmData ip
+swarmSinglton a = SwarmData
+    { peers = PSQ.singleton a 0
+    , name = Nothing }
+
+swarmInsert :: Ord ip => SwarmData ip -> SwarmData ip -> SwarmData ip
+swarmInsert old new = SwarmData
+    { peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith (\p _ -> p) a t q) (peers old) (PSQ.toList $ peers new)
+    , name = name new <|> name old -- TODO: decodeUtf8' check
+    }
+
 
 -- | Empty store.
 instance Default (PeerStore a) where
   def = PeerStore HM.empty
   {-# INLINE def #-}
 
--- | Monoid under union operation.
-instance Eq a => Monoid (PeerStore a) where
-  mempty  = def
-  {-# INLINE mempty #-}
-
-  mappend (PeerStore a) (PeerStore b) =
-    PeerStore (HM.unionWith L.union a b)
-  {-# INLINE mappend #-}
+-- -- | Monoid under union operation.
+-- instance Eq a => Monoid (PeerStore a) where
+--   mempty  = def
+--   {-# INLINE mempty #-}
+-- 
+--   mappend (PeerStore a) (PeerStore b) =
+--     PeerStore (HM.unionWith L.union a b)
+--   {-# INLINE mappend #-}
 
 -- | Can be used to store peers between invocations of the client
 -- software.
@@ -130,9 +157,12 @@ instance Serialize (PeerStore a) where
   put = undefined
 
 -- | Used in 'get_peers' DHT queries.
-lookup :: InfoHash -> PeerStore a -> [PeerAddr a]
-lookup ih (PeerStore m) = fromMaybe [] $ HM.lookup ih m
+lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
+lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
 
 -- | Used in 'announce_peer' DHT queries.
-insert :: Eq a => InfoHash -> PeerAddr a -> PeerStore a -> PeerStore a
-insert ih a (PeerStore m) = PeerStore (HM.insertWith L.union ih [a] m)
+insertPeer :: Ord a => InfoHash -> Maybe ByteString -> PeerAddr a -> PeerStore a -> PeerStore a
+insertPeer ih name a (PeerStore m) = PeerStore (HM.insertWith swarmInsert ih a' m)
+ where
+    a' = SwarmData { peers = PSQ.singleton a 0
+                   , name = name }