PQueue.hs

{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, GeneralizedNewtypeDeriving #-}
{-# OPTIONS -fno-warn-missing-methods #-}
{- | An efficient implementation of a priority queue.

The implementation of 'PQueue' is based on a /pairing heap/, a simple and efficient implementation of a general-purpose priority queue.  'PQueue' supports 'insert', 'merge', and 'peek' in constant time, and 'extract' and 'delete' in logarithmic time.
-}

module Data.Queue.PQueue (PQueue) where

import Data.Monoid
import Data.Maybe
import Control.Monad
import Data.Tree
import Data.Queue.Class
import GHC.Exts

data PHeap e = PH {elts :: {-# UNPACK #-} !Int, heap :: {-# UNPACK #-} !(Tree e)} deriving (Read, Show)
newtype PQueue e = PQ {getQ :: Maybe (PHeap e)} deriving (Monoid, Read, Show)

instance Ord e => Monoid (PHeap e) where
	PH n1 h1 `mappend` PH n2 h2 = PH (n1 + n2) (h1 `meld` h2)

instance Ord e => Queuelike (PQueue e) e where
	singleton x = mkQ 1 (single x)
	peek (PQ h) = fmap (rootLabel . heap) h
	delete (PQ h) = fmap (\ (PH (n+1) (Node _ ts)) -> mkQ n (fuser ts)) h
	isEmpty = isNothing . getQ
	size = maybe 0 elts . getQ
	fromList xs = case xs of	[] -> PQ Nothing
					_ -> case foldr (\ x (n, ys) -> (n+1, single x : ys)) (0, []) xs of
						(n, ys) -> mkQ n (fuser ys)
	{-# INLINE toList_ #-}
	toList_ (PQ h) = maybe [] (flatten . heap) h
	merge = mappend
	empty = PQ Nothing
	xs `insertAll` q = q `merge` fromList xs

{-# INLINE mkQ #-}
mkQ :: Int -> Tree e -> PQueue e
mkQ n t = PQ (Just (PH n t))

meld :: Ord e => Tree e -> Tree e -> Tree e
meld t1@(Node x1 ts1) t2@(Node x2 ts2) =
	if x1 > x2 then Node x2 (t1:ts2) else Node x1 (t2:ts1)

fuser :: Ord e => Forest e -> Tree e
fuser [t] = t
fuser l = fuser (fuser' l) where
	fuser' (t1:t2:ts) = t1 `meld` t2 : fuser' ts
	fuser' [t1] = [t1]
	fuser' [] = []

{-# INLINE single #-}
single :: e -> Tree e
single x = Node x []

mergeAllPH :: Ord e => [PQueue e] -> PQueue e
mergeAllPH qs = let (n, ts) = foldr (\ (PH n t) (m, ts) -> (n+m, t:ts)) (0, []) [ph | PQ (Just ph) <- qs] in
	if n == 0 then PQ Nothing else mkQ n (fuser ts)

{-# NOINLINE [0] flattenFB #-}
flattenFB :: Tree a -> (a -> b -> b) -> b -> b
flattenFB (Node x ts) c n = x `c` foldr (flip flattenFB c) n ts

{-# RULES
	"flatten" [~1] forall t . flatten t = build (flattenFB t);
	"flattenList" [1] forall t . build (flattenFB t) = flatten t;
	#-}