2 files changed, 53 insertions, 30 deletions

diff --git a/src/Data/MinMaxPSQ.hs b/src/Data/MinMaxPSQ.hs
index f385f258..a48e62f9 100644
--- a/src/Data/MinMaxPSQ.hs
+++ b/src/Data/MinMaxPSQ.hs
@@ -3,56 +3,61 @@ module Data.MinMaxPSQ where
 
 import Data.Ord
 import qualified Data.Wrapper.PSQ as PSQ
-         ;import Data.Wrapper.PSQ as PSQ hiding (insert, null, size)
+         ;import Data.Wrapper.PSQ as PSQ hiding (insert, insert', null, size)
 import Prelude hiding (null, take)
 
-data MinMaxPSQ k p = MinMaxPSQ !(PSQ k p) !(PSQ k (Down p))
+data MinMaxPSQ' k p v = MinMaxPSQ !(PSQ' k p v) !(PSQ' k (Down p) v)
+type MinMaxPSQ k p = MinMaxPSQ' k p ()
 
-empty :: MinMaxPSQ k p
+empty :: MinMaxPSQ' k p v
 empty = MinMaxPSQ PSQ.empty PSQ.empty
 
-null :: MinMaxPSQ k p -> Bool
+null :: MinMaxPSQ' k p v -> Bool
 null (MinMaxPSQ nq xq) = PSQ.null nq
 
-size :: MinMaxPSQ k p -> Int
+size :: MinMaxPSQ' k p v -> Int
 size (MinMaxPSQ nq xq) = PSQ.size nq
 
-toList :: (PSQKey k, Ord p) => MinMaxPSQ k p -> [Binding k p]
+toList :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> [Binding' k p v]
 toList (MinMaxPSQ nq xq) = PSQ.toList nq
 
-fromList :: (PSQKey k, Ord p) => [Binding k p] -> MinMaxPSQ k p
+fromList :: (PSQKey k, Ord p) => [Binding' k p v] -> MinMaxPSQ' k p v
 fromList kps = MinMaxPSQ (PSQ.fromList kps)
-                         (PSQ.fromList $ map (\(k :-> p) -> (k :-> Down p)) kps)
+                         (PSQ.fromList $ map (\(Binding k v p) -> Binding k v (Down p)) kps)
 
-findMin :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
+findMin :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v)
 findMin (MinMaxPSQ nq xq) = PSQ.findMin nq
 
-findMax :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
-findMax (MinMaxPSQ nq xq) = fmap (\(k :-> Down p) -> k :-> p) $ PSQ.findMin xq
+findMax :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v)
+findMax (MinMaxPSQ nq xq) = fmap (\(Binding k v (Down p)) -> Binding k v p) $ PSQ.findMin xq
 
 insert :: (PSQKey k, Ord p) => k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
 insert k p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert k p        nq)
                                          (PSQ.insert k (Down p) xq)
 
-delete :: (PSQKey k, Ord p) => k -> MinMaxPSQ k p -> MinMaxPSQ k p
+insert' :: (PSQKey k, Ord p) => k -> v -> p -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
+insert' k v p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert' k v p        nq)
+                                            (PSQ.insert' k v (Down p) xq)
+
+delete :: (PSQKey k, Ord p) => k -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
 delete k (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.delete k nq) (PSQ.delete k xq)
 
-deleteMin :: (PSQKey k, Ord p) => MinMaxPSQ k p -> MinMaxPSQ k p
+deleteMin :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> MinMaxPSQ' k p v
 deleteMin (MinMaxPSQ nq xq) = case PSQ.minView nq of
-    Just (k :-> _, nq') -> MinMaxPSQ nq' (PSQ.delete k xq)
-    Nothing             -> MinMaxPSQ nq xq
+    Just (Binding k _ _, nq') -> MinMaxPSQ nq' (PSQ.delete k xq)
+    Nothing                   -> MinMaxPSQ nq xq
 
-deleteMax :: (PSQKey k, Ord p) => MinMaxPSQ k p -> MinMaxPSQ k p
+deleteMax :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> MinMaxPSQ' k p v
 deleteMax (MinMaxPSQ nq xq) = case PSQ.minView xq of
-    Just (k :-> _, xq') -> MinMaxPSQ (PSQ.delete k nq) xq'
-    Nothing             -> MinMaxPSQ nq xq
+    Just (Binding k _ _, xq') -> MinMaxPSQ (PSQ.delete k nq) xq'
+    Nothing                   -> MinMaxPSQ nq xq
 
-minView :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
-minView (MinMaxPSQ nq xq) = fmap (\(k :-> p, nq') -> (k :-> p, MinMaxPSQ nq' (PSQ.delete k xq)))
+minView :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v, MinMaxPSQ' k p v)
+minView (MinMaxPSQ nq xq) = fmap (\(Binding k v p, nq') -> (Binding k v p, MinMaxPSQ nq' (PSQ.delete k xq)))
                             $ PSQ.minView nq
 
-maxView :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
-maxView (MinMaxPSQ nq xq) = fmap (\(k :-> Down p, xq') -> (k :-> p, MinMaxPSQ (PSQ.delete k nq) xq'))
+maxView :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v, MinMaxPSQ' k p v)
+maxView (MinMaxPSQ nq xq) = fmap (\(Binding k v (Down p), xq') -> (Binding k v p, MinMaxPSQ (PSQ.delete k nq) xq'))
                             $ PSQ.minView xq
 
 -- | Maintains a bounded 'MinMaxPSQ' by deleting the maximum element if the
@@ -60,8 +65,14 @@ maxView (MinMaxPSQ nq xq) = fmap (\(k :-> Down p, xq') -> (k :-> p, MinMaxPSQ (P
 insertTake :: (PSQKey k, Ord p) => Int -> k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
 insertTake n k p q = take n $ insert k p q
 
+-- | Maintains a bounded 'MinMaxPSQ\'' by deleting the maximum element if the
+-- insertion would cause the queue to have too many elements.
+insertTake' :: (PSQKey k, Ord p) => Int -> k -> v -> p -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
+insertTake' n k v p q = take n $ insert' k v p q
+
+
 -- | Truncate the 'MinMaxPSQ' to the given number of lowest priority elements.
-take :: (PSQKey k, Ord p) => Int -> MinMaxPSQ k p -> MinMaxPSQ k p
+take :: (PSQKey k, Ord p) => Int -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
 take !n !q | (size q <= n) = q
            | null q        = q
            | otherwise     = take n $ deleteMax q
diff --git a/src/Data/Wrapper/PSQ.hs b/src/Data/Wrapper/PSQ.hs
index 54365a1d..87648e84 100644
--- a/src/Data/Wrapper/PSQ.hs
+++ b/src/Data/Wrapper/PSQ.hs
@@ -26,19 +26,27 @@ import Data.HashPSQ as HashPSQ hiding (insert, map, singleton, minView)
 import qualified Data.HashPSQ as Q
 import Data.Hashable
 
-type PSQ k p = HashPSQ k p ()
-type Binding k p = (k,p,())
+type PSQ' k p v = HashPSQ k p v
+type PSQ k p    = PSQ' k p ()
+
+type Binding' k p v = (k,p,v)
+type Binding k p    = Binding' k p ()
 
 type PSQKey k = (Hashable k, Ord k)
 
 pattern (:->) :: k -> p -> Binding k p
-pattern k :-> p <- (k,p,()) where k :-> p = (k,p,())
+pattern k :-> p <- (k,p,_) where k :-> p = (k,p,())
+
+-- I tried defining (::->) :: (k,v) -> p -> Binding' k p v
+-- but no luck...
+pattern Binding :: k -> v -> p -> Binding' k p v
+pattern Binding k v p <- (k,p,v) where Binding k v p = (k,p,v)
 
-key :: Binding k p -> k
+key :: (k,p,v) -> k
 key (k,p,v) = k
 {-# INLINE key #-}
 
-prio :: Binding k p -> p
+prio :: (k,p,v) -> p
 prio (k,p,v) = p
 {-# INLINE prio #-}
 
@@ -46,6 +54,10 @@ insert :: (PSQKey k, Ord p) => k -> p -> PSQ k p -> PSQ k p
 insert k p q = Q.insert k p () q
 {-# INLINE insert #-}
 
+insert' :: (PSQKey k, Ord p) => k -> v -> p -> PSQ' k p v -> PSQ' k p v
+insert' k v p q = Q.insert k p v q
+{-# INLINE insert' #-}
+
 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
@@ -57,8 +69,8 @@ singleton :: (PSQKey k, Ord p) => k -> p -> PSQ k p
 singleton k p = Q.singleton k p ()
 {-# INLINE singleton #-}
 
-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
+minView :: (PSQKey k, Ord p) => PSQ' k p v -> Maybe (Binding' k p v, PSQ' k p v)
+minView q = fmap (\(k,p,v,q') -> (Binding k v p, q')) $ Q.minView q
 {-# INLINE minView #-}
 
 #endif