1 files changed, 34 insertions, 23 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