SkewHeap.hs

{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving #-}
{-# OPTIONS -fspec-constr -fspec-constr-count=8 -fspec-constr-threshold=50 #-}

module Data.Queue.Fuse.SkewHeap (FuseSHeap, extractSingle, replace) where

import Data.Semigroup
import Data.Maybe
import Data.Queue.Class
import Data.Queue.QueueHelpers
import Data.Queue.Fuse.SplitList

data SHeap k v = SH k v (MSHeap k v) (MSHeap k v) deriving (Show)
type MSHeap k v = Point (SHeap k v)
newtype FuseSHeap k v = FSH (MSHeap k v) deriving (Monoid, Show)

type Comparator e = e -> e -> Ordering

instance (Ord k, Semigroup v) => Semigroup (SHeap k v) where
	sappend = mergeSH
	sconcat = mergeSHs

{-# INLINE mergeSH #-}
mergeSH :: (Ord k, Semigroup v) => Endo (SHeap k v)
mergeSH = (mergeFuncs compare sappend)

mergeFuncs :: Comparator k -> Endo v -> Endo (SHeap k v)
mergeFuncs cmp (><) = (>!<) where
	(>?<) = endoPoint (>!<)
	h1@(SH k1 v1 l1 r1) >!< h2@(SH k2 v2 l2 r2) = case cmp k1 k2 of
		LT	-> SH k1 v1 (pJust h2 >?< r1) l1
		EQ	-> SH k1 (v1 >< v2) (l1 >?< r2) (l2 >?< r1)
		GT	-> SH k2 v2 (pJust h1 >?< r2) l2

data Merge k v = Mrg k v [v] {-# UNPACK #-} !(Split (SHeap k v)) 

{-# INLINE sh #-}
sh :: k -> v -> Maybe (SHeap k v) -> Maybe (SHeap k v) -> SHeap k v
sh k v l r = SH k v (Pt l) (Pt r)

{-# INLINE mergeSHs #-}
mergeSHs :: (Ord k, Semigroup v) => Fusion (SHeap k v)
mergeSHs = mergeSHs0 compare sappend sconcat

mergeSHs0 :: Comparator k -> Endo v -> Fusion v -> Fusion (SHeap k v)
mergeSHs0 cmp (><) cat0 = mergeSHs' where
	mergeSHs' [] = Nothing
	mergeSHs' (SH k v l r:hs0) = Just $ mrgToSH $ foldl merger (Mrg k v [] (consLR l r emptySpl)) hs0
	consLR l r spl = [q | Pt (Just q) <- [l,r]] <<| spl
	(>!<) = mergeFuncs cmp (><)
	v `cat` vs = fromJust (cat0 (v:vs))
	{-# INLINE mrgToSH #-}
	mrgToSH (Mrg k v vs subTs) = let (ls, rs) = split subTs in 
		sh k (v `cat` vs) (mergeSHs' ls) (mergeSHs' rs)
	mrg@(Mrg k0 v0 vs0 subTs) `merger` h@(SH k v l r) = case cmp k k0 of
		LT	-> let (ls, rs) = split subTs in
			Mrg k v [] $ singleSpl (mrgToSH mrg)
		EQ	-> Mrg k0 v0 (v:vs0) (consLR l r subTs)
		GT	-> Mrg k0 v0 vs0 (h <| subTs)

instance (Ord k, Semigroup v) => IQueue (FuseSHeap k v) where
	type QueueKey (FuseSHeap k v) = (k, v)
	empty = mempty
	merge = mappend
	mergeAll = mconcat

	singleton = FSH . pJust . single
	fromList = fsh . sconcat . map single

	extract (FSH (Pt h)) = fmap extract' h where
		extract' (SH k v l r) = ((k,v), FSH (l `mappend` r))

	null (FSH (Pt Nothing)) = True
	null _ = False

	toList_ (FSH (Pt h)) = maybe [] (unfoldList unHeap) h where
		unHeap (SH k v l r) = ((k, v), [t | Pt (Just t) <- [l,r]])

single :: (Ord k, Semigroup v) => (k, v) -> SHeap k v
single (k, v) = SH k v mempty mempty

extractSingle :: FuseSHeap k v -> Maybe (k, v)
extractSingle (FSH (Pt (Just (SH k v (Pt Nothing) (Pt Nothing)))))
	= Just (k, v)
extractSingle _ = Nothing

replace :: v -> FuseSHeap k v -> FuseSHeap k v
replace v (FSH (Pt (Just (SH k _ l r)))) = FSH (pJust (SH k v l r))

fsh = FSH . Pt