{-# Language ImportQualifiedPost #-}
{-|
Module      : Advent.PQueue
Description : Int-priority min queue
Copyright   : (c) Eric Mertens, 2018
License     : ISC
Maintainer  : emertens@gmail.com

Priority queue with 'Int' priorities returning smallest priority first.

-}
{-# Language PatternSynonyms, ViewPatterns, DeriveTraversable #-}
{-# Options_GHC -Wno-name-shadowing #-}
module Advent.PQueue
  ( PQueue(Empty, (:<|))

  -- * Construction
  , singleton
  , fromList

  -- * Insertion
  , insert

  -- * Query
  , Advent.PQueue.null
  , view
  , viewWithPriority
  ) where

import Data.IntMap (IntMap)
import Data.IntMap.Strict qualified as IntMap

-- | Priority queue. No guarantees are made regarding the order
-- entries with the same priority are returned in.
newtype PQueue a = PQ (IntMap [a]) -- invariant: all values non-empty
  deriving ((forall a b. (a -> b) -> PQueue a -> PQueue b)
-> (forall a b. a -> PQueue b -> PQueue a) -> Functor PQueue
forall a b. a -> PQueue b -> PQueue a
forall a b. (a -> b) -> PQueue a -> PQueue b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> PQueue a -> PQueue b
fmap :: forall a b. (a -> b) -> PQueue a -> PQueue b
$c<$ :: forall a b. a -> PQueue b -> PQueue a
<$ :: forall a b. a -> PQueue b -> PQueue a
Functor, (forall m. Monoid m => PQueue m -> m)
-> (forall m a. Monoid m => (a -> m) -> PQueue a -> m)
-> (forall m a. Monoid m => (a -> m) -> PQueue a -> m)
-> (forall a b. (a -> b -> b) -> b -> PQueue a -> b)
-> (forall a b. (a -> b -> b) -> b -> PQueue a -> b)
-> (forall b a. (b -> a -> b) -> b -> PQueue a -> b)
-> (forall b a. (b -> a -> b) -> b -> PQueue a -> b)
-> (forall a. (a -> a -> a) -> PQueue a -> a)
-> (forall a. (a -> a -> a) -> PQueue a -> a)
-> (forall a. PQueue a -> [a])
-> (forall a. PQueue a -> Bool)
-> (forall a. PQueue a -> Int)
-> (forall a. Eq a => a -> PQueue a -> Bool)
-> (forall a. Ord a => PQueue a -> a)
-> (forall a. Ord a => PQueue a -> a)
-> (forall a. Num a => PQueue a -> a)
-> (forall a. Num a => PQueue a -> a)
-> Foldable PQueue
forall a. Eq a => a -> PQueue a -> Bool
forall a. Num a => PQueue a -> a
forall a. Ord a => PQueue a -> a
forall m. Monoid m => PQueue m -> m
forall a. PQueue a -> Bool
forall a. PQueue a -> Int
forall a. PQueue a -> [a]
forall a. (a -> a -> a) -> PQueue a -> a
forall m a. Monoid m => (a -> m) -> PQueue a -> m
forall b a. (b -> a -> b) -> b -> PQueue a -> b
forall a b. (a -> b -> b) -> b -> PQueue a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => PQueue m -> m
fold :: forall m. Monoid m => PQueue m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> PQueue a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> PQueue a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> PQueue a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> PQueue a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> PQueue a -> b
foldr :: forall a b. (a -> b -> b) -> b -> PQueue a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> PQueue a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> PQueue a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> PQueue a -> b
foldl :: forall b a. (b -> a -> b) -> b -> PQueue a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> PQueue a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> PQueue a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> PQueue a -> a
foldr1 :: forall a. (a -> a -> a) -> PQueue a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> PQueue a -> a
foldl1 :: forall a. (a -> a -> a) -> PQueue a -> a
$ctoList :: forall a. PQueue a -> [a]
toList :: forall a. PQueue a -> [a]
$cnull :: forall a. PQueue a -> Bool
null :: forall a. PQueue a -> Bool
$clength :: forall a. PQueue a -> Int
length :: forall a. PQueue a -> Int
$celem :: forall a. Eq a => a -> PQueue a -> Bool
elem :: forall a. Eq a => a -> PQueue a -> Bool
$cmaximum :: forall a. Ord a => PQueue a -> a
maximum :: forall a. Ord a => PQueue a -> a
$cminimum :: forall a. Ord a => PQueue a -> a
minimum :: forall a. Ord a => PQueue a -> a
$csum :: forall a. Num a => PQueue a -> a
sum :: forall a. Num a => PQueue a -> a
$cproduct :: forall a. Num a => PQueue a -> a
product :: forall a. Num a => PQueue a -> a
Foldable, Functor PQueue
Foldable PQueue
(Functor PQueue, Foldable PQueue) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> PQueue a -> f (PQueue b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    PQueue (f a) -> f (PQueue a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> PQueue a -> m (PQueue b))
-> (forall (m :: * -> *) a.
    Monad m =>
    PQueue (m a) -> m (PQueue a))
-> Traversable PQueue
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => PQueue (m a) -> m (PQueue a)
forall (f :: * -> *) a.
Applicative f =>
PQueue (f a) -> f (PQueue a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> PQueue a -> m (PQueue b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PQueue a -> f (PQueue b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PQueue a -> f (PQueue b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PQueue a -> f (PQueue b)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
PQueue (f a) -> f (PQueue a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
PQueue (f a) -> f (PQueue a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> PQueue a -> m (PQueue b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> PQueue a -> m (PQueue b)
$csequence :: forall (m :: * -> *) a. Monad m => PQueue (m a) -> m (PQueue a)
sequence :: forall (m :: * -> *) a. Monad m => PQueue (m a) -> m (PQueue a)
Traversable)

-- | Show a 'PQueue' using 'fromList'
--
-- >>> show (singleton 1 'a')
-- "fromList [(1,'a')]"
instance Show a => Show (PQueue a) where
  showsPrec :: Int -> PQueue a -> ShowS
showsPrec Int
prec (PQ IntMap [a]
q)
    = Bool -> ShowS -> ShowS
showParen (Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
11)
    (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String -> ShowS
showString String
"fromList "
    ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, a)] -> ShowS
forall a. Show a => a -> ShowS
shows [ (Int
p,a
v) | (Int
p, [a]
vs) <- IntMap [a] -> [(Int, [a])]
forall a. IntMap a -> [(Int, a)]
IntMap.toList IntMap [a]
q, a
v <- [a]
vs ]

instance Read a => Read (PQueue a) where
  readsPrec :: Int -> ReadS (PQueue a)
readsPrec Int
prec
    = Bool -> ReadS (PQueue a) -> ReadS (PQueue a)
forall a. Bool -> ReadS a -> ReadS a
readParen (Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
11) (ReadS (PQueue a) -> ReadS (PQueue a))
-> ReadS (PQueue a) -> ReadS (PQueue a)
forall a b. (a -> b) -> a -> b
$ \String
str ->
      do (String
"fromList", String
str) <- ReadS String
lex String
str
         ([(Int, a)]
xs,         String
str) <- ReadS [(Int, a)]
forall a. Read a => ReadS a
reads String
str
         (PQueue a, String) -> [(PQueue a, String)]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Int, a)] -> PQueue a
forall a. [(Int, a)] -> PQueue a
fromList [(Int, a)]
xs, String
str)

{-# Complete Empty, (:<|) #-}

-- | Empty priority queue
pattern Empty :: PQueue a
pattern $mEmpty :: forall {r} {a}. PQueue a -> ((# #) -> r) -> ((# #) -> r) -> r
$bEmpty :: forall a. PQueue a
Empty <- (Advent.PQueue.null -> True)
  where
    Empty = IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ IntMap [a]
forall a. IntMap a
IntMap.empty

-- | Pattern for extracting an element with the minimum priority
-- from the queue. See also: 'view'
pattern (:<|) :: a -> PQueue a -> PQueue a
pattern v $m:<| :: forall {r} {a}.
PQueue a -> (a -> PQueue a -> r) -> ((# #) -> r) -> r
:<| q <- (view -> Just (v,q))

-- | Test if a queue has no elements.
null :: PQueue a -> Bool
null :: forall a. PQueue a -> Bool
null (PQ IntMap [a]
q) = IntMap [a] -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap [a]
q

-- | Construct a priority queue from a single priority and value.
singleton :: Int {- ^ priority -} -> a {- ^ value -} -> PQueue a
singleton :: forall a. Int -> a -> PQueue a
singleton Int
p a
v = IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ (Int -> [a] -> IntMap [a]
forall a. Int -> a -> IntMap a
IntMap.singleton Int
p [a
v])

-- | Insert a new value into the queue given a priority.
insert :: Int {- ^ priority -} -> a {- ^ value -} -> PQueue a -> PQueue a
insert :: forall a. Int -> a -> PQueue a -> PQueue a
insert Int
k a
v (PQ IntMap [a]
q) = IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ ((Maybe [a] -> Maybe [a]) -> Int -> IntMap [a] -> IntMap [a]
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IntMap.alter Maybe [a] -> Maybe [a]
aux Int
k IntMap [a]
q)
  where
    aux :: Maybe [a] -> Maybe [a]
aux Maybe [a]
Nothing   = [a] -> Maybe [a]
forall a. a -> Maybe a
Just [a
v]
    aux (Just [a]
vs) = [a] -> Maybe [a]
forall a. a -> Maybe a
Just (a
va -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
vs)

-- | Match the lowest priority entry in a queue returning the corresponding
-- value and queue without that entry. See also: (':<|')
view :: PQueue a -> Maybe (a, PQueue a)
view :: forall a. PQueue a -> Maybe (a, PQueue a)
view (PQ IntMap [a]
q) =
  do ((Int
k,[a]
xs),IntMap [a]
q1) <- IntMap [a] -> Maybe ((Int, [a]), IntMap [a])
forall a. IntMap a -> Maybe ((Int, a), IntMap a)
IntMap.minViewWithKey IntMap [a]
q
     case [a]
xs of
       [] -> String -> Maybe (a, PQueue a)
forall a. HasCallStack => String -> a
error String
"Advent.PQueue.view: Malformed queue"
       [a
x] -> (a, PQueue a) -> Maybe (a, PQueue a)
forall a. a -> Maybe a
Just (a
x, IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ IntMap [a]
q1)
       a
x:[a]
xs -> let q2 :: PQueue a
q2 = IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ (Int -> [a] -> IntMap [a] -> IntMap [a]
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k [a]
xs IntMap [a]
q1)
               in PQueue a
q2 PQueue a -> Maybe (a, PQueue a) -> Maybe (a, PQueue a)
forall a b. a -> b -> b
`seq` (a, PQueue a) -> Maybe (a, PQueue a)
forall a. a -> Maybe a
Just (a
x,PQueue a
q2)

-- | Match the lowest priority entry in a queue returning the corresponding
-- priority, value and queue without that entry.
viewWithPriority :: PQueue a -> Maybe (Int, a, PQueue a)
viewWithPriority :: forall a. PQueue a -> Maybe (Int, a, PQueue a)
viewWithPriority (PQ IntMap [a]
q) =
  do ((Int
k,[a]
xs),IntMap [a]
q1) <- IntMap [a] -> Maybe ((Int, [a]), IntMap [a])
forall a. IntMap a -> Maybe ((Int, a), IntMap a)
IntMap.minViewWithKey IntMap [a]
q
     case [a]
xs of
       [] -> String -> Maybe (Int, a, PQueue a)
forall a. HasCallStack => String -> a
error String
"Advent.PQueue.view: Malformed queue"
       [a
x] -> (Int, a, PQueue a) -> Maybe (Int, a, PQueue a)
forall a. a -> Maybe a
Just (Int
k, a
x, IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ IntMap [a]
q1)
       a
x:[a]
xs -> let q2 :: PQueue a
q2 = IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ (Int -> [a] -> IntMap [a] -> IntMap [a]
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k [a]
xs IntMap [a]
q1)
               in PQueue a
q2 PQueue a -> Maybe (Int, a, PQueue a) -> Maybe (Int, a, PQueue a)
forall a b. a -> b -> b
`seq` (Int, a, PQueue a) -> Maybe (Int, a, PQueue a)
forall a. a -> Maybe a
Just (Int
k,a
x,PQueue a
q2)

-- | Construct a priority queue from a list of priorities and values.
fromList :: [(Int, a)] -> PQueue a
fromList :: forall a. [(Int, a)] -> PQueue a
fromList [(Int, a)]
xs = IntMap [a] -> PQueue a
forall a. IntMap [a] -> PQueue a
PQ (([a] -> [a] -> [a]) -> [(Int, [a])] -> IntMap [a]
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) [ (Int
p, [a
v]) | (Int
p, a
v) <- [(Int, a)]
xs ])