{-# LANGUAGE ScopedTypeVariables #-}
------------------------------------------------------------------------
-- |
-- Module      :  Data.DEPQ
-- Copyright   :  (c) Marco Zocca 2020
-- License     :  BSD3-style (see LICENSE)
--
-- Maintainer  :  @ocramz
-- Stability   :  experimental
-- Portability :  portable
--
-- Double-ended priority queue (DEPQ)
--
-- Allows for efficiently finding and removing both the minimum and maximum priority elements, due to the min-heap invariant property of the underlying representation.
--
-- See https://en.wikipedia.org/wiki/Double-ended_priority_queue for definitions; the current implementation is based on the "dual structure" method outlined in the wikipedia page.
--
-- Based on `P.IntPSQ` : https://hackage.haskell.org/package/psqueues-0.2.7.2/docs/Data-IntPSQ.html
--
-- = Usage
--
-- Populate a DEPQ (either from a `Foldable` collection such as a list or array or by `insert`ing incrementally) and query either of its extremes (with `findMin`, `findMax`, `popMin`, `popMax`, `topK`, `bottomK`).
--
-- = Note
--
-- Import this module qualified (e.g. @import qualified Data.DEPQ as DQ@ or similar), as some of the function names are pretty common (e.g. `lookup`, `empty`), and might collide with similar functions imported from other libraries.
------------------------------------------------------------------------
module Data.DEPQ (
   DEPQ, 
   -- * Creation
   empty,
   -- * Conversion from/to lists
   fromList, toList,
   -- * Predicates
   null,
   valid,
   -- * Properties
   size,
   -- * Modification
   insert, delete, deleteMin, deleteMax, popMin, popMax,
   -- * Lookup
   lookup, findMin, findMax,
   -- ** Top-K lookup
   topK, bottomK
  ) where

import Data.Maybe (fromMaybe)
import Data.Ord (Down(..))

-- containers
import qualified Data.Sequence as S (Seq, empty, (|>))
-- deepseq
import Control.DeepSeq     (NFData (rnf))
-- psqueues
import qualified Data.IntPSQ as P (IntPSQ, empty, null, size, insert, delete, toList, findMin, delete, deleteMin, valid, lookup)

import Prelude hiding (null, lookup)

import Test.QuickCheck (Arbitrary(..), Gen)

-- | A double-ended priority queue
data DEPQ p a = DEPQ {
    DEPQ p a -> IntPSQ p a
minHeap :: P.IntPSQ p a
  , DEPQ p a -> IntPSQ (Down p) a
maxHeap :: P.IntPSQ (Down p) a
                     } deriving (DEPQ p a -> DEPQ p a -> Bool
(DEPQ p a -> DEPQ p a -> Bool)
-> (DEPQ p a -> DEPQ p a -> Bool) -> Eq (DEPQ p a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall p a. (Ord p, Eq a) => DEPQ p a -> DEPQ p a -> Bool
/= :: DEPQ p a -> DEPQ p a -> Bool
$c/= :: forall p a. (Ord p, Eq a) => DEPQ p a -> DEPQ p a -> Bool
== :: DEPQ p a -> DEPQ p a -> Bool
$c== :: forall p a. (Ord p, Eq a) => DEPQ p a -> DEPQ p a -> Bool
Eq, Int -> DEPQ p a -> ShowS
[DEPQ p a] -> ShowS
DEPQ p a -> String
(Int -> DEPQ p a -> ShowS)
-> (DEPQ p a -> String) -> ([DEPQ p a] -> ShowS) -> Show (DEPQ p a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall p a. (Show p, Show a) => Int -> DEPQ p a -> ShowS
forall p a. (Show p, Show a) => [DEPQ p a] -> ShowS
forall p a. (Show p, Show a) => DEPQ p a -> String
showList :: [DEPQ p a] -> ShowS
$cshowList :: forall p a. (Show p, Show a) => [DEPQ p a] -> ShowS
show :: DEPQ p a -> String
$cshow :: forall p a. (Show p, Show a) => DEPQ p a -> String
showsPrec :: Int -> DEPQ p a -> ShowS
$cshowsPrec :: forall p a. (Show p, Show a) => Int -> DEPQ p a -> ShowS
Show)

instance Foldable (DEPQ p) where
  foldr :: (a -> b -> b) -> b -> DEPQ p a -> b
foldr a -> b -> b
f b
z (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
_) = (a -> b -> b) -> b -> IntPSQ p a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
z IntPSQ p a
mi

instance (NFData p, NFData a) => NFData (DEPQ p a) where
  rnf :: DEPQ p a -> ()
rnf (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = IntPSQ p a -> ()
forall a. NFData a => a -> ()
rnf IntPSQ p a
mi () -> () -> ()
`seq` IntPSQ (Down p) a -> ()
forall a. NFData a => a -> ()
rnf IntPSQ (Down p) a
ma

instance (Ord p, Arbitrary p, Arbitrary a) => Arbitrary (DEPQ p a) where
  arbitrary :: Gen (DEPQ p a)
arbitrary = [(Int, p, a)] -> DEPQ p a
forall (t :: * -> *) p a.
(Foldable t, Ord p) =>
t (Int, p, a) -> DEPQ p a
fromList ([(Int, p, a)] -> DEPQ p a) -> Gen [(Int, p, a)] -> Gen (DEPQ p a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Gen [(Int, p, a)]
forall a. Arbitrary a => Gen a
arbitrary :: Gen [(Int, p, a)])
  -- Convert given DEPQ into list, shrink it, then convert it back
  shrink :: DEPQ p a -> [DEPQ p a]
shrink DEPQ p a
depq = ([(Int, p, a)] -> DEPQ p a) -> [[(Int, p, a)]] -> [DEPQ p a]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, p, a)] -> DEPQ p a
forall (t :: * -> *) p a.
(Foldable t, Ord p) =>
t (Int, p, a) -> DEPQ p a
fromList ([[(Int, p, a)]] -> [DEPQ p a]) -> [[(Int, p, a)]] -> [DEPQ p a]
forall a b. (a -> b) -> a -> b
$ [(Int, p, a)] -> [[(Int, p, a)]]
forall a. Arbitrary a => a -> [a]
shrink ([(Int, p, a)] -> [[(Int, p, a)]])
-> [(Int, p, a)] -> [[(Int, p, a)]]
forall a b. (a -> b) -> a -> b
$ DEPQ p a -> [(Int, p, a)]
forall p v. DEPQ p v -> [(Int, p, v)]
toList DEPQ p a
depq


-- | Insert an element
insert :: (Ord p) =>
          Int -- ^ key
       -> p -- ^ priority
       -> a -- ^ value
       -> DEPQ p a -> DEPQ p a
insert :: Int -> p -> a -> DEPQ p a -> DEPQ p a
insert Int
k p
p a
v (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma'
  where
    mi' :: IntPSQ p a
mi' = Int -> p -> a -> IntPSQ p a -> IntPSQ p a
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
P.insert Int
k p
p  a
v IntPSQ p a
mi
    ma' :: IntPSQ (Down p) a
ma' = Int -> Down p -> a -> IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
P.insert Int
k (p -> Down p
forall a. a -> Down a
Down p
p) a
v IntPSQ (Down p) a
ma
{-# INLINE insert #-}

-- | The empty DEPQ
empty :: DEPQ p a
empty :: DEPQ p a
empty = IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
forall p v. IntPSQ p v
P.empty IntPSQ (Down p) a
forall p v. IntPSQ p v
P.empty

-- | Number of elements in the DEPQ
size :: DEPQ p a -> Int
size :: DEPQ p a -> Int
size (DEPQ IntPSQ p a
p IntPSQ (Down p) a
_) = IntPSQ p a -> Int
forall p v. IntPSQ p v -> Int
P.size IntPSQ p a
p

-- | Lookup a key
lookup :: Int -- ^ lookup key
       -> DEPQ p v
       -> Maybe (p, v)
lookup :: Int -> DEPQ p v -> Maybe (p, v)
lookup Int
k (DEPQ IntPSQ p v
p IntPSQ (Down p) v
_) = Int -> IntPSQ p v -> Maybe (p, v)
forall p v. Int -> IntPSQ p v -> Maybe (p, v)
P.lookup Int
k IntPSQ p v
p
{-# inline lookup #-}

-- | Populate a DEPQ from a 'Foldable' container (e.g. a list)
fromList :: (Foldable t, Ord p) =>
            t (Int, p, a) -- ^ (key, priority, value)
         -> DEPQ p a
fromList :: t (Int, p, a) -> DEPQ p a
fromList = (DEPQ p a -> (Int, p, a) -> DEPQ p a)
-> DEPQ p a -> t (Int, p, a) -> DEPQ p a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl DEPQ p a -> (Int, p, a) -> DEPQ p a
forall p a. Ord p => DEPQ p a -> (Int, p, a) -> DEPQ p a
insf DEPQ p a
forall p a. DEPQ p a
empty where
  insf :: DEPQ p a -> (Int, p, a) -> DEPQ p a
insf DEPQ p a
acc (Int
k,p
p,a
v) = Int -> p -> a -> DEPQ p a -> DEPQ p a
forall p a. Ord p => Int -> p -> a -> DEPQ p a -> DEPQ p a
insert Int
k p
p a
v DEPQ p a
acc
{-# inline fromList #-}

-- | Produce a list of (key, priority, value) triples with the entries of the DEPQ
--
-- Note : the order of the output list is undefined
toList :: DEPQ p v -> [(Int, p, v)]
toList :: DEPQ p v -> [(Int, p, v)]
toList (DEPQ IntPSQ p v
p IntPSQ (Down p) v
_) = IntPSQ p v -> [(Int, p, v)]
forall p v. IntPSQ p v -> [(Int, p, v)]
P.toList IntPSQ p v
p
{-# inline toList #-}

-- | Is the DEPQ empty ?
null :: DEPQ p v -> Bool
null :: DEPQ p v -> Bool
null (DEPQ IntPSQ p v
mi IntPSQ (Down p) v
ma) = IntPSQ p v -> Bool
forall p v. IntPSQ p v -> Bool
P.null IntPSQ p v
mi Bool -> Bool -> Bool
&& IntPSQ (Down p) v -> Bool
forall p v. IntPSQ p v -> Bool
P.null IntPSQ (Down p) v
ma

-- | Is the DEPQ valid ?
valid :: (Ord p) => DEPQ p v -> Bool
valid :: DEPQ p v -> Bool
valid (DEPQ IntPSQ p v
mi IntPSQ (Down p) v
ma) = IntPSQ p v -> Bool
forall p v. Ord p => IntPSQ p v -> Bool
P.valid IntPSQ p v
mi Bool -> Bool -> Bool
&& IntPSQ (Down p) v -> Bool
forall p v. Ord p => IntPSQ p v -> Bool
P.valid IntPSQ (Down p) v
ma

-- | Delete a (key, priority, value) triple from the queue. When
-- the key is not a member of the queue, the original queue is returned.
delete :: Ord p => Int -- ^ key of the triple to be deleted
       -> DEPQ p a -> DEPQ p a
delete :: Int -> DEPQ p a -> DEPQ p a
delete Int
k (DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma'
  where
    mi' :: IntPSQ p a
mi' = Int -> IntPSQ p a -> IntPSQ p a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
k IntPSQ p a
mi
    ma' :: IntPSQ (Down p) a
ma' = Int -> IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
k IntPSQ (Down p) a
ma

-- | Delete the minimum-priority element in the DEPQ
deleteMin :: Ord p => DEPQ p a -> DEPQ p a
deleteMin :: DEPQ p a -> DEPQ p a
deleteMin de :: DEPQ p a
de@(DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = case IntPSQ p a -> Maybe (Int, p, a)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ p a
mi of
  Maybe (Int, p, a)
Nothing -> DEPQ p a
de
  Just (Int
imin, p
_, a
_) -> IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma' where
    mi' :: IntPSQ p a
mi' = IntPSQ p a -> IntPSQ p a
forall p v. Ord p => IntPSQ p v -> IntPSQ p v
P.deleteMin IntPSQ p a
mi
    ma' :: IntPSQ (Down p) a
ma' = Int -> IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
imin IntPSQ (Down p) a
ma
{-# INLINE deleteMin #-}

-- | Delete the maximum-priority element in the DEPQ
deleteMax :: Ord p => DEPQ p a -> DEPQ p a
deleteMax :: DEPQ p a -> DEPQ p a
deleteMax de :: DEPQ p a
de@(DEPQ IntPSQ p a
mi IntPSQ (Down p) a
ma) = case IntPSQ (Down p) a -> Maybe (Int, Down p, a)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ (Down p) a
ma of
  Maybe (Int, Down p, a)
Nothing -> DEPQ p a
de
  Just (Int
imax, Down p
_, a
_) -> IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
forall p a. IntPSQ p a -> IntPSQ (Down p) a -> DEPQ p a
DEPQ IntPSQ p a
mi' IntPSQ (Down p) a
ma' where
    ma' :: IntPSQ (Down p) a
ma' = IntPSQ (Down p) a -> IntPSQ (Down p) a
forall p v. Ord p => IntPSQ p v -> IntPSQ p v
P.deleteMin IntPSQ (Down p) a
ma
    mi' :: IntPSQ p a
mi' = Int -> IntPSQ p a -> IntPSQ p a
forall p v. Ord p => Int -> IntPSQ p v -> IntPSQ p v
P.delete Int
imax IntPSQ p a
mi
{-# INLINE deleteMax #-}

-- | /O(1)/ Find the minimum-priority element in the DEPQ
findMin :: Ord p => DEPQ p v -> Maybe (Int, p, v)
findMin :: DEPQ p v -> Maybe (Int, p, v)
findMin (DEPQ IntPSQ p v
mi IntPSQ (Down p) v
_) = IntPSQ p v -> Maybe (Int, p, v)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ p v
mi
{-# inline findMin #-}

-- | /O(1)/ Find the maximum-priority element in the DEPQ
findMax :: Ord p => DEPQ p v -> Maybe (Int, p, v)
findMax :: DEPQ p v -> Maybe (Int, p, v)
findMax (DEPQ IntPSQ p v
_ IntPSQ (Down p) v
ma) = (Int, Down p, v) -> (Int, p, v)
forall a b c. (a, Down b, c) -> (a, b, c)
f ((Int, Down p, v) -> (Int, p, v))
-> Maybe (Int, Down p, v) -> Maybe (Int, p, v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntPSQ (Down p) v -> Maybe (Int, Down p, v)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v)
P.findMin IntPSQ (Down p) v
ma
  where
    f :: (a, Down b, c) -> (a, b, c)
f (a
i, Down b
p, c
v) = (a
i, b
p, c
v)
{-# inline findMax #-}


-- | Return the minimum along with a new DEPQ without that element
popMin :: Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMin :: DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMin DEPQ p v
q = do
  (Int, p, v)
x <- DEPQ p v -> Maybe (Int, p, v)
forall p v. Ord p => DEPQ p v -> Maybe (Int, p, v)
findMin DEPQ p v
q
  let q' :: DEPQ p v
q' = DEPQ p v -> DEPQ p v
forall p a. Ord p => DEPQ p a -> DEPQ p a
deleteMin DEPQ p v
q
  ((Int, p, v), DEPQ p v) -> Maybe ((Int, p, v), DEPQ p v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, p, v)
x, DEPQ p v
q')

-- | Return the maximum along with a new DEPQ without that element
popMax :: Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMax :: DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMax DEPQ p v
q = do
  (Int, p, v)
x <- DEPQ p v -> Maybe (Int, p, v)
forall p v. Ord p => DEPQ p v -> Maybe (Int, p, v)
findMax DEPQ p v
q
  let q' :: DEPQ p v
q' = DEPQ p v -> DEPQ p v
forall p a. Ord p => DEPQ p a -> DEPQ p a
deleteMax DEPQ p v
q
  ((Int, p, v), DEPQ p v) -> Maybe ((Int, p, v), DEPQ p v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, p, v)
x, DEPQ p v
q')

-- | K highest-scoring entries in the DEPQ
--
-- NB : this returns an empty sequence if there are fewer than K elements in the DEPQ
topK :: Ord p => Int -> DEPQ p v -> S.Seq (Int, p, v)
topK :: Int -> DEPQ p v -> Seq (Int, p, v)
topK = (DEPQ p v -> Maybe ((Int, p, v), DEPQ p v))
-> Int -> DEPQ p v -> Seq (Int, p, v)
forall q a. (q -> Maybe (a, q)) -> Int -> q -> Seq a
popK DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
forall p v. Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMax

-- | K lowest-scoring entries in the DEPQ
--
-- NB : this returns an empty sequence if there are fewer than K elements in the DEPQ
bottomK :: Ord p => Int -> DEPQ p v -> S.Seq (Int, p, v)
bottomK :: Int -> DEPQ p v -> Seq (Int, p, v)
bottomK = (DEPQ p v -> Maybe ((Int, p, v), DEPQ p v))
-> Int -> DEPQ p v -> Seq (Int, p, v)
forall q a. (q -> Maybe (a, q)) -> Int -> q -> Seq a
popK DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
forall p v. Ord p => DEPQ p v -> Maybe ((Int, p, v), DEPQ p v)
popMin

popK :: (q -> Maybe (a, q))
     -> Int
     -> q
     -> S.Seq a
popK :: (q -> Maybe (a, q)) -> Int -> q -> Seq a
popK q -> Maybe (a, q)
pop Int
kk q
qq = Seq a -> Maybe (Seq a) -> Seq a
forall a. a -> Maybe a -> a
fromMaybe Seq a
forall a. Seq a
S.empty (Maybe (Seq a) -> Seq a) -> Maybe (Seq a) -> Seq a
forall a b. (a -> b) -> a -> b
$ q -> Int -> Seq a -> Maybe (Seq a)
forall t. (Eq t, Num t) => q -> t -> Seq a -> Maybe (Seq a)
go q
qq Int
kk Seq a
forall a. Seq a
S.empty where
  go :: q -> t -> Seq a -> Maybe (Seq a)
go q
_ t
0 Seq a
acc = Seq a -> Maybe (Seq a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Seq a
acc
  go q
q t
k Seq a
acc = do
    (a
x, q
q') <- q -> Maybe (a, q)
pop q
q
    q -> t -> Seq a -> Maybe (Seq a)
go q
q' (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Seq a
acc Seq a -> a -> Seq a
forall a. Seq a -> a -> Seq a
S.|> a
x)