module Data.PriorityQueue (PQ, singleton, insert, minView, minViewWithKey, minAlterWithKeyF,
                           from, toOrdList, mapMaybeWithKeyA, foldMapWithKey, toV,
                           foldrWithKey, foldlWithKey, foldrWithKeyM, foldlWithKeyM) where

import Control.Applicative
import Control.Monad
import Control.Monad.Primitive
import Data.Bits
import Data.Bool
import Data.Filtrable
import Data.Foldable
import Data.Function (on)
import Data.Functor.Classes
import Data.Tuple (swap)
import Data.Vector.Generic ((!?), Vector)
import qualified Data.Vector.Generic as V
import Data.Vector.Generic.Mutable (MVector)
import qualified Data.Vector.Generic.Mutable as MV
import Util.Vector as V
import Util.Vector.Mutable as MV
import Util

newtype PQ v k a = PQ { PQ v k a -> v (k, a)
toV :: v (k, a) }

foldMapWithKey :: (Ord k, Vector v (k, a)) => Monoid b => (k -> a -> b) -> PQ v k a -> b
foldMapWithKey :: (k -> a -> b) -> PQ v k a -> b
foldMapWithKey f :: k -> a -> b
f = (k -> a -> b -> b) -> b -> PQ v k a -> b
forall k (v :: * -> *) a b.
(Ord k, Vector v (k, a)) =>
(k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey (b -> b -> b
forall a. Semigroup a => a -> a -> a
(<>) (b -> b -> b) -> (k -> a -> b) -> k -> a -> b -> b
forall c d a b. (c -> d) -> (a -> b -> c) -> a -> b -> d
∘∘ k -> a -> b
f) b
forall a. Monoid a => a
mempty

foldrWithKey :: (Ord k, Vector v (k, a)) => (k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey :: (k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey f :: k -> a -> b -> b
f z :: b
z = PQ v k a -> Maybe (k, a, PQ v k a)
forall k (v :: * -> *) a.
(Ord k, Vector v (k, a)) =>
PQ v k a -> Maybe (k, a, PQ v k a)
minViewWithKey (PQ v k a -> Maybe (k, a, PQ v k a))
-> (Maybe (k, a, PQ v k a) -> b) -> PQ v k a -> b
forall (p :: * -> * -> *) a b c.
Category p =>
p a b -> p b c -> p a c
& \ case Nothing -> b
z
                                           Just (k :: k
k, a :: a
a, pq :: PQ v k a
pq) -> k -> a -> b -> b
f k
k a
a (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$ (k -> a -> b -> b) -> b -> PQ v k a -> b
forall k (v :: * -> *) a b.
(Ord k, Vector v (k, a)) =>
(k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey k -> a -> b -> b
f b
z PQ v k a
pq

foldlWithKey :: (Ord k, Vector v (k, a)) => (k -> b -> a -> b) -> b -> PQ v k a -> b
foldlWithKey :: (k -> b -> a -> b) -> b -> PQ v k a -> b
foldlWithKey f :: k -> b -> a -> b
f z :: b
z xs :: PQ v k a
xs = (k -> a -> (b -> b) -> b -> b) -> (b -> b) -> PQ v k a -> b -> b
forall k (v :: * -> *) a b.
(Ord k, Vector v (k, a)) =>
(k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey (\ k :: k
k a :: a
a c :: b -> b
c x :: b
x -> b -> b
c (k -> b -> a -> b
f k
k b
x a
a)) b -> b
forall a. a -> a
id PQ v k a
xs b
z

foldrWithKeyM :: (Ord k, Vector v (k, a)) => Monad m => (k -> a -> b -> m b) -> b -> PQ v k a -> m b
foldrWithKeyM :: (k -> a -> b -> m b) -> b -> PQ v k a -> m b
foldrWithKeyM f :: k -> a -> b -> m b
f z :: b
z xs :: PQ v k a
xs = (k -> (b -> m b) -> a -> b -> m b)
-> (b -> m b) -> PQ v k a -> b -> m b
forall k (v :: * -> *) a b.
(Ord k, Vector v (k, a)) =>
(k -> b -> a -> b) -> b -> PQ v k a -> b
foldlWithKey k -> (b -> m b) -> a -> b -> m b
forall b. k -> (b -> m b) -> a -> b -> m b
f' b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure PQ v k a
xs b
z where f' :: k -> (b -> m b) -> a -> b -> m b
f' k :: k
k c :: b -> m b
c x :: a
x z :: b
z = k -> a -> b -> m b
f k
k a
x b
z m b -> (b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
c

foldlWithKeyM :: (Ord k, Vector v (k, a)) => Monad m => (k -> b -> a -> m b) -> b -> PQ v k a -> m b
foldlWithKeyM :: (k -> b -> a -> m b) -> b -> PQ v k a -> m b
foldlWithKeyM f :: k -> b -> a -> m b
f z :: b
z xs :: PQ v k a
xs = (k -> a -> (b -> m b) -> b -> m b)
-> (b -> m b) -> PQ v k a -> b -> m b
forall k (v :: * -> *) a b.
(Ord k, Vector v (k, a)) =>
(k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey k -> a -> (b -> m b) -> b -> m b
forall b. k -> a -> (b -> m b) -> b -> m b
f' b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure PQ v k a
xs b
z where f' :: k -> a -> (b -> m b) -> b -> m b
f' k :: k
k x :: a
x c :: b -> m b
c z :: b
z = k -> b -> a -> m b
f k
k b
z a
x m b -> (b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
c

mapMaybeWithKeyA :: (Ord k, Vector v (k, a), Vector v (k, b), Applicative p)
                 => (k -> a -> p (Maybe (k, b))) -> PQ v k a -> p (PQ v k b)
mapMaybeWithKeyA :: (k -> a -> p (Maybe (k, b))) -> PQ v k a -> p (PQ v k b)
mapMaybeWithKeyA f :: k -> a -> p (Maybe (k, b))
f = ([(k, b)] -> PQ v k b) -> p [(k, b)] -> p (PQ v k b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(k, b)] -> PQ v k b
forall k (v :: * -> *) a (f :: * -> *).
(Ord k, Vector v (k, a), Foldable f) =>
f (k, a) -> PQ v k a
from (p [(k, b)] -> p (PQ v k b))
-> (PQ v k a -> p [(k, b)]) -> PQ v k a -> p (PQ v k b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, a) -> p (Maybe (k, b))) -> [(k, a)] -> p [(k, b)]
forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA ((k -> a -> p (Maybe (k, b))) -> (k, a) -> p (Maybe (k, b))
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> a -> p (Maybe (k, b))
f) ([(k, a)] -> p [(k, b)])
-> (PQ v k a -> [(k, a)]) -> PQ v k a -> p [(k, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v (k, a) -> [(k, a)]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList (v (k, a) -> [(k, a)])
-> (PQ v k a -> v (k, a)) -> PQ v k a -> [(k, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PQ v k a -> v (k, a)
forall (v :: * -> *) k a. PQ v k a -> v (k, a)
toV

toOrdList :: (Ord k, Vector v (k, a)) => PQ v k a -> [(k, a)]
toOrdList :: PQ v k a -> [(k, a)]
toOrdList = (k -> a -> [(k, a)] -> [(k, a)])
-> [(k, a)] -> PQ v k a -> [(k, a)]
forall k (v :: * -> *) a b.
(Ord k, Vector v (k, a)) =>
(k -> a -> b -> b) -> b -> PQ v k a -> b
foldrWithKey (((k, a) -> [(k, a)] -> [(k, a)]) -> k -> a -> [(k, a)] -> [(k, a)]
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (:)) []

from :: (Ord k, Vector v (k, a), Foldable f) => f (k, a) -> PQ v k a
from :: f (k, a) -> PQ v k a
from = v (k, a) -> PQ v k a
forall k (v :: * -> *) a.
(Ord k, Vector v (k, a)) =>
v (k, a) -> PQ v k a
fromV (v (k, a) -> PQ v k a)
-> (f (k, a) -> v (k, a)) -> f (k, a) -> PQ v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, a)] -> v (k, a)
forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList ([(k, a)] -> v (k, a))
-> (f (k, a) -> [(k, a)]) -> f (k, a) -> v (k, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (k, a) -> [(k, a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

fromV :: (Ord k, Vector v (k, a)) => v (k, a) -> PQ v k a
fromV :: v (k, a) -> PQ v k a
fromV = v (k, a) -> PQ v k a
forall (v :: * -> *) k a. v (k, a) -> PQ v k a
PQ (v (k, a) -> PQ v k a)
-> (v (k, a) -> v (k, a)) -> v (k, a) -> PQ v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall s. Mutable v s (k, a) -> ST s ()) -> v (k, a) -> v (k, a)
forall (v :: * -> *) a.
Vector v a =>
(forall s. Mutable v s a -> ST s ()) -> v a -> v a
V.modify (((k, a) -> (k, a) -> Ordering)
-> Mutable v (PrimState (ST s)) (k, a) -> ST s ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
(a -> a -> Ordering) -> v (PrimState m) a -> m ()
buildBy (k -> k -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (k -> k -> Ordering)
-> ((k, a) -> k) -> (k, a) -> (k, a) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (k, a) -> k
forall a b. (a, b) -> a
fst))

singleton :: Vector v (k, a) => k -> a -> PQ v k a
singleton :: k -> a -> PQ v k a
singleton = ((k, a) -> PQ v k a) -> k -> a -> PQ v k a
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (((k, a) -> PQ v k a) -> k -> a -> PQ v k a)
-> ((k, a) -> PQ v k a) -> k -> a -> PQ v k a
forall a b. (a -> b) -> a -> b
$ v (k, a) -> PQ v k a
forall (v :: * -> *) k a. v (k, a) -> PQ v k a
PQ (v (k, a) -> PQ v k a)
-> ((k, a) -> v (k, a)) -> (k, a) -> PQ v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k, a) -> v (k, a)
forall (v :: * -> *) a. Vector v a => a -> v a
V.singleton

insert :: (Ord k, Vector v (k, a)) => k -> a -> PQ v k a -> PQ v k a
insert :: k -> a -> PQ v k a -> PQ v k a
insert k :: k
k a :: a
a (PQ xs :: v (k, a)
xs) = v (k, a) -> PQ v k a
forall (v :: * -> *) k a. v (k, a) -> PQ v k a
PQ (v (k, a) -> PQ v k a) -> v (k, a) -> PQ v k a
forall a b. (a -> b) -> a -> b
$ (forall s. Mutable v s (k, a) -> ST s ()) -> v (k, a) -> v (k, a)
forall (v :: * -> *) a.
Vector v a =>
(forall s. Mutable v s a -> ST s ()) -> v a -> v a
V.modify forall s. Mutable v s (k, a) -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v (k, a)) =>
v (PrimState m) (k, a) -> m ()
go v (k, a)
xs
  where go :: v (PrimState m) (k, a) -> m ()
go xs :: v (PrimState m) (k, a)
xs = do
            let l :: Int
l = v (PrimState m) (k, a) -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v (PrimState m) (k, a)
xs
            v (PrimState m) (k, a)
xs <- v (PrimState m) (k, a) -> Int -> m (v (PrimState m) (k, a))
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
MV.unsafeGrow v (PrimState m) (k, a)
xs 1
            v (PrimState m) (k, a) -> Int -> (k, a) -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite v (PrimState m) (k, a)
xs Int
l (k
k, a
a)
            ((k, a) -> (k, a) -> Ordering)
-> v (PrimState m) (k, a) -> Int -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
(a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftUpBy (k -> k -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (k -> k -> Ordering)
-> ((k, a) -> k) -> (k, a) -> (k, a) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (k, a) -> k
forall a b. (a, b) -> a
fst) v (PrimState m) (k, a)
xs Int
l

minAlterWithKeyF :: (Ord k, Vector v (k, a), Functor f)
                 => (Maybe (k, a) -> f (Maybe (k, a))) -> PQ v k a -> f (PQ v k a)
minAlterWithKeyF :: (Maybe (k, a) -> f (Maybe (k, a))) -> PQ v k a -> f (PQ v k a)
minAlterWithKeyF f :: Maybe (k, a) -> f (Maybe (k, a))
f (PQ xs :: v (k, a)
xs) = v (k, a) -> PQ v k a
forall (v :: * -> *) k a. v (k, a) -> PQ v k a
PQ (v (k, a) -> PQ v k a)
-> (Maybe (k, a) -> v (k, a)) -> Maybe (k, a) -> PQ v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe (k, a) -> v (k, a))
-> (Maybe (k, a) -> v (k, a)) -> Bool -> Maybe (k, a) -> v (k, a)
forall a. a -> a -> Bool -> a
bool (((v (k, a) -> v (k, a)) -> v (k, a) -> v (k, a))
-> (v (k, a) -> v (k, a), v (k, a)) -> v (k, a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (v (k, a) -> v (k, a)) -> v (k, a) -> v (k, a)
forall a. a -> a
id ((v (k, a) -> v (k, a), v (k, a)) -> v (k, a))
-> (Maybe (k, a) -> (v (k, a) -> v (k, a), v (k, a)))
-> Maybe (k, a)
-> v (k, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \ a :: Maybe (k, a)
a -> (forall s. Mutable v s (k, a) -> ST s (v (k, a) -> v (k, a)))
-> v (k, a) -> (v (k, a) -> v (k, a), v (k, a))
forall (v :: * -> *) a b.
Vector v a =>
(forall s. Mutable v s a -> ST s b) -> v a -> (b, v a)
V.modify' (Maybe (k, a)
-> Mutable v (PrimState (ST s)) (k, a)
-> ST s (v (k, a) -> v (k, a))
forall (m :: * -> *) (v :: * -> *) a a (v :: * -> * -> *) b.
(PrimMonad m, Vector v a, Ord a, MVector v (a, b)) =>
Maybe (a, b) -> v (PrimState m) (a, b) -> m (v a -> v a)
go Maybe (k, a)
a) v (k, a)
xs) (v (k, a) -> ((k, a) -> v (k, a)) -> Maybe (k, a) -> v (k, a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe v (k, a)
xs (k, a) -> v (k, a)
forall (v :: * -> *) a. Vector v a => a -> v a
V.singleton) (v (k, a) -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v (k, a)
xs) (Maybe (k, a) -> PQ v k a) -> f (Maybe (k, a)) -> f (PQ v k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (k, a) -> f (Maybe (k, a))
f (v (k, a)
xs v (k, a) -> Int -> Maybe (k, a)
forall (v :: * -> *) a. Vector v a => v a -> Int -> Maybe a
!? 0)
  where go :: Maybe (a, b) -> v (PrimState m) (a, b) -> m (v a -> v a)
go Nothing xs :: v (PrimState m) (a, b)
xs = do
            v (PrimState m) (a, b) -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
MV.unsafeSwap v (PrimState m) (a, b)
xs 0 (v (PrimState m) (a, b) -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v (PrimState m) (a, b)
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
            v a -> v a
forall (v :: * -> *) a. Vector v a => v a -> v a
V.init (v a -> v a) -> m () -> m (v a -> v a)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ((a, b) -> (a, b) -> Ordering)
-> v (PrimState m) (a, b) -> Int -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
(a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftDownBy (a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (a -> a -> Ordering)
-> ((a, b) -> a) -> (a, b) -> (a, b) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, b) -> a
forall a b. (a, b) -> a
fst) (v (PrimState m) (a, b) -> v (PrimState m) (a, b)
forall (v :: * -> * -> *) a s. MVector v a => v s a -> v s a
MV.init v (PrimState m) (a, b)
xs) 0
        go (Just (k :: a
k, a :: b
a)) xs :: v (PrimState m) (a, b)
xs = do
            v (PrimState m) (a, b) -> Int -> (a, b) -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite v (PrimState m) (a, b)
xs 0 (a
k, b
a)
            v a -> v a
forall a. a -> a
id     (v a -> v a) -> m () -> m (v a -> v a)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ((a, b) -> (a, b) -> Ordering)
-> v (PrimState m) (a, b) -> Int -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
(a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftDownBy (a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (a -> a -> Ordering)
-> ((a, b) -> a) -> (a, b) -> (a, b) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, b) -> a
forall a b. (a, b) -> a
fst) v (PrimState m) (a, b)
xs           0

minView :: (Ord k, Vector v (k, a)) => PQ v k a -> Maybe (a, PQ v k a)
minView :: PQ v k a -> Maybe (a, PQ v k a)
minView = PQ v k a -> Maybe (k, a, PQ v k a)
forall k (v :: * -> *) a.
(Ord k, Vector v (k, a)) =>
PQ v k a -> Maybe (k, a, PQ v k a)
minViewWithKey (PQ v k a -> Maybe (k, a, PQ v k a))
-> (Maybe (k, a, PQ v k a) -> Maybe (a, PQ v k a))
-> PQ v k a
-> Maybe (a, PQ v k a)
forall (p :: * -> * -> *) a b c.
Category p =>
p a b -> p b c -> p a c
& ((k, a, PQ v k a) -> (a, PQ v k a))
-> Maybe (k, a, PQ v k a) -> Maybe (a, PQ v k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (_, a :: a
a, pq :: PQ v k a
pq) -> (a
a, PQ v k a
pq))

minViewWithKey :: (Ord k, Vector v (k, a)) => PQ v k a -> Maybe (k, a, PQ v k a)
minViewWithKey :: PQ v k a -> Maybe (k, a, PQ v k a)
minViewWithKey = (Maybe (k, a) -> (Maybe (k, a), Maybe (k, a)))
-> PQ v k a -> (Maybe (k, a), PQ v k a)
forall k (v :: * -> *) a (f :: * -> *).
(Ord k, Vector v (k, a), Functor f) =>
(Maybe (k, a) -> f (Maybe (k, a))) -> PQ v k a -> f (PQ v k a)
minAlterWithKeyF ((Maybe (k, a) -> Maybe (k, a) -> (Maybe (k, a), Maybe (k, a)))
-> Maybe (k, a) -> Maybe (k, a) -> (Maybe (k, a), Maybe (k, a))
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Maybe (k, a)
forall a. Maybe a
Nothing) (PQ v k a -> (Maybe (k, a), PQ v k a))
-> ((Maybe (k, a), PQ v k a) -> Maybe (k, a, PQ v k a))
-> PQ v k a
-> Maybe (k, a, PQ v k a)
forall (p :: * -> * -> *) a b c.
Category p =>
p a b -> p b c -> p a c
& (Maybe (k, a), PQ v k a) -> (PQ v k a, Maybe (k, a))
forall a b. (a, b) -> (b, a)
swap ((Maybe (k, a), PQ v k a) -> (PQ v k a, Maybe (k, a)))
-> ((PQ v k a, Maybe (k, a)) -> Maybe (k, a, PQ v k a))
-> (Maybe (k, a), PQ v k a)
-> Maybe (k, a, PQ v k a)
forall (p :: * -> * -> *) a b c.
Category p =>
p a b -> p b c -> p a c
& (PQ v k a, Maybe (k, a)) -> Maybe (PQ v k a, (k, a))
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA ((PQ v k a, Maybe (k, a)) -> Maybe (PQ v k a, (k, a)))
-> (Maybe (PQ v k a, (k, a)) -> Maybe (k, a, PQ v k a))
-> (PQ v k a, Maybe (k, a))
-> Maybe (k, a, PQ v k a)
forall (p :: * -> * -> *) a b c.
Category p =>
p a b -> p b c -> p a c
& ((PQ v k a, (k, a)) -> (k, a, PQ v k a))
-> Maybe (PQ v k a, (k, a)) -> Maybe (k, a, PQ v k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (pq :: PQ v k a
pq, (k :: k
k, a :: a
a)) -> (k
k, a
a, PQ v k a
pq))

siftUpBy :: (MVector v a, PrimMonad m) => (a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftUpBy :: (a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftUpBy cmp :: a -> a -> Ordering
cmp xs :: v (PrimState m) a
xs = Int -> m ()
go
  where go :: Int -> m ()
go 0 = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        go n :: Int
n = do
            let m :: Int
m = (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` 1
            a
i <- v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MV.unsafeRead v (PrimState m) a
xs Int
m
            a
j <- v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MV.unsafeRead v (PrimState m) a
xs Int
n
            Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ordering
GT Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== a -> a -> Ordering
cmp a
i a
j) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
MV.unsafeSwap v (PrimState m) a
xs Int
m Int
n m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go Int
m

siftDownBy :: (MVector v a, PrimMonad m) => (a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftDownBy :: (a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftDownBy cmp :: a -> a -> Ordering
cmp xs :: v (PrimState m) a
xs = Int -> m ()
go
  where go :: Int -> m ()
go m :: Int
m = ((Maybe (Int, a) -> Maybe (Int, a) -> Maybe (Int, a))
-> m (Maybe (Int, a)) -> m (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ((Maybe (Int, a) -> Maybe (Int, a) -> Ordering)
-> Maybe (Int, a) -> Maybe (Int, a) -> Maybe (Int, a)
forall a. (a -> a -> Ordering) -> a -> a -> a
minBy Maybe (Int, a) -> Maybe (Int, a) -> Ordering
forall a. Maybe (a, a) -> Maybe (a, a) -> Ordering
cmp') (m (Maybe (Int, a)) -> m (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> (Int -> m (Maybe (Int, a))) -> Int -> Int -> m (Maybe (Int, a))
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` \ n :: Int
n ->
                (a -> (Int, a)) -> Maybe a -> Maybe (Int, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) Int
n) (Maybe a -> Maybe (Int, a)) -> m (Maybe a) -> m (Maybe (Int, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v (PrimState m) a -> Int -> m (Maybe a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (Maybe a)
MV.readMaybe v (PrimState m) a
xs Int
n) (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 2) m (Maybe (Int, a)) -> (Maybe (Int, a) -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
            Nothing -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just (n, j) -> do
                a
i <- v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MV.unsafeRead v (PrimState m) a
xs Int
m
                Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ordering
GT Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== a -> a -> Ordering
cmp a
i a
j) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
MV.unsafeSwap v (PrimState m) a
xs Int
m Int
n m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go Int
n
        cmp' :: Maybe (a, a) -> Maybe (a, a) -> Ordering
cmp' = (Maybe (a, a) -> Maybe (a, a) -> Ordering)
-> Maybe (a, a) -> Maybe (a, a) -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Maybe (a, a) -> Maybe (a, a) -> Ordering)
 -> Maybe (a, a) -> Maybe (a, a) -> Ordering)
-> (Maybe (a, a) -> Maybe (a, a) -> Ordering)
-> Maybe (a, a)
-> Maybe (a, a)
-> Ordering
forall a b. (a -> b) -> a -> b
$ ((a, a) -> (a, a) -> Ordering)
-> Maybe (a, a) -> Maybe (a, a) -> Ordering
forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare ((a -> a -> Ordering) -> a -> a -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> Ordering
cmp (a -> a -> Ordering)
-> ((a, a) -> a) -> (a, a) -> (a, a) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, a) -> a
forall a b. (a, b) -> b
snd)

minBy :: (a -> a -> Ordering) -> a -> a -> a
minBy :: (a -> a -> Ordering) -> a -> a -> a
minBy cmp :: a -> a -> Ordering
cmp x :: a
x y :: a
y | Ordering
GT <- a -> a -> Ordering
cmp a
x a
y = a
y | Bool
otherwise = a
x

buildBy :: (MVector v a, PrimMonad m) => (a -> a -> Ordering) -> v (PrimState m) a -> m ()
buildBy :: (a -> a -> Ordering) -> v (PrimState m) a -> m ()
buildBy cmp :: a -> a -> Ordering
cmp xs :: v (PrimState m) a
xs = [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ ([Int] -> [Int]
forall a. [a] -> [a]
reverse [0 .. v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length v (PrimState m) a
xs Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` 1]) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ (a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
(a -> a -> Ordering) -> v (PrimState m) a -> Int -> m ()
siftDownBy a -> a -> Ordering
cmp v (PrimState m) a
xs