{-# LANGUAGE CPP                   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections         #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.PriorityQueue.FingerTree
-- Copyright   :  (c) Ross Paterson 2008
-- License     :  BSD-style
-- Maintainer  :  R.Paterson@city.ac.uk
-- Stability   :  experimental
-- Portability :  non-portable (MPTCs and functional dependencies)
--
-- Min-priority queues implemented using the 'FingerTree' type,
-- following section 4.6 of
--
--  * Ralf Hinze and Ross Paterson,
--    \"Finger trees: a simple general-purpose data structure\",
--    /Journal of Functional Programming/ 16:2 (2006) pp 197-217.
--    <http://staff.city.ac.uk/~ross/papers/FingerTree.html>
--
-- These have the same big-O complexity as skew heap implementations,
-- but are approximately an order of magnitude slower.
-- On the other hand, they are stable, so they can be used for fair
-- queueing.  They are also shallower, so that 'fmap' consumes less
-- space.
--
-- An amortized running time is given for each operation, with /n/
-- referring to the size of the priority queue.  These bounds hold even
-- in a persistent (shared) setting.
--
-- /Note/: Many of these operations have the same names as similar
-- operations on lists in the "Prelude".  The ambiguity may be resolved
-- using either qualification or the @hiding@ clause.
--
-----------------------------------------------------------------------------

module HaskellWorks.Data.PriorityQueue.FingerTree
  ( PQueue
    -- * Construction
  , empty
  , singleton
  , union
  , insert
  , add
  , fromList
    -- * Deconstruction
  , null
  , minView
  , minViewWithKey
  ) where

import Data.Bifunctor               (first)
import HaskellWorks.Data.FingerTree (FingerTree, Measured (..), ViewL (..), (<|), (><), (|>))
import Prelude                      hiding (null)

import qualified Data.Semigroup               as S
import qualified HaskellWorks.Data.FingerTree 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

-- | Priority queues.
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. Ord k => PQueue k v
empty
  {-# INLINE mempty #-}

-- | /O(1)/. The empty priority queue.
empty :: Ord k => 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. Measured v a => FingerTree v a
FT.empty

-- | /O(1)/. A singleton priority queue.
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 v a. Measured v a => a -> FingerTree v a
FT.singleton (k -> v -> Entry k v
forall k v. k -> v -> Entry k v
Entry k
k v
v))

-- | /O(log n)/. Add a (priority, value) pair to the front of a priority queue.
--
-- * @'insert' k v q = 'union' ('singleton' k v) q@
--
-- If @q@ contains entries with the same priority @k@, 'minView' of
-- @'insert' k v q@ will return them after this one.
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 Elem (FingerTree (Prio k v) (Entry k v))
-> FingerTree (Prio k v) (Entry k v)
-> FingerTree (Prio k v) (Entry k v)
forall v. Cons v => Elem v -> v -> v
<| FingerTree (Prio k v) (Entry k v)
q)

-- | /O(log n)/. Add a (priority, value) pair to the back of a priority queue.
--
-- * @'add' k v q = 'union' q ('singleton' k v)@
--
-- If @q@ contains entries with the same priority @k@, 'minView' of
-- @'add' k v q@ will return them before this one.
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)
-> Elem (FingerTree (Prio k v) (Entry k v))
-> FingerTree (Prio k v) (Entry k v)
forall v. Snoc v => v -> Elem v -> v
|> k -> v -> Entry k v
forall k v. k -> v -> Entry k v
Entry k
k v
v)

-- | /O(log(min(n1,n2)))/. Concatenate two priority queues.
-- 'union' is associative, with identity 'empty'.
--
-- If there are entries with the same priority in both arguments, 'minView'
-- of @'union' xs ys@ will return those from @xs@ before those from @ys@.
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. (Semigroup v, Container v) => v -> v -> v
>< FingerTree (Prio k v) (Entry k v)
ys)

-- | /O(n)/. Create a priority queue from a finite list of priorities
-- and values.
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. Ord k => PQueue k v
empty

-- | /O(1)/. Is this the empty priority queue?
null :: Ord k => 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. Measured v a => FingerTree v a -> Bool
FT.null FingerTree (Prio k v) (Entry k v)
q

-- | /O(1)/ for the element, /O(log(n))/ for the reduced queue.
-- Returns 'Nothing' for an empty map, or the value associated with the
-- minimal priority together with the rest of the priority queue.
--
--  * @'minView' 'empty' = 'Nothing'@
--
--  * @'minView' ('singleton' k v) = 'Just' (v, 'empty')@
--
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)

-- | /O(1)/ for the element, /O(log(n))/ for the reduced queue.
-- Returns 'Nothing' for an empty map, or the minimal (priority, value)
-- pair together with the rest of the priority queue.
--
--  * @'minViewWithKey' 'empty' = 'Nothing'@
--
--  * @'minViewWithKey' ('singleton' k v) = 'Just' ((k, v), 'empty')@
--
--  * If @'minViewWithKey' qi = 'Just' ((ki, vi), qi')@ and @k1 <= k2@,
--    then @'minViewWithKey' ('union' q1 q2) = 'Just' ((k1, v1), 'union' q1' q2)@
--
--  * If @'minViewWithKey' qi = 'Just' ((ki, vi), qi')@ and @k2 < k1@,
--    then @'minViewWithKey' ('union' q1 q2) = 'Just' ((k2, v2), 'union' q1 q2')@
--
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. Measured 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 = case 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 of
      Prio k
k v
v -> case (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 of
        (FingerTree (Prio k v) (Entry k v)
l, FingerTree (Prio k v) (Entry k v)
r) -> ((k, v), PQueue k v) -> Maybe ((k, v), PQueue k v)
forall a. a -> Maybe a
Just (((k, v), PQueue k v) -> Maybe ((k, v), PQueue k v))
-> (PQueue k v -> ((k, v), PQueue k v))
-> PQueue k v
-> Maybe ((k, v), PQueue k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k
k, v
v),) (PQueue k v -> Maybe ((k, v), PQueue k v))
-> PQueue k v -> Maybe ((k, v), PQueue k v)
forall a b. (a -> b) -> a -> b
$ 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. (Semigroup v, Container v) => v -> v -> v
>< 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"
      Prio k v
NoPrio -> [Char] -> Maybe ((k, v), PQueue k v)
forall a. HasCallStack => [Char] -> a
error [Char]
"can't happen: queue has no priority"

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