Vanilla.hs

{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving #-}

module Data.Queue.Fuse.Vanilla (FuseVHeap, extractSingle, replace) where

import Data.Queue.Class
import Data.Queue.PQueue
import Data.Semigroup
import Prelude hiding (null)

data Assoc k v = A k v deriving (Show)
data VHeap k v = VH k v {-# UNPACK #-} !(PQueue (Assoc k v)) deriving (Show)
newtype FuseVHeap k v = FVH (Point (VHeap k v)) deriving (Monoid, Show)

instance Eq k => Eq (Assoc k v) where
	A x _ == A y _ = x == y

instance Ord k => Ord (Assoc k v) where
	A x _ `compare` A y _ = x `compare` y
	A x _ <= A y _ = x <= y

instance (Ord k, Semigroup v) => Semigroup (VHeap k v) where
	sappend = mergeVH
	sconcat [] = Nothing
	sconcat (q:qs) = Just (mergeVHs q qs)
	sconcat_ = mergeVHs

mergeVH :: (Ord k, Semigroup v) => Endo (VHeap k v)
VH k1 v1 q1 `mergeVH` VH k2 v2 q2 = case compare k1 k2 of
	LT	-> VH k1 v1 (A k2 v2 `insert` q)
	EQ	-> VH k1 (v1 `sappend` v2) q
	GT	-> VH k2 v2 (A k1 v1 `insert` q)
	where	q = q1 `merge` q2

data Acc k v = Ac k v [v] [Assoc k v] [PQueue (Assoc k v)]

mergeVHs :: (Ord k, Semigroup v) => VHeap k v -> [VHeap k v] -> VHeap k v
mergeVHs (VH k0 v0 q0) hs = case foldl merger (Ac k0 v0 [] [] []) hs of
	Ac k v vs as qs	-> VH k (v <<| vs) (as `insertAll` mergeAll qs)
	where	Ac k v vs as qs `merger` VH k' v' q' = let qs' = q':qs in case compare k' k of
			LT -> Ac k' v' [] (A k (v <<| vs):as) qs'
			EQ -> Ac k v (v':vs) as qs'
			GT -> Ac k v vs (A k' v':as) qs'
		{-# NOINLINE (<<|) #-}
		(<<|) = sconcat_

toVH :: (Ord k, Semigroup v) => PQueue (Assoc k v) -> Maybe (VHeap k v)
toVH = fmap toVH1 . extract where
	toVH1 (A k v, q) = let	toVH' vs q = case extract q of	Just (A k' v', q') | k == k'
									-> toVH' (v':vs) q'
								_	-> VH k (sconcat_ v vs) q
				in toVH' [] q

instance (Ord k, Semigroup v) => IQueue (FuseVHeap k v) where
	type QueueKey (FuseVHeap k v) = (k, v)
	empty = fvh Nothing
	merge = mappend
	mergeAll = mconcat

	singleton = fvh . Just . single
	fromList = fvh . fusing . map single

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

	extract (FVH (Pt h)) = fmap extract' h where
		extract' (VH k v q) = ((k, v), fvh $ toVH q)

single :: (Ord k, Semigroup v) => (k, v) -> VHeap k v
single (k, v) = VH k v empty

fvh :: (Ord k, Semigroup v) => Maybe (VHeap k v) -> FuseVHeap k v
fvh = FVH . Pt

extractSingle :: Ord k => FuseVHeap k v -> Maybe (k, v)
extractSingle (FVH (Pt (Just (VH k v q))))
	| null q	= Just (k, v)
extractSingle _ = Nothing

replace :: (Ord k, Semigroup v) => v -> FuseVHeap k v -> FuseVHeap k v
replace v (FVH (Pt (Just (VH k _ q)))) = fvh (Just (VH k v q))
replace _ q = q