module Data.FingerTree.PSQueue
(Binding (..), PSQ, size, Data.FingerTree.PSQueue.null, Data.FingerTree.PSQueue.lookup,
empty, singleton, alter, Data.FingerTree.PSQueue.delete, adjust, adjustWithKey,
update, updateWithKey, toList, keys, fromList, fromAscList, minView,
findMin, deleteMin, range, atMost, Data.FingerTree.PSQueue.foldr,
Data.FingerTree.PSQueue.foldl) where
import qualified Data.Foldable as F
import qualified Data.FingerTree as FT
import Data.FingerTree (FingerTree, ViewL (..), (<|), (|>),
Measured (..), split, viewl, (><))
import Data.Monoid
import Data.Ord
import Data.List hiding (lookup)
import Prelude hiding (lookup)
data Prio k a = Prio (Binding k a) | PMax
deriving (Eq, Ord, Show)
instance (Ord a) => Monoid (Prio k a)
where mempty = PMax
mappend PMax y = y
mappend x PMax = x
mappend a@(Prio (u :-> x)) b@(Prio (v :-> y)) = if x <= y then a else b
data Key a = NoKey | Key a
deriving (Eq, Ord, Show)
instance Monoid (Key k)
where mempty = NoKey
mappend x NoKey = x
mappend x y = y
leqKey k NoKey = False
leqKey k (Key k') = k <= k'
data Binding k p = k :-> p
deriving (Eq, Ord, Show)
data KPS k p = KPS { kpsKey :: !(Key k),
kpsPrio :: !(Prio k p),
kpsSize :: !(Sum Int) }
deriving (Show)
instance Eq k => Eq (KPS k p) where
x == y = kpsKey x == kpsKey y
instance Ord k => Ord (KPS k p) where
compare = comparing kpsKey
instance (Ord p) => Monoid (KPS k p)
where mempty = KPS mempty mempty mempty
mappend (KPS k p s) (KPS k' p' s') = KPS (mappend k k')
(mappend p p')
(mappend s s')
instance (Ord k, Ord p) => Measured (KPS k p) (Binding k p)
where measure a@(k :-> p) = KPS (Key k) (Prio a) (Sum 1)
newtype PSQ k p = PSQ (FingerTree (KPS k p) (Binding k p))
deriving (Eq, Ord, Show, Measured (KPS k p))
size :: (Ord k, Ord p) => PSQ k p -> Int
size = getSum . kpsSize . measure
null :: (Ord k, Ord p) => PSQ k p -> Bool
null (PSQ q) = FT.null q
lookup :: (Ord k, Ord p) => k -> PSQ k p -> Maybe p
lookup k (PSQ q) = let (u,v) = split ((>= Key k) . kpsKey) q
in case viewl v of
EmptyL -> Nothing
(k' :-> p) :< v' | k == k' -> Just p
| otherwise -> Nothing
empty :: (Ord k, Ord p) => PSQ k p
empty = PSQ (FT.empty)
singleton :: (Ord k, Ord p) => k -> p -> PSQ k p
singleton k p = PSQ (FT.singleton (k :-> p))
alter :: (Ord k, Ord p) => (Maybe p -> Maybe p) -> k -> PSQ k p -> PSQ k p
alter f k (PSQ q) =
PSQ $ let (u,v) = split (leqKey k . kpsKey) q
in case viewl v of
EmptyL -> case f Nothing of
Nothing -> q
Just p -> q |> (k :-> p)
(k' :-> p') :< v'
| k == k' -> case f (Just p') of
Nothing -> u >< v'
Just p -> u >< ((k :-> p) <| v')
| otherwise -> case f Nothing of
Nothing -> u >< v
Just p -> u >< ((k :-> p) <| v)
delete :: (Ord k, Ord p) => k -> PSQ k p -> PSQ k p
delete = alter (const Nothing)
adjust :: (Ord k, Ord p) => (p -> p) -> k -> PSQ k p -> PSQ k p
adjust f = alter (fmap f)
adjustWithKey :: (Ord k, Ord p) => (k -> p -> p) -> k -> PSQ k p -> PSQ k p
adjustWithKey f k = adjust (f k) k
update :: (Ord k, Ord p) => (p -> Maybe p) -> k -> PSQ k p -> PSQ k p
update f = alter (>>= f)
updateWithKey :: (Ord k, Ord p) => (k -> p -> Maybe p) -> k -> PSQ k p -> PSQ k p
updateWithKey f k = update (f k) k
toList :: (Ord k, Ord p) => PSQ k p -> [Binding k p]
toList (PSQ q) = F.toList q
keys :: (Ord k, Ord p) => PSQ k p -> [k]
keys = map (\(k :-> p) -> k) . toList
fromList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromList = PSQ . FT.fromList . sort
fromAscList :: (Ord k, Ord p) => [Binding k p] -> PSQ k p
fromAscList = PSQ . FT.fromList
minView :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p, PSQ k p)
minView (PSQ q) =
let minPrio = kpsPrio . measure $ q
(u,v) = split ((== minPrio) . kpsPrio) q
in case viewl v of
EmptyL -> Nothing
((k :-> p) :< v') -> Just (k :-> p, PSQ (u >< v'))
findMin :: (Ord k, Ord p) => PSQ k p -> Maybe (Binding k p)
findMin q = case kpsPrio . measure $ q of
PMax -> Nothing
Prio b -> Just b
deleteMin :: (Ord k, Ord p) => PSQ k p -> PSQ k p
deleteMin q = maybe q id . fmap snd . minView $ q
range :: (Ord k, Ord p) => (k, k) -> PSQ k p -> PSQ k p
range (l,u) (PSQ q) = PSQ (fst . split ((> Key u) . kpsKey) . snd . split ((>= Key l) . kpsKey) $ q)
atMost :: (Ord k, Ord p) => p -> PSQ k p -> [Binding k p]
atMost p (PSQ q) =
let less (Prio (k :-> p')) = p' < p
(u,v) = split (less . kpsPrio) q
in case viewl v of
EmptyL -> []
(b :< v') -> b : atMost p (PSQ v')
foldr :: (Ord k, Ord p) => (Binding k p -> b -> b) -> b -> PSQ k p -> b
foldr f z = Prelude.foldr f z . toList
foldl :: (Ord k, Ord p) => (b -> Binding k p -> b) -> b -> PSQ k p -> b
foldl f z = Prelude.foldl f z . toList