Switched PSQueue to psqueues for strict spine.

3 files changed, 49 insertions, 11 deletions

diff --git a/src/Data/Wrapper/PSQ.hs b/src/Data/Wrapper/PSQ.hs
new file mode 100644
index 00000000..e8fa2d98
--- /dev/null
+++ b/src/Data/Wrapper/PSQ.hs
@@ -0,0 +1,41 @@
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE CPP #-}
+module Data.Wrapper.PSQ
+#if 0
+     ( module Data.PSQueue ) where ; import Data.PSQueue
+#else
+    ( module Data.Wrapper.PSQ , module OrdPSQ ) where
+
+import Data.OrdPSQ as OrdPSQ hiding (insert, map, singleton, minView)
+import qualified Data.OrdPSQ as OrdPSQ
+
+type PSQ k p = OrdPSQ k p ()
+type Binding k p = (k,p,())
+
+pattern (:->) :: k -> p -> Binding k p
+pattern k :-> p <- (k,p,()) where k :-> p = (k,p,())
+
+key :: Binding k v -> k
+key (k,p,v) = k
+{-# INLINE key #-}
+
+insert :: (Ord k, Ord p) => k -> p -> PSQ k p -> PSQ k p
+insert k p q = OrdPSQ.insert k p () q
+{-# INLINE insert #-}
+
+insertWith :: (Ord k, Ord p) => (p -> p -> p) -> k -> p -> PSQ k p -> PSQ k p
+insertWith f k p0 q = snd $ OrdPSQ.alter f' k q
+ where
+    f' (Just (p,())) = ((),Just (f p0 p, ()))
+    f' Nothing      = ((),Nothing)
+{-# INLINE insertWith #-}
+
+singleton :: (Ord k, Ord p) => k -> p -> PSQ k p
+singleton k p = OrdPSQ.singleton k p ()
+{-# INLINE singleton #-}
+
+minView :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p, PSQ k p)
+minView q = fmap (\(k,p,(),q') -> (k :-> p, q')) $ OrdPSQ.minView q
+{-# INLINE minView #-}
+
+#endif
diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index ea3b96f0..26bdeae9 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -12,7 +12,7 @@ import Data.List as L
 import Data.Maybe
 import Data.HashMap.Strict as HM
 import Data.Serialize
-import Data.PSQueue as PSQ
+import Data.Wrapper.PSQ as PSQ
 import Data.Time.Clock.POSIX
 import Data.ByteString (ByteString)
 import Data.Word
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index d64e415e..3c2e30aa 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -70,7 +70,7 @@ import Data.Functor.Identity
 import Data.List as L hiding (insert)
 import Data.Maybe
 import Data.Monoid
-import Data.PSQueue as PSQ
+import Data.Wrapper.PSQ as PSQ
 import Data.Serialize as S hiding (Result, Done)
 import qualified Data.Sequence as Seq
 import Data.Time
@@ -181,10 +181,6 @@ runRouting ping_node find_nodes timestamper = go
 -- | Timestamp - last time this node is pinged.
 type NodeEntry ip = Binding (NodeInfo ip) Timestamp
 
-instance (Serialize k, Serialize v) => Serialize (Binding k v) where
-  get = (:->) <$> get <*> get
-  put (k :-> v) = put k >> put v
-
 -- TODO instance Pretty where
 
 -- | Number of nodes in a bucket.
@@ -239,12 +235,13 @@ data Bucket ip = Bucket { bktNodes :: PSQ (NodeInfo ip) Timestamp
                         , bktQ     :: BucketQueue ip
                         } deriving (Show,Generic)
 
-instance (Eq ip, Serialize ip) => Serialize (Bucket ip)
+instance (Eq ip, Serialize ip) => Serialize (Bucket ip) where
+    get = Bucket . psqFromPairList <$> get <*> pure (runIdentity $ emptyQueue bucketQ)
+    put = put . psqToPairList . bktNodes
+
+psqFromPairList xs = PSQ.fromList $ map (\(a,b) -> a :-> b) xs
 
-instance (Serialize k, Serialize v, Ord k, Ord v)
-       => Serialize (PSQ k v) where
-  get = PSQ.fromList <$> get
-  put = put . PSQ.toList
+psqToPairList psq = map (\(a :-> b) -> (a,b)) $ PSQ.toList psq
 
 -- | Update interval, in seconds.
 delta :: NominalDiffTime