{-# LANGUAGE Trustworthy       #-}
{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE CPP               #-}
{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE MagicHash         #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE UnboxedTuples     #-}

module GHC.Event.PSQ
    (
    -- * Binding Type
      Elem(..)
    , Key
    , Prio

    -- * Priority Search Queue Type
    , PSQ

    -- * Query
    , size
    , null
    , lookup

    -- * Construction
    , empty
    , singleton

    -- * Insertion
    , unsafeInsertNew

    -- * Delete/Update
    , delete
    , adjust

    -- * Conversion
    , toList

    -- * Min
    , findMin
    , deleteMin
    , minView
    , atMost
    ) where

import GHC.Base hiding (Nat, empty)
import GHC.Event.Unique
import GHC.Word (Word64)
import GHC.Num (Num(..))
import GHC.Real (fromIntegral)

#include "MachDeps.h"

-- TODO (SM): get rid of bang patterns

{-
-- Use macros to define strictness of functions.
-- STRICT_x_OF_y denotes a y-ary function strict in the x-th parameter.
-- We do not use BangPatterns, because they are not in any standard and we
-- want the compilers to be compiled by as many compilers as possible.
#define STRICT_1_OF_2(fn) fn arg _ | arg `seq` False = undefined
-}


------------------------------------------------------------------------------
-- Types
------------------------------------------------------------------------------

type Prio = Word64

type Nat = Word

type Key = Unique

-- | We store masks as the index of the bit that determines the branching.
type Mask = Int

type PSQ a = IntPSQ a

-- | @E k p@ binds the key @k@ with the priority @p@.
data Elem a = E
    { Elem a -> Key
key   :: {-# UNPACK #-} !Key
    , Elem a -> Prio
prio  :: {-# UNPACK #-} !Prio
    , Elem a -> a
value :: a
    }

-- | A priority search queue with @Int@ keys and priorities of type @p@ and
-- values of type @v@. It is strict in keys, priorities and values.
data IntPSQ v
    = Bin {-# UNPACK #-} !Key {-# UNPACK #-} !Prio !v {-# UNPACK #-} !Mask !(IntPSQ v) !(IntPSQ v)
    | Tip {-# UNPACK #-} !Key {-# UNPACK #-} !Prio !v
    | Nil

-- bit twiddling
----------------

(.&.) :: Nat -> Nat -> Nat
.&. :: Nat -> Nat -> Nat
(.&.) (W# Word#
w1) (W# Word#
w2) = Word# -> Nat
W# (Word#
w1 Word# -> Word# -> Word#
`and#` Word#
w2)
{-# INLINE (.&.) #-}

xor :: Nat -> Nat -> Nat
xor :: Nat -> Nat -> Nat
xor (W# Word#
w1) (W# Word#
w2) = Word# -> Nat
W# (Word#
w1 Word# -> Word# -> Word#
`xor#` Word#
w2)
{-# INLINE xor #-}

complement :: Nat -> Nat
complement :: Nat -> Nat
complement (W# Word#
w) = Word# -> Nat
W# (Word#
w Word# -> Word# -> Word#
`xor#` Word#
mb)
  where
#if WORD_SIZE_IN_BITS == 32
    mb = 0xFFFFFFFF##
#elif WORD_SIZE_IN_BITS == 64
    mb :: Word#
mb = Word#
0xFFFFFFFFFFFFFFFF##
#else
#error Unhandled value for WORD_SIZE_IN_BITS
#endif
{-# INLINE complement #-}

{-# INLINE natFromInt #-}
natFromInt :: Int -> Nat
natFromInt :: Int -> Nat
natFromInt = Int -> Nat
forall a b. (Integral a, Num b) => a -> b
fromIntegral

{-# INLINE intFromNat #-}
intFromNat :: Nat -> Int
intFromNat :: Nat -> Int
intFromNat = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral

{-# INLINE zero #-}
zero :: Key -> Mask -> Bool
zero :: Key -> Int -> Bool
zero Key
i Int
m
  = (Int -> Nat
natFromInt (Key -> Int
asInt Key
i)) Nat -> Nat -> Nat
.&. (Int -> Nat
natFromInt Int
m) Nat -> Nat -> Bool
forall a. Eq a => a -> a -> Bool
== Nat
0

{-# INLINE nomatch #-}
nomatch :: Key -> Key -> Mask -> Bool
nomatch :: Key -> Key -> Int -> Bool
nomatch Key
k1 Key
k2 Int
m =
    Int -> Nat
natFromInt (Key -> Int
asInt Key
k1) Nat -> Nat -> Nat
.&. Nat
m' Nat -> Nat -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> Nat
natFromInt (Key -> Int
asInt Key
k2) Nat -> Nat -> Nat
.&. Nat
m'
  where
    m' :: Nat
m' = Nat -> Nat
maskW (Int -> Nat
natFromInt Int
m)

{-# INLINE maskW #-}
maskW :: Nat -> Nat
maskW :: Nat -> Nat
maskW Nat
m = Nat -> Nat
complement (Nat
mNat -> Nat -> Nat
forall a. Num a => a -> a -> a
-Nat
1) Nat -> Nat -> Nat
`xor` Nat
m

{-# INLINE branchMask #-}
branchMask :: Key -> Key -> Mask
branchMask :: Key -> Key -> Int
branchMask Key
k1' Key
k2' =
    Nat -> Int
intFromNat (Nat -> Nat
highestBitMask (Int -> Nat
natFromInt Int
k1 Nat -> Nat -> Nat
`xor` Int -> Nat
natFromInt Int
k2))
  where
    k1 :: Int
k1 = Key -> Int
asInt Key
k1'
    k2 :: Int
k2 = Key -> Int
asInt Key
k2'

highestBitMask :: Nat -> Nat
highestBitMask :: Nat -> Nat
highestBitMask (W# Word#
x) =
    Word# -> Nat
W# (Word# -> Int# -> Word#
uncheckedShiftL# Word#
1## (Word# -> Int#
word2Int# (WORD_SIZE_IN_BITS## `minusWord#` 1## `minusWord#` clz# x)))
{-# INLINE highestBitMask #-}

------------------------------------------------------------------------------
-- Query
------------------------------------------------------------------------------

-- | /O(1)/ True if the queue is empty.
null :: IntPSQ v -> Bool
null :: IntPSQ v -> Bool
null IntPSQ v
Nil = Bool
True
null IntPSQ v
_   = Bool
False

-- | /O(n)/ The number of elements stored in the queue.
size :: IntPSQ v -> Int
size :: IntPSQ v -> Int
size IntPSQ v
Nil               = Int
0
size (Tip Key
_ Prio
_ v
_)       = Int
1
size (Bin Key
_ Prio
_ v
_ Int
_ IntPSQ v
l IntPSQ v
r) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IntPSQ v -> Int
forall v. IntPSQ v -> Int
size IntPSQ v
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IntPSQ v -> Int
forall v. IntPSQ v -> Int
size IntPSQ v
r
-- TODO (SM): benchmark this against a tail-recursive variant

-- | /O(min(n,W))/ The priority and value of a given key, or 'Nothing' if the
-- key is not bound.
lookup :: Key -> IntPSQ v -> Maybe (Prio, v)
lookup :: Key -> IntPSQ v -> Maybe (Prio, v)
lookup Key
k = IntPSQ v -> Maybe (Prio, v)
forall b. IntPSQ b -> Maybe (Prio, b)
go
  where
    go :: IntPSQ b -> Maybe (Prio, b)
go IntPSQ b
t = case IntPSQ b
t of
        IntPSQ b
Nil                -> Maybe (Prio, b)
forall a. Maybe a
Nothing

        Tip Key
k' Prio
p' b
x'
          | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k'        -> (Prio, b) -> Maybe (Prio, b)
forall a. a -> Maybe a
Just (Prio
p', b
x')
          | Bool
otherwise      -> Maybe (Prio, b)
forall a. Maybe a
Nothing

        Bin Key
k' Prio
p' b
x' Int
m IntPSQ b
l IntPSQ b
r
          | Key -> Key -> Int -> Bool
nomatch Key
k Key
k' Int
m -> Maybe (Prio, b)
forall a. Maybe a
Nothing
          | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k'        -> (Prio, b) -> Maybe (Prio, b)
forall a. a -> Maybe a
Just (Prio
p', b
x')
          | Key -> Int -> Bool
zero Key
k Int
m       -> IntPSQ b -> Maybe (Prio, b)
go IntPSQ b
l
          | Bool
otherwise      -> IntPSQ b -> Maybe (Prio, b)
go IntPSQ b
r

-- | /O(1)/ The element with the lowest priority.
findMin :: IntPSQ v -> Maybe (Elem v)
findMin :: IntPSQ v -> Maybe (Elem v)
findMin IntPSQ v
t = case IntPSQ v
t of
    IntPSQ v
Nil             -> Maybe (Elem v)
forall a. Maybe a
Nothing
    Tip Key
k Prio
p v
x       -> Elem v -> Maybe (Elem v)
forall a. a -> Maybe a
Just (Key -> Prio -> v -> Elem v
forall a. Key -> Prio -> a -> Elem a
E Key
k Prio
p v
x)
    Bin Key
k Prio
p v
x Int
_ IntPSQ v
_ IntPSQ v
_ -> Elem v -> Maybe (Elem v)
forall a. a -> Maybe a
Just (Key -> Prio -> v -> Elem v
forall a. Key -> Prio -> a -> Elem a
E Key
k Prio
p v
x)


------------------------------------------------------------------------------
--- Construction
------------------------------------------------------------------------------

-- | /O(1)/ The empty queue.
empty :: IntPSQ v
empty :: IntPSQ v
empty = IntPSQ v
forall v. IntPSQ v
Nil

-- | /O(1)/ Build a queue with one element.
singleton :: Key -> Prio -> v -> IntPSQ v
singleton :: Key -> Prio -> v -> IntPSQ v
singleton = Key -> Prio -> v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v
Tip


------------------------------------------------------------------------------
-- Insertion
------------------------------------------------------------------------------

-- | /O(min(n,W))/ Insert a new key that is *not* present in the priority queue.
{-# INLINABLE unsafeInsertNew #-}
unsafeInsertNew :: Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew :: Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew Key
k Prio
p v
x = IntPSQ v -> IntPSQ v
go
  where
    go :: IntPSQ v -> IntPSQ v
go IntPSQ v
t = case IntPSQ v
t of
      IntPSQ v
Nil       -> Key -> Prio -> v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v
Tip Key
k Prio
p v
x

      Tip Key
k' Prio
p' v
x'
        | (Prio
p, Key
k) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
p', Key
k') -> Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
link Key
k  Prio
p  v
x  Key
k' IntPSQ v
t           IntPSQ v
forall v. IntPSQ v
Nil
        | Bool
otherwise         -> Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
link Key
k' Prio
p' v
x' Key
k  (Key -> Prio -> v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v
Tip Key
k Prio
p v
x) IntPSQ v
forall v. IntPSQ v
Nil

      Bin Key
k' Prio
p' v
x' Int
m IntPSQ v
l IntPSQ v
r
        | Key -> Key -> Int -> Bool
nomatch Key
k Key
k' Int
m ->
            if (Prio
p, Key
k) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
p', Key
k')
              then Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
link Key
k  Prio
p  v
x  Key
k' IntPSQ v
t           IntPSQ v
forall v. IntPSQ v
Nil
              else Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
link Key
k' Prio
p' v
x' Key
k  (Key -> Prio -> v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v
Tip Key
k Prio
p v
x) (Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
m IntPSQ v
l IntPSQ v
r)

        | Bool
otherwise ->
            if (Prio
p, Key
k) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
p', Key
k')
              then
                if Key -> Int -> Bool
zero Key
k' Int
m
                  then Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k  Prio
p  v
x  Int
m (Key -> Prio -> v -> IntPSQ v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew Key
k' Prio
p' v
x' IntPSQ v
l) IntPSQ v
r
                  else Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k  Prio
p  v
x  Int
m IntPSQ v
l (Key -> Prio -> v -> IntPSQ v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew Key
k' Prio
p' v
x' IntPSQ v
r)
              else
                if Key -> Int -> Bool
zero Key
k Int
m
                  then Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k' Prio
p' v
x' Int
m (Key -> Prio -> v -> IntPSQ v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew Key
k  Prio
p  v
x  IntPSQ v
l) IntPSQ v
r
                  else Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k' Prio
p' v
x' Int
m IntPSQ v
l (Key -> Prio -> v -> IntPSQ v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew Key
k  Prio
p  v
x  IntPSQ v
r)

-- | Link
link :: Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
link :: Key -> Prio -> v -> Key -> IntPSQ v -> IntPSQ v -> IntPSQ v
link Key
k Prio
p v
x Key
k' IntPSQ v
k't IntPSQ v
otherTree
  | Key -> Int -> Bool
zero (Int -> Key
Unique Int
m) (Key -> Int
asInt Key
k') = Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k Prio
p v
x Int
m IntPSQ v
k't IntPSQ v
otherTree
  | Bool
otherwise                  = Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k Prio
p v
x Int
m IntPSQ v
otherTree IntPSQ v
k't
  where
    m :: Int
m = Key -> Key -> Int
branchMask Key
k Key
k'


------------------------------------------------------------------------------
-- Delete/Alter
------------------------------------------------------------------------------

-- | /O(min(n,W))/ Delete a key and its priority and value from the queue. When
-- the key is not a member of the queue, the original queue is returned.
{-# INLINABLE delete #-}
delete :: Key -> IntPSQ v -> IntPSQ v
delete :: Key -> IntPSQ v -> IntPSQ v
delete Key
k = IntPSQ v -> IntPSQ v
forall v. IntPSQ v -> IntPSQ v
go
  where
    go :: IntPSQ v -> IntPSQ v
go IntPSQ v
t = case IntPSQ v
t of
        IntPSQ v
Nil           -> IntPSQ v
forall v. IntPSQ v
Nil

        Tip Key
k' Prio
_ v
_
          | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k'   -> IntPSQ v
forall v. IntPSQ v
Nil
          | Bool
otherwise -> IntPSQ v
t

        Bin Key
k' Prio
p' v
x' Int
m IntPSQ v
l IntPSQ v
r
          | Key -> Key -> Int -> Bool
nomatch Key
k Key
k' Int
m -> IntPSQ v
t
          | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k'        -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
m IntPSQ v
l IntPSQ v
r
          | Key -> Int -> Bool
zero Key
k Int
m       -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkL Key
k' Prio
p' v
x' Int
m (IntPSQ v -> IntPSQ v
go IntPSQ v
l) IntPSQ v
r
          | Bool
otherwise      -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkR Key
k' Prio
p' v
x' Int
m IntPSQ v
l      (IntPSQ v -> IntPSQ v
go IntPSQ v
r)

-- | /O(min(n,W))/ Delete the binding with the least priority, and return the
-- rest of the queue stripped of that binding. In case the queue is empty, the
-- empty queue is returned again.
{-# INLINE deleteMin #-}
deleteMin :: IntPSQ v -> IntPSQ v
deleteMin :: IntPSQ v -> IntPSQ v
deleteMin IntPSQ v
t = case IntPSQ v -> Maybe (Elem v, IntPSQ v)
forall v. IntPSQ v -> Maybe (Elem v, IntPSQ v)
minView IntPSQ v
t of
    Maybe (Elem v, IntPSQ v)
Nothing      -> IntPSQ v
t
    Just (Elem v
_, IntPSQ v
t') -> IntPSQ v
t'


adjust
    :: (Prio -> Prio)
    -> Key
    -> PSQ a
    -> PSQ a
adjust :: (Prio -> Prio) -> Key -> PSQ a -> PSQ a
adjust Prio -> Prio
f Key
k PSQ a
q = case (Maybe (Prio, a) -> ((), Maybe (Prio, a)))
-> Key -> PSQ a -> ((), PSQ a)
forall v b.
(Maybe (Prio, v) -> (b, Maybe (Prio, v)))
-> Key -> IntPSQ v -> (b, IntPSQ v)
alter Maybe (Prio, a) -> ((), Maybe (Prio, a))
forall b. Maybe (Prio, b) -> ((), Maybe (Prio, b))
g Key
k PSQ a
q of (()
_, PSQ a
q') -> PSQ a
q'
  where g :: Maybe (Prio, b) -> ((), Maybe (Prio, b))
g (Just (Prio
p, b
v)) = ((), (Prio, b) -> Maybe (Prio, b)
forall a. a -> Maybe a
Just ((Prio -> Prio
f Prio
p), b
v))
        g Maybe (Prio, b)
Nothing       = ((), Maybe (Prio, b)
forall a. Maybe a
Nothing)

{-# INLINE adjust #-}

-- | /O(min(n,W))/ The expression @alter f k queue@ alters the value @x@ at @k@,
-- or absence thereof. 'alter' can be used to insert, delete, or update a value
-- in a queue. It also allows you to calculate an additional value @b@.
{-# INLINE alter #-}
alter
    :: (Maybe (Prio, v) -> (b, Maybe (Prio, v)))
    -> Key
    -> IntPSQ v
    -> (b, IntPSQ v)
alter :: (Maybe (Prio, v) -> (b, Maybe (Prio, v)))
-> Key -> IntPSQ v -> (b, IntPSQ v)
alter Maybe (Prio, v) -> (b, Maybe (Prio, v))
f = \Key
k IntPSQ v
t0 ->
    let (IntPSQ v
t, Maybe (Prio, v)
mbX) = case Key -> IntPSQ v -> Maybe (Prio, v, IntPSQ v)
forall v. Key -> IntPSQ v -> Maybe (Prio, v, IntPSQ v)
deleteView Key
k IntPSQ v
t0 of
                            Maybe (Prio, v, IntPSQ v)
Nothing          -> (IntPSQ v
t0, Maybe (Prio, v)
forall a. Maybe a
Nothing)
                            Just (Prio
p, v
v, IntPSQ v
t0') -> (IntPSQ v
t0', (Prio, v) -> Maybe (Prio, v)
forall a. a -> Maybe a
Just (Prio
p, v
v))
    in case Maybe (Prio, v) -> (b, Maybe (Prio, v))
f Maybe (Prio, v)
mbX of
          (b
b, Maybe (Prio, v)
mbX') ->
            (b
b, IntPSQ v -> ((Prio, v) -> IntPSQ v) -> Maybe (Prio, v) -> IntPSQ v
forall p t. p -> (t -> p) -> Maybe t -> p
maybe IntPSQ v
t (\(Prio
p, v
v) -> Key -> Prio -> v -> IntPSQ v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v -> IntPSQ v
unsafeInsertNew Key
k Prio
p v
v IntPSQ v
t) Maybe (Prio, v)
mbX')
    where
        maybe :: p -> (t -> p) -> Maybe t -> p
maybe p
_ t -> p
g (Just t
x)  = t -> p
g t
x
        maybe p
def t -> p
_ Maybe t
Nothing = p
def

-- | Smart constructor for a 'Bin' node whose left subtree could have become
-- 'Nil'.
{-# INLINE binShrinkL #-}
binShrinkL :: Key -> Prio -> v -> Mask -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkL :: Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkL Key
k Prio
p v
x Int
m IntPSQ v
Nil IntPSQ v
r = case IntPSQ v
r of IntPSQ v
Nil -> Key -> Prio -> v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v
Tip Key
k Prio
p v
x; IntPSQ v
_ -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k Prio
p v
x Int
m IntPSQ v
forall v. IntPSQ v
Nil IntPSQ v
r
binShrinkL Key
k Prio
p v
x Int
m IntPSQ v
l   IntPSQ v
r = Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k Prio
p v
x Int
m IntPSQ v
l IntPSQ v
r

-- | Smart constructor for a 'Bin' node whose right subtree could have become
-- 'Nil'.
{-# INLINE binShrinkR #-}
binShrinkR :: Key -> Prio -> v -> Mask -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkR :: Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkR Key
k Prio
p v
x Int
m IntPSQ v
l IntPSQ v
Nil = case IntPSQ v
l of IntPSQ v
Nil -> Key -> Prio -> v -> IntPSQ v
forall v. Key -> Prio -> v -> IntPSQ v
Tip Key
k Prio
p v
x; IntPSQ v
_ -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k Prio
p v
x Int
m IntPSQ v
l IntPSQ v
forall v. IntPSQ v
Nil
binShrinkR Key
k Prio
p v
x Int
m IntPSQ v
l IntPSQ v
r   = Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
k Prio
p v
x Int
m IntPSQ v
l IntPSQ v
r

------------------------------------------------------------------------------
-- Lists
------------------------------------------------------------------------------

-- | /O(n)/ Convert a queue to a list of (key, priority, value) tuples. The
-- order of the list is not specified.
toList :: IntPSQ v -> [Elem v]
toList :: IntPSQ v -> [Elem v]
toList =
    [Elem v] -> IntPSQ v -> [Elem v]
forall a. [Elem a] -> IntPSQ a -> [Elem a]
go []
  where
    go :: [Elem a] -> IntPSQ a -> [Elem a]
go [Elem a]
acc IntPSQ a
Nil                   = [Elem a]
acc
    go [Elem a]
acc (Tip Key
k' Prio
p' a
x')        = (Key -> Prio -> a -> Elem a
forall a. Key -> Prio -> a -> Elem a
E Key
k' Prio
p' a
x') Elem a -> [Elem a] -> [Elem a]
forall a. a -> [a] -> [a]
: [Elem a]
acc
    go [Elem a]
acc (Bin Key
k' Prio
p' a
x' Int
_m IntPSQ a
l IntPSQ a
r) = (Key -> Prio -> a -> Elem a
forall a. Key -> Prio -> a -> Elem a
E Key
k' Prio
p' a
x') Elem a -> [Elem a] -> [Elem a]
forall a. a -> [a] -> [a]
: [Elem a] -> IntPSQ a -> [Elem a]
go ([Elem a] -> IntPSQ a -> [Elem a]
go [Elem a]
acc IntPSQ a
r) IntPSQ a
l


------------------------------------------------------------------------------
-- Views
------------------------------------------------------------------------------

-- | /O(min(n,W))/ Delete a key and its priority and value from the queue. If
-- the key was present, the associated priority and value are returned in
-- addition to the updated queue.
{-# INLINABLE deleteView #-}
deleteView :: Key -> IntPSQ v -> Maybe (Prio, v, IntPSQ v)
deleteView :: Key -> IntPSQ v -> Maybe (Prio, v, IntPSQ v)
deleteView Key
k IntPSQ v
t0 =
    case IntPSQ v -> (# IntPSQ v, Maybe (Prio, v) #)
forall b. IntPSQ b -> (# IntPSQ b, Maybe (Prio, b) #)
delFrom IntPSQ v
t0 of
      (# IntPSQ v
_, Maybe (Prio, v)
Nothing     #) -> Maybe (Prio, v, IntPSQ v)
forall a. Maybe a
Nothing
      (# IntPSQ v
t, Just (Prio
p, v
x) #) -> (Prio, v, IntPSQ v) -> Maybe (Prio, v, IntPSQ v)
forall a. a -> Maybe a
Just (Prio
p, v
x, IntPSQ v
t)
  where
    delFrom :: IntPSQ b -> (# IntPSQ b, Maybe (Prio, b) #)
delFrom IntPSQ b
t = case IntPSQ b
t of
      IntPSQ b
Nil -> (# IntPSQ b
forall v. IntPSQ v
Nil, Maybe (Prio, b)
forall a. Maybe a
Nothing #)

      Tip Key
k' Prio
p' b
x'
        | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k'   -> (# IntPSQ b
forall v. IntPSQ v
Nil, (Prio, b) -> Maybe (Prio, b)
forall a. a -> Maybe a
Just (Prio
p', b
x') #)
        | Bool
otherwise -> (# IntPSQ b
t,   Maybe (Prio, b)
forall a. Maybe a
Nothing       #)

      Bin Key
k' Prio
p' b
x' Int
m IntPSQ b
l IntPSQ b
r
        | Key -> Key -> Int -> Bool
nomatch Key
k Key
k' Int
m -> (# IntPSQ b
t, Maybe (Prio, b)
forall a. Maybe a
Nothing #)
        | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k'   -> let t' :: IntPSQ b
t' = Int -> IntPSQ b -> IntPSQ b -> IntPSQ b
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
m IntPSQ b
l IntPSQ b
r
                       in  IntPSQ b
t' IntPSQ b
-> (# IntPSQ b, Maybe (Prio, b) #)
-> (# IntPSQ b, Maybe (Prio, b) #)
`seq` (# IntPSQ b
t', (Prio, b) -> Maybe (Prio, b)
forall a. a -> Maybe a
Just (Prio
p', b
x') #)

        | Key -> Int -> Bool
zero Key
k Int
m  -> case IntPSQ b -> (# IntPSQ b, Maybe (Prio, b) #)
delFrom IntPSQ b
l of
                         (# IntPSQ b
l', Maybe (Prio, b)
mbPX #) -> let t' :: IntPSQ b
t' = Key -> Prio -> b -> Int -> IntPSQ b -> IntPSQ b -> IntPSQ b
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkL Key
k' Prio
p' b
x' Int
m IntPSQ b
l' IntPSQ b
r
                                           in  IntPSQ b
t' IntPSQ b
-> (# IntPSQ b, Maybe (Prio, b) #)
-> (# IntPSQ b, Maybe (Prio, b) #)
`seq` (# IntPSQ b
t', Maybe (Prio, b)
mbPX #)

        | Bool
otherwise -> case IntPSQ b -> (# IntPSQ b, Maybe (Prio, b) #)
delFrom IntPSQ b
r of
                         (# IntPSQ b
r', Maybe (Prio, b)
mbPX #) -> let t' :: IntPSQ b
t' = Key -> Prio -> b -> Int -> IntPSQ b -> IntPSQ b -> IntPSQ b
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
binShrinkR Key
k' Prio
p' b
x' Int
m IntPSQ b
l  IntPSQ b
r'
                                           in  IntPSQ b
t' IntPSQ b
-> (# IntPSQ b, Maybe (Prio, b) #)
-> (# IntPSQ b, Maybe (Prio, b) #)
`seq` (# IntPSQ b
t', Maybe (Prio, b)
mbPX #)

-- | /O(min(n,W))/ Retrieve the binding with the least priority, and the
-- rest of the queue stripped of that binding.
{-# INLINE minView #-}
minView :: IntPSQ v -> Maybe (Elem v, IntPSQ v)
minView :: IntPSQ v -> Maybe (Elem v, IntPSQ v)
minView IntPSQ v
t = case IntPSQ v
t of
    IntPSQ v
Nil             -> Maybe (Elem v, IntPSQ v)
forall a. Maybe a
Nothing
    Tip Key
k Prio
p v
x       -> (Elem v, IntPSQ v) -> Maybe (Elem v, IntPSQ v)
forall a. a -> Maybe a
Just (Key -> Prio -> v -> Elem v
forall a. Key -> Prio -> a -> Elem a
E Key
k Prio
p v
x, IntPSQ v
forall v. IntPSQ v
Nil)
    Bin Key
k Prio
p v
x Int
m IntPSQ v
l IntPSQ v
r -> (Elem v, IntPSQ v) -> Maybe (Elem v, IntPSQ v)
forall a. a -> Maybe a
Just (Key -> Prio -> v -> Elem v
forall a. Key -> Prio -> a -> Elem a
E Key
k Prio
p v
x, Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
m IntPSQ v
l IntPSQ v
r)

-- | Return a list of elements ordered by key whose priorities are at most @pt@,
-- and the rest of the queue stripped of these elements.  The returned list of
-- elements can be in any order: no guarantees there.
{-# INLINABLE atMost #-}
atMost :: Prio -> IntPSQ v -> ([Elem v], IntPSQ v)
atMost :: Prio -> IntPSQ v -> ([Elem v], IntPSQ v)
atMost Prio
pt IntPSQ v
t0 = [Elem v] -> IntPSQ v -> ([Elem v], IntPSQ v)
forall a. [Elem a] -> IntPSQ a -> ([Elem a], IntPSQ a)
go [] IntPSQ v
t0
  where
    go :: [Elem a] -> IntPSQ a -> ([Elem a], IntPSQ a)
go [Elem a]
acc IntPSQ a
t = case IntPSQ a
t of
        IntPSQ a
Nil             -> ([Elem a]
acc, IntPSQ a
t)
        Tip Key
k Prio
p a
x
            | Prio
p Prio -> Prio -> Bool
forall a. Ord a => a -> a -> Bool
> Prio
pt    -> ([Elem a]
acc, IntPSQ a
t)
            | Bool
otherwise -> ((Key -> Prio -> a -> Elem a
forall a. Key -> Prio -> a -> Elem a
E Key
k Prio
p a
x) Elem a -> [Elem a] -> [Elem a]
forall a. a -> [a] -> [a]
: [Elem a]
acc, IntPSQ a
forall v. IntPSQ v
Nil)

        Bin Key
k Prio
p a
x Int
m IntPSQ a
l IntPSQ a
r
            | Prio
p Prio -> Prio -> Bool
forall a. Ord a => a -> a -> Bool
> Prio
pt    -> ([Elem a]
acc, IntPSQ a
t)
            | Bool
otherwise ->
                let ([Elem a]
acc',  IntPSQ a
l') = [Elem a] -> IntPSQ a -> ([Elem a], IntPSQ a)
go [Elem a]
acc  IntPSQ a
l
                    ([Elem a]
acc'', IntPSQ a
r') = [Elem a] -> IntPSQ a -> ([Elem a], IntPSQ a)
go [Elem a]
acc' IntPSQ a
r
                in  ((Key -> Prio -> a -> Elem a
forall a. Key -> Prio -> a -> Elem a
E Key
k Prio
p a
x) Elem a -> [Elem a] -> [Elem a]
forall a. a -> [a] -> [a]
: [Elem a]
acc'', Int -> IntPSQ a -> IntPSQ a -> IntPSQ a
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
m IntPSQ a
l' IntPSQ a
r')


------------------------------------------------------------------------------
-- Traversal
------------------------------------------------------------------------------

-- | Internal function that merges two *disjoint* 'IntPSQ's that share the
-- same prefix mask.
{-# INLINABLE merge #-}
merge :: Mask -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge :: Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
m IntPSQ v
l IntPSQ v
r = case IntPSQ v
l of
    IntPSQ v
Nil -> IntPSQ v
r

    Tip Key
lk Prio
lp v
lx ->
      case IntPSQ v
r of
        IntPSQ v
Nil                     -> IntPSQ v
l
        Tip Key
rk Prio
rp v
rx
          | (Prio
lp, Key
lk) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
rp, Key
rk) -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
lk Prio
lp v
lx Int
m IntPSQ v
forall v. IntPSQ v
Nil IntPSQ v
r
          | Bool
otherwise           -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
rk Prio
rp v
rx Int
m IntPSQ v
l   IntPSQ v
forall v. IntPSQ v
Nil
        Bin Key
rk Prio
rp v
rx Int
rm IntPSQ v
rl IntPSQ v
rr
          | (Prio
lp, Key
lk) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
rp, Key
rk) -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
lk Prio
lp v
lx Int
m IntPSQ v
forall v. IntPSQ v
Nil IntPSQ v
r
          | Bool
otherwise           -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
rk Prio
rp v
rx Int
m IntPSQ v
l   (Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
rm IntPSQ v
rl IntPSQ v
rr)

    Bin Key
lk Prio
lp v
lx Int
lm IntPSQ v
ll IntPSQ v
lr ->
      case IntPSQ v
r of
        IntPSQ v
Nil                     -> IntPSQ v
l
        Tip Key
rk Prio
rp v
rx
          | (Prio
lp, Key
lk) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
rp, Key
rk) -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
lk Prio
lp v
lx Int
m (Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
lm IntPSQ v
ll IntPSQ v
lr) IntPSQ v
r
          | Bool
otherwise           -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
rk Prio
rp v
rx Int
m IntPSQ v
l                IntPSQ v
forall v. IntPSQ v
Nil
        Bin Key
rk Prio
rp v
rx Int
rm IntPSQ v
rl IntPSQ v
rr
          | (Prio
lp, Key
lk) (Prio, Key) -> (Prio, Key) -> Bool
forall a. Ord a => a -> a -> Bool
< (Prio
rp, Key
rk) -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
lk Prio
lp v
lx Int
m (Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
lm IntPSQ v
ll IntPSQ v
lr) IntPSQ v
r
          | Bool
otherwise           -> Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v.
Key -> Prio -> v -> Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
Bin Key
rk Prio
rp v
rx Int
m IntPSQ v
l                (Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
forall v. Int -> IntPSQ v -> IntPSQ v -> IntPSQ v
merge Int
rm IntPSQ v
rl IntPSQ v
rr)