{-# LANGUAGE CPP #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module HaskellWorks.Data.PriorityQueue.Strict
( PQueue
, empty
, singleton
, union
, insert
, add
, fromList
, null
, minView
, minViewWithKey
, takeWithKeys
, take
) where
import Control.Arrow ((***))
import Control.DeepSeq (NFData)
import Data.Bifunctor (first)
import Data.Foldable (Foldable (foldMap))
import GHC.Generics (Generic)
import HaskellWorks.Data.FingerTree.Strict (FingerTree, Measured (..), ViewL (..), (<|), (><), (|>))
import Prelude hiding (null, take)
import qualified Data.Semigroup as S
import qualified HaskellWorks.Data.FingerTree.Strict as FT
data Entry k v = Entry k v
instance Functor (Entry k) where
fmap :: (a -> b) -> Entry k a -> Entry k b
fmap a -> b
f (Entry k
k a
v) = k -> b -> Entry k b
forall k v. k -> v -> Entry k v
Entry k
k (a -> b
f a
v)
instance Foldable (Entry k) where
foldMap :: (a -> m) -> Entry k a -> m
foldMap a -> m
f (Entry k
_ a
v) = a -> m
f a
v
data Prio k v = NoPrio | Prio k v
appendPrio :: Ord k => Prio k v -> Prio k v -> Prio k v
appendPrio :: Prio k v -> Prio k v -> Prio k v
appendPrio Prio k v
x Prio k v
NoPrio = Prio k v
x
appendPrio Prio k v
NoPrio Prio k v
y = Prio k v
y
appendPrio x :: Prio k v
x@(Prio k
kx v
_) y :: Prio k v
y@(Prio k
ky v
_) = if k
kx k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
ky then Prio k v
x else Prio k v
y
{-# INLINE appendPrio #-}
instance Ord k => S.Semigroup (Prio k v) where
<> :: Prio k v -> Prio k v -> Prio k v
(<>) = Prio k v -> Prio k v -> Prio k v
forall k v. Ord k => Prio k v -> Prio k v -> Prio k v
appendPrio
{-# INLINE (<>) #-}
instance Ord k => Monoid (Prio k v) where
mempty :: Prio k v
mempty = Prio k v
forall k v. Prio k v
NoPrio
{-# INLINE mempty #-}
instance Ord k => Measured (Prio k v) (Entry k v) where
measure :: Entry k v -> Prio k v
measure (Entry k
k v
v) = k -> v -> Prio k v
forall k v. k -> v -> Prio k v
Prio k
k v
v
newtype PQueue k v = PQueue (FingerTree (Prio k v) (Entry k v))
instance Ord k => Functor (PQueue k) where
fmap :: (a -> b) -> PQueue k a -> PQueue k b
fmap a -> b
f (PQueue FingerTree (Prio k a) (Entry k a)
xs) = FingerTree (Prio k b) (Entry k b) -> PQueue k b
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue ((Entry k a -> Entry k b)
-> FingerTree (Prio k a) (Entry k a)
-> FingerTree (Prio k b) (Entry k b)
forall v1 a1 v2 a2.
(Measured v1 a1, Measured v2 a2) =>
(a1 -> a2) -> FingerTree v1 a1 -> FingerTree v2 a2
FT.fmap' ((a -> b) -> Entry k a -> Entry k b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) FingerTree (Prio k a) (Entry k a)
xs)
instance Ord k => Foldable (PQueue k) where
foldMap :: (a -> m) -> PQueue k a -> m
foldMap a -> m
f PQueue k a
q = case PQueue k a -> Maybe (a, PQueue k a)
forall k v. Ord k => PQueue k v -> Maybe (v, PQueue k v)
minView PQueue k a
q of
Maybe (a, PQueue k a)
Nothing -> m
forall a. Monoid a => a
mempty
Just (a
v, PQueue k a
q') -> a -> m
f a
v m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` (a -> m) -> PQueue k a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f PQueue k a
q'
instance Ord k => S.Semigroup (PQueue k v) where
<> :: PQueue k v -> PQueue k v -> PQueue k v
(<>) = PQueue k v -> PQueue k v -> PQueue k v
forall k v. Ord k => PQueue k v -> PQueue k v -> PQueue k v
union
{-# INLINE (<>) #-}
instance Ord k => Monoid (PQueue k v) where
mempty :: PQueue k v
mempty = PQueue k v
forall k v. PQueue k v
empty
{-# INLINE mempty #-}
mappend :: PQueue k v -> PQueue k v -> PQueue k v
mappend = PQueue k v -> PQueue k v -> PQueue k v
forall a. Semigroup a => a -> a -> a
(<>)
{-# INLINE mappend #-}
empty :: PQueue k v
empty :: PQueue k v
empty = FingerTree (Prio k v) (Entry k v) -> PQueue k v
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue FingerTree (Prio k v) (Entry k v)
forall v a. FingerTree v a
FT.empty
singleton :: Ord k => k -> v -> PQueue k v
singleton :: k -> v -> PQueue k v
singleton k
k v
v = FingerTree (Prio k v) (Entry k v) -> PQueue k v
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue (Entry k v -> FingerTree (Prio k v) (Entry k v)
forall a v. a -> FingerTree v a
FT.singleton (k -> v -> Entry k v
forall k v. k -> v -> Entry k v
Entry k
k v
v))
insert :: Ord k => k -> v -> PQueue k v -> PQueue k v
insert :: k -> v -> PQueue k v -> PQueue k v
insert k
k v
v (PQueue FingerTree (Prio k v) (Entry k v)
q) = FingerTree (Prio k v) (Entry k v) -> PQueue k v
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue (k -> v -> Entry k v
forall k v. k -> v -> Entry k v
Entry k
k v
v Entry k v
-> FingerTree (Prio k v) (Entry k v)
-> FingerTree (Prio k v) (Entry k v)
forall v a. Measured v a => a -> FingerTree v a -> FingerTree v a
<| FingerTree (Prio k v) (Entry k v)
q)
add :: Ord k => k -> v -> PQueue k v -> PQueue k v
add :: k -> v -> PQueue k v -> PQueue k v
add k
k v
v (PQueue FingerTree (Prio k v) (Entry k v)
q) = FingerTree (Prio k v) (Entry k v) -> PQueue k v
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue (FingerTree (Prio k v) (Entry k v)
q FingerTree (Prio k v) (Entry k v)
-> Entry k v -> FingerTree (Prio k v) (Entry k v)
forall v a. Measured v a => FingerTree v a -> a -> FingerTree v a
|> k -> v -> Entry k v
forall k v. k -> v -> Entry k v
Entry k
k v
v)
union :: Ord k => PQueue k v -> PQueue k v -> PQueue k v
union :: PQueue k v -> PQueue k v -> PQueue k v
union (PQueue FingerTree (Prio k v) (Entry k v)
xs) (PQueue FingerTree (Prio k v) (Entry k v)
ys) = FingerTree (Prio k v) (Entry k v) -> PQueue k v
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue (FingerTree (Prio k v) (Entry k v)
xs FingerTree (Prio k v) (Entry k v)
-> FingerTree (Prio k v) (Entry k v)
-> FingerTree (Prio k v) (Entry k v)
forall v a.
Measured v a =>
FingerTree v a -> FingerTree v a -> FingerTree v a
>< FingerTree (Prio k v) (Entry k v)
ys)
fromList :: Ord k => [(k, v)] -> PQueue k v
fromList :: [(k, v)] -> PQueue k v
fromList = ((k, v) -> PQueue k v -> PQueue k v)
-> PQueue k v -> [(k, v)] -> PQueue k v
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((k -> v -> PQueue k v -> PQueue k v)
-> (k, v) -> PQueue k v -> PQueue k v
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> v -> PQueue k v -> PQueue k v
forall k v. Ord k => k -> v -> PQueue k v -> PQueue k v
insert) PQueue k v
forall k v. PQueue k v
empty
null :: PQueue k v -> Bool
null :: PQueue k v -> Bool
null (PQueue FingerTree (Prio k v) (Entry k v)
q) = FingerTree (Prio k v) (Entry k v) -> Bool
forall v a. FingerTree v a -> Bool
FT.null FingerTree (Prio k v) (Entry k v)
q
minView :: Ord k => PQueue k v -> Maybe (v, PQueue k v)
minView :: PQueue k v -> Maybe (v, PQueue k v)
minView PQueue k v
q = (((k, v), PQueue k v) -> (v, PQueue k v))
-> Maybe ((k, v), PQueue k v) -> Maybe (v, PQueue k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((k, v) -> v) -> ((k, v), PQueue k v) -> (v, PQueue k v)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (k, v) -> v
forall a b. (a, b) -> b
snd) (PQueue k v -> Maybe ((k, v), PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe ((k, v), PQueue k v)
minViewWithKey PQueue k v
q)
takeWithKeys :: Ord k => Int -> PQueue k v -> ([(k, v)], PQueue k v)
takeWithKeys :: Int -> PQueue k v -> ([(k, v)], PQueue k v)
takeWithKeys = [(k, v)] -> Int -> PQueue k v -> ([(k, v)], PQueue k v)
forall k v.
Ord k =>
[(k, v)] -> Int -> PQueue k v -> ([(k, v)], PQueue k v)
go []
where go :: Ord k => [(k, v)] -> Int -> PQueue k v -> ([(k, v)], PQueue k v)
go :: [(k, v)] -> Int -> PQueue k v -> ([(k, v)], PQueue k v)
go [(k, v)]
as Int
n PQueue k v
q | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = case PQueue k v -> Maybe ((k, v), PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe ((k, v), PQueue k v)
minViewWithKey PQueue k v
q of
Just ((k, v)
a, PQueue k v
r) -> [(k, v)] -> Int -> PQueue k v -> ([(k, v)], PQueue k v)
forall k v.
Ord k =>
[(k, v)] -> Int -> PQueue k v -> ([(k, v)], PQueue k v)
go ((k, v)
a(k, v) -> [(k, v)] -> [(k, v)]
forall a. a -> [a] -> [a]
:[(k, v)]
as) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) PQueue k v
r
Maybe ((k, v), PQueue k v)
_ -> ([(k, v)] -> [(k, v)]
forall a. [a] -> [a]
reverse [(k, v)]
as, PQueue k v
q)
go [(k, v)]
as Int
_ PQueue k v
q = ([(k, v)] -> [(k, v)]
forall a. [a] -> [a]
reverse [(k, v)]
as, PQueue k v
q)
take :: Ord k => Int -> PQueue k v -> ([v], PQueue k v)
take :: Int -> PQueue k v -> ([v], PQueue k v)
take = [v] -> Int -> PQueue k v -> ([v], PQueue k v)
forall k v. Ord k => [v] -> Int -> PQueue k v -> ([v], PQueue k v)
go []
where go :: Ord k => [v] -> Int -> PQueue k v -> ([v], PQueue k v)
go :: [v] -> Int -> PQueue k v -> ([v], PQueue k v)
go [v]
as Int
n PQueue k v
q | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = case PQueue k v -> Maybe (v, PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe (v, PQueue k v)
minView PQueue k v
q of
Just (v
a, PQueue k v
r) -> [v] -> Int -> PQueue k v -> ([v], PQueue k v)
forall k v. Ord k => [v] -> Int -> PQueue k v -> ([v], PQueue k v)
go (v
av -> [v] -> [v]
forall a. a -> [a] -> [a]
:[v]
as) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) PQueue k v
r
Maybe (v, PQueue k v)
_ -> ([v] -> [v]
forall a. [a] -> [a]
reverse [v]
as, PQueue k v
q)
go [v]
as Int
_ PQueue k v
q = ([v] -> [v]
forall a. [a] -> [a]
reverse [v]
as, PQueue k v
q)
minViewWithKey :: Ord k => PQueue k v -> Maybe ((k, v), PQueue k v)
minViewWithKey :: PQueue k v -> Maybe ((k, v), PQueue k v)
minViewWithKey (PQueue FingerTree (Prio k v) (Entry k v)
q)
| FingerTree (Prio k v) (Entry k v) -> Bool
forall v a. FingerTree v a -> Bool
FT.null FingerTree (Prio k v) (Entry k v)
q = Maybe ((k, v), PQueue k v)
forall a. Maybe a
Nothing
| Bool
otherwise = ((k, v), PQueue k v) -> Maybe ((k, v), PQueue k v)
forall a. a -> Maybe a
Just ((k
k, v
v), case FingerTree (Prio k v) (Entry k v)
-> ViewL (FingerTree (Prio k v)) (Entry k v)
forall v a.
Measured v a =>
FingerTree v a -> ViewL (FingerTree v) a
FT.viewl FingerTree (Prio k v) (Entry k v)
r of
Entry k v
_ :< FingerTree (Prio k v) (Entry k v)
r' -> FingerTree (Prio k v) (Entry k v) -> PQueue k v
forall k v. FingerTree (Prio k v) (Entry k v) -> PQueue k v
PQueue (FingerTree (Prio k v) (Entry k v)
l FingerTree (Prio k v) (Entry k v)
-> FingerTree (Prio k v) (Entry k v)
-> FingerTree (Prio k v) (Entry k v)
forall v a.
Measured v a =>
FingerTree v a -> FingerTree v a -> FingerTree v a
>< FingerTree (Prio k v) (Entry k v)
r')
ViewL (FingerTree (Prio k v)) (Entry k v)
_ -> [Char] -> PQueue k v
forall a. HasCallStack => [Char] -> a
error [Char]
"can't happen")
where
Prio k
k v
v = FingerTree (Prio k v) (Entry k v) -> Prio k v
forall v a. Measured v a => a -> v
measure FingerTree (Prio k v) (Entry k v)
q
(FingerTree (Prio k v) (Entry k v)
l, FingerTree (Prio k v) (Entry k v)
r) = (Prio k v -> Bool)
-> FingerTree (Prio k v) (Entry k v)
-> (FingerTree (Prio k v) (Entry k v),
FingerTree (Prio k v) (Entry k v))
forall v a.
Measured v a =>
(v -> Bool) -> FingerTree v a -> (FingerTree v a, FingerTree v a)
FT.split (k -> Prio k v -> Bool
forall k v. Ord k => k -> Prio k v -> Bool
below k
k) FingerTree (Prio k v) (Entry k v)
q
below :: Ord k => k -> Prio k v -> Bool
below :: k -> Prio k v -> Bool
below k
_ Prio k v
NoPrio = Bool
False
below k
k (Prio k
k' v
_) = k
k' k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
k