{-# LANGUAGE TypeOperators, TypeFamilies #-}

-- | A small collection of specialized 'Int'-indexed priority queues dealing with both untagged 'Int's and association pairs with 'Int' keys.  The implementation is a simple bootstrap from 'IntMap'.  (Note: Duplicate keys /will/ be counted separately.  No guarantees are made on the order in which values associated with equal keys are returned.)
module Data.Queue.IntQueue (IntQueue, IntAssocQueue) where

import Data.Queue.Class
import Data.Maybe
import qualified Data.IntMap as IM

-- | A 'Queuelike' type with @QueueKey IntQueue ~ Int@.
newtype IntQueue = IQ (IM.IntMap Int)

-- | A 'Queuelike' type with @QueueKey (IntAssocQueue e) ~ e :-> Int@.
newtype IntAssocQueue e = IAQ (IM.IntMap [e])

instance Queuelike IntQueue where
	type QueueKey IntQueue = Int
	x `insert` IQ m = IQ (IM.alter (Just . maybe 1 (+1)) x m)
	xs `insertAll` q = q `merge` fromList xs
	extract (IQ m) = fmap (\ ((k, ct), m') -> (k, IQ (if ct > 1 then IM.insert k (ct - 1) m' else m'))) (IM.minViewWithKey m)
	isEmpty (IQ m) = IM.null m
	empty = IQ IM.empty
	singleton x = IQ (IM.singleton x 1)
	fromList xs = IQ (IM.fromListWith (+) [(x, 1) | x <- xs])
	IQ q1 `merge` IQ q2 = IQ (IM.unionWith (+) q1 q2)
	mergeAll qs = IQ (IM.unionsWith (+) [m | IQ m <- qs])

instance Queuelike (IntAssocQueue e) where
	type QueueKey (IntAssocQueue e) = e :-> Int
	(v :-> k) `insert` IAQ m = IAQ (IM.alter (Just . maybe [v] (v:)) k m)
	xs `insertAll` q = q `merge` fromList xs
	extract (IAQ m) = fmap (\ ((k, v:vs), m') -> (v :-> k, IAQ (if null vs then m' else IM.insert k vs m'))) (IM.minViewWithKey m)
	isEmpty (IAQ m) = IM.null m
	empty = IAQ IM.empty
	singleton (v :-> k) = IAQ (IM.singleton k [v])
	fromList xs = IAQ (IM.fromListWith (++) [(k, [v]) | (v :-> k) <- xs])
	IAQ q1 `merge` IAQ q2 = IAQ (IM.unionWith (++) q1 q2)
	mergeAll qs = IAQ (IM.unionsWith (++) [m | IAQ m <- qs])