module Data.Queue.FibQueue (FQueue) where
import Data.Queue.Class
import Control.Monad.Array
import Data.Tree(Tree(..))
import Data.Maybe
import Control.Monad
import GHC.Exts(build)
import Prelude hiding (getContents)
data Rk e = Rk {rk :: !Int, lab :: e}
type RkTree e = Tree (Rk e)
data FQueue e = FQueue {elts :: !Int, maxRank :: !Int, heap :: [RkTree e]}
instance Ord e => Queuelike (FQueue e) e where
FQueue n1 r1 h1 `merge` FQueue n2 r2 h2 = FQueue (n1 + n2) (max r1 r2) (case (h1, h2) of
((treeMin -> x1):_, (treeMin -> x2):_) -> if x1 <= x2 then h1 ++ h2 else h2 ++ h1
(_, []) -> h1
([], _) -> h2)
empty = FQueue 0 0 []
singleton x = FQueue 1 0 [Node (Rk 0 x) []]
toList_ = concatMap (map lab . flatten) . heap
size = elts
peek = liftM treeMin . listToMaybe . heap
delete (FQueue n mR (Node (Rk _ x) ts : tss)) = Just $ rebuild (MA (n1) mR (mapM_ meld tss >> mapM_ meld ts))
delete _ = Nothing
treeMin :: RkTree e -> e
treeMin (Node (Rk _ x) _) = x
flatten :: Tree e -> [e]
flatten t = build (\ c n -> flatten' c t n) where
flatten' c (Node x ts) n = x `c` foldr (flatten' c) n ts
meldTree :: Ord e => RkTree e -> RkTree e -> RkTree e
t1@(Node (Rk d x1) ts1) `meldTree` t2@(Node (Rk _ x2) ts2)
| x1 <= x2 = Node (Rk (d+1) x1) (t2:ts1)
| otherwise = Node (Rk (d+1) x2) (t1:ts2)
meld :: Ord e => RkTree e -> ArrayM s (Maybe (RkTree e)) ()
meld t@(rk . rootLabel -> d) =
ensureSize (d+2) >> readAt d >>= maybe (writeAt d (Just t)) (\ t' -> writeAt d Nothing >> meld (t `meldTree` t'))
extractMin :: Ord e => [Maybe (RkTree e)] -> (Int, [RkTree e])
extractMin ls = case foldr exM (Nothing, 0, []) ls of (mi, rk, ts) -> maybe (0, []) ((,) rk . (:ts)) mi ; where
exM Nothing p = p
exM (Just t@(Node (Rk d x) _)) (mi, rk, ts) = let rk' = max d rk in maybe (Just t, rk', ts)
(\ t'@(lab . rootLabel -> y) -> if x <= y then (Just t, rk', t':ts) else (Just t', rk', t:ts)) mi
rebuild :: Ord e => MergeAccum e -> FQueue e
rebuild (MA n mR melder) = runArrayM mR Nothing $ melder >> liftM ((\ (mR', h') -> FQueue n mR' h') . extractMin) getContents
data MergeAccum e = MA !Int !Int (forall s . ArrayM s (Maybe (RkTree e)) ())
mergeAllFH :: Ord e => [FQueue e] -> FQueue e
mergeAllFH qs = rebuild (foldr merger (MA 0 0 (return ())) qs) where
merger :: Ord e => FQueue e -> MergeAccum e -> MergeAccum e
merger (FQueue n r ts) (MA m mR toMerge) = MA (n+m) (max r mR) (mapM_ meld ts >> toMerge)