2 files changed, 29 insertions, 12 deletions

diff --git a/src/Data/Wrapper/PSQ.hs b/src/Data/Wrapper/PSQ.hs
index 2c08011b..7e14008b 100644
--- a/src/Data/Wrapper/PSQ.hs
+++ b/src/Data/Wrapper/PSQ.hs
@@ -1,17 +1,34 @@
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE ConstraintKinds #-}
 module Data.Wrapper.PSQ
 #if 0
-     ( module Data.PSQueue ) where ; import Data.PSQueue
+     ( module Data.PSQueue ) where
+
+import Data.PSQueue hiding (foldr, foldl)
+import qualified Data.PSQueue as PSQueue
+
+-- | Wrapper over PSQueue-style foldr to implement a psqueues-style interface.
+fold' :: (Ord k, Ord p) => (k -> p -> () -> a -> a) -> a -> PSQ k p -> a
+fold' f a q = PSQueue.foldr f' a q
+ where
+    f' (k :-> prio) x = f k prio () x
+
 #else
     ( module Data.Wrapper.PSQ , module OrdPSQ ) where
 
 import Data.OrdPSQ as OrdPSQ hiding (insert, map, singleton, minView)
-import qualified Data.OrdPSQ as OrdPSQ
+import qualified Data.OrdPSQ as Q
+
+import Data.HashPSQ as HashPSQ hiding (insert, map, singleton, minView)
+import qualified Data.HashPSQ as HashPSQ
+import Data.Hashable
 
 type PSQ k p = OrdPSQ k p ()
 type Binding k p = (k,p,())
 
+type PSQKey k = (Ord k)
+
 pattern (:->) :: k -> p -> Binding k p
 pattern k :-> p <- (k,p,()) where k :-> p = (k,p,())
 
@@ -23,23 +40,23 @@ prio :: Binding k p -> p
 prio (k,p,v) = p
 {-# INLINE prio #-}
 
-insert :: (Ord k, Ord p) => k -> p -> PSQ k p -> PSQ k p
-insert k p q = OrdPSQ.insert k p () q
+insert :: (PSQKey k, Ord p) => k -> p -> PSQ k p -> PSQ k p
+insert k p q = Q.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
+insertWith :: (PSQKey k, Ord p) => (p -> p -> p) -> k -> p -> PSQ k p -> PSQ k p
+insertWith f k p0 q = snd $ Q.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 ()
+singleton :: (PSQKey k, Ord p) => k -> p -> PSQ k p
+singleton k p = Q.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
+minView :: (PSQKey k, Ord p) => PSQ k p -> Maybe (Binding k p, PSQ k p)
+minView q = fmap (\(k,p,(),q') -> (k :-> p, q')) $ Q.minView q
 {-# INLINE minView #-}
 
 #endif
diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 273cf9dc..d380c6f2 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -248,10 +248,10 @@ instance (Eq ip, Ord (NodeId), Serialize (NodeId), Serialize ip, Serialize u) =>
 
 #endif
 
-psqFromPairList :: (Ord p, Ord k) => [(k, p)] -> OrdPSQ k p ()
+psqFromPairList :: (Ord p, Ord k) => [(k, p)] -> PSQ k p
 psqFromPairList xs = PSQ.fromList $ map (\(a,b) -> a :-> b) xs
 
-psqToPairList :: OrdPSQ t t1 () -> [(t, t1)]
+psqToPairList :: PSQ t t1 -> [(t, t1)]
 psqToPairList psq = map (\(a :-> b) -> (a,b)) $ PSQ.toList psq
 
 -- | Update interval, in seconds.