{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}

module Network.HTTP2.Priority.PSQ (
    Key
  , Precedence(..)
  , newPrecedence
  , PriorityQueue(..)
  , empty
  , isEmpty
  , enqueue
  , dequeue
  , delete
  ) where

import Data.Array (Array, listArray, (!))
import Data.IntPSQ (IntPSQ)
import qualified Data.IntPSQ as P

----------------------------------------------------------------

type Key = Int
type Weight = Int
type Deficit = Word -- Deficit can be overflowed

-- | Internal representation of priority in priority queues.
--   The precedence of a dequeued entry should be specified
--   to enqueue when the entry is enqueued again.
data Precedence = Precedence {
    Precedence -> Deficit
deficit    :: {-# UNPACK #-} !Deficit
  , Precedence -> Weight
weight     :: {-# UNPACK #-} !Weight
  -- stream dependency, used by the upper layer
  , Precedence -> Weight
dependency :: {-# UNPACK #-} !Key
  } deriving Weight -> Precedence -> ShowS
[Precedence] -> ShowS
Precedence -> String
(Weight -> Precedence -> ShowS)
-> (Precedence -> String)
-> ([Precedence] -> ShowS)
-> Show Precedence
forall a.
(Weight -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Precedence] -> ShowS
$cshowList :: [Precedence] -> ShowS
show :: Precedence -> String
$cshow :: Precedence -> String
showsPrec :: Weight -> Precedence -> ShowS
$cshowsPrec :: Weight -> Precedence -> ShowS
Show

-- | For test only
newPrecedence :: Weight -> Precedence
newPrecedence :: Weight -> Precedence
newPrecedence w :: Weight
w = Deficit -> Weight -> Weight -> Precedence
Precedence 0 Weight
w 0

instance Eq Precedence where
  Precedence d1 :: Deficit
d1 _ _ == :: Precedence -> Precedence -> Bool
== Precedence d2 :: Deficit
d2 _ _ = Deficit
d1 Deficit -> Deficit -> Bool
forall a. Eq a => a -> a -> Bool
== Deficit
d2

instance Ord Precedence where
  -- This is correct even if one of them is overflowed
  Precedence d1 :: Deficit
d1 _ _ < :: Precedence -> Precedence -> Bool
<  Precedence d2 :: Deficit
d2 _ _ = Deficit
d1 Deficit -> Deficit -> Bool
forall a. Eq a => a -> a -> Bool
/= Deficit
d2 Bool -> Bool -> Bool
&& Deficit
d2 Deficit -> Deficit -> Deficit
forall a. Num a => a -> a -> a
- Deficit
d1 Deficit -> Deficit -> Bool
forall a. Ord a => a -> a -> Bool
<= Deficit
deficitStepsW
  Precedence d1 :: Deficit
d1 _ _ <= :: Precedence -> Precedence -> Bool
<= Precedence d2 :: Deficit
d2 _ _ = Deficit
d2 Deficit -> Deficit -> Deficit
forall a. Num a => a -> a -> a
- Deficit
d1 Deficit -> Deficit -> Bool
forall a. Ord a => a -> a -> Bool
<= Deficit
deficitStepsW

type Heap a = IntPSQ Precedence a

data PriorityQueue a = PriorityQueue {
    PriorityQueue a -> Deficit
baseDeficit :: {-# UNPACK #-} !Deficit
  , PriorityQueue a -> Heap a
queue :: !(Heap a)
  }

----------------------------------------------------------------

deficitSteps :: Int
deficitSteps :: Weight
deficitSteps = 65536

deficitStepsW :: Word
deficitStepsW :: Deficit
deficitStepsW = Weight -> Deficit
forall a b. (Integral a, Num b) => a -> b
fromIntegral Weight
deficitSteps

deficitList :: [Deficit]
deficitList :: [Deficit]
deficitList = (Double -> Deficit) -> [Double] -> [Deficit]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Deficit
forall a b. (RealFrac a, Integral b) => a -> b
calc [Double]
idxs
  where
    idxs :: [Double]
idxs = [1..256] :: [Double]
    calc :: a -> b
calc w :: a
w = a -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Weight -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Weight
deficitSteps a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
w)

deficitTable :: Array Int Deficit
deficitTable :: Array Weight Deficit
deficitTable = (Weight, Weight) -> [Deficit] -> Array Weight Deficit
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (1,256) [Deficit]
deficitList

weightToDeficit :: Weight -> Deficit
weightToDeficit :: Weight -> Deficit
weightToDeficit w :: Weight
w = Array Weight Deficit
deficitTable Array Weight Deficit -> Weight -> Deficit
forall i e. Ix i => Array i e -> i -> e
! Weight
w

----------------------------------------------------------------

empty :: PriorityQueue a
empty :: PriorityQueue a
empty = Deficit -> Heap a -> PriorityQueue a
forall a. Deficit -> Heap a -> PriorityQueue a
PriorityQueue 0 Heap a
forall p v. IntPSQ p v
P.empty

isEmpty :: PriorityQueue a -> Bool
isEmpty :: PriorityQueue a -> Bool
isEmpty PriorityQueue{..} = Heap a -> Bool
forall p v. IntPSQ p v -> Bool
P.null Heap a
queue

enqueue :: Key -> Precedence -> a -> PriorityQueue a -> PriorityQueue a
enqueue :: Weight -> Precedence -> a -> PriorityQueue a -> PriorityQueue a
enqueue k :: Weight
k p :: Precedence
p@Precedence{..} v :: a
v PriorityQueue{..} =
    Deficit -> Heap a -> PriorityQueue a
forall a. Deficit -> Heap a -> PriorityQueue a
PriorityQueue Deficit
baseDeficit Heap a
queue'
  where
    !d :: Deficit
d = Weight -> Deficit
weightToDeficit Weight
weight
    !b :: Deficit
b = if Deficit
deficit Deficit -> Deficit -> Bool
forall a. Eq a => a -> a -> Bool
== 0 then Deficit
baseDeficit else Deficit
deficit
    !deficit' :: Deficit
deficit' = Deficit -> Deficit -> Deficit
forall a. Ord a => a -> a -> a
max (Deficit
b Deficit -> Deficit -> Deficit
forall a. Num a => a -> a -> a
+ Deficit
d) Deficit
baseDeficit
    !p' :: Precedence
p' = Precedence
p { deficit :: Deficit
deficit = Deficit
deficit' }
    !queue' :: Heap a
queue' = Weight -> Precedence -> a -> Heap a -> Heap a
forall p v. Ord p => Weight -> p -> v -> IntPSQ p v -> IntPSQ p v
P.insert Weight
k Precedence
p' a
v Heap a
queue

dequeue :: PriorityQueue a -> Maybe (Key, Precedence, a, PriorityQueue a)
dequeue :: PriorityQueue a -> Maybe (Weight, Precedence, a, PriorityQueue a)
dequeue PriorityQueue{..} = case Heap a -> Maybe (Weight, Precedence, a, Heap a)
forall p v. Ord p => IntPSQ p v -> Maybe (Weight, p, v, IntPSQ p v)
P.minView Heap a
queue of
    Nothing                -> Maybe (Weight, Precedence, a, PriorityQueue a)
forall a. Maybe a
Nothing
    Just (k :: Weight
k, p :: Precedence
p, v :: a
v, queue' :: Heap a
queue') -> let !base :: Deficit
base = Precedence -> Deficit
deficit Precedence
p
                              in (Weight, Precedence, a, PriorityQueue a)
-> Maybe (Weight, Precedence, a, PriorityQueue a)
forall a. a -> Maybe a
Just (Weight
k, Precedence
p, a
v, Deficit -> Heap a -> PriorityQueue a
forall a. Deficit -> Heap a -> PriorityQueue a
PriorityQueue Deficit
base Heap a
queue')

delete :: Key -> PriorityQueue a -> (Maybe a, PriorityQueue a)
delete :: Weight -> PriorityQueue a -> (Maybe a, PriorityQueue a)
delete k :: Weight
k q :: PriorityQueue a
q@PriorityQueue{..} = case (Maybe (Precedence, a) -> (Maybe a, Maybe (Precedence, a)))
-> Weight -> Heap a -> (Maybe a, Heap a)
forall p v b.
Ord p =>
(Maybe (p, v) -> (b, Maybe (p, v)))
-> Weight -> IntPSQ p v -> (b, IntPSQ p v)
P.alter Maybe (Precedence, a) -> (Maybe a, Maybe (Precedence, a))
forall a a a. Maybe (a, a) -> (Maybe a, Maybe a)
f Weight
k Heap a
queue of
    (mv :: Maybe a
mv@(Just _), queue' :: Heap a
queue') -> case Heap a -> Maybe (Weight, Precedence, a, Heap a)
forall p v. Ord p => IntPSQ p v -> Maybe (Weight, p, v, IntPSQ p v)
P.minView Heap a
queue of
        Nothing          -> String -> (Maybe a, PriorityQueue a)
forall a. HasCallStack => String -> a
error "delete"
        Just (k' :: Weight
k',p' :: Precedence
p',_,_)
          | Weight
k' Weight -> Weight -> Bool
forall a. Eq a => a -> a -> Bool
== Weight
k      -> (Maybe a
mv, Deficit -> Heap a -> PriorityQueue a
forall a. Deficit -> Heap a -> PriorityQueue a
PriorityQueue (Precedence -> Deficit
deficit Precedence
p') Heap a
queue')
          | Bool
otherwise    -> (Maybe a
mv, Deficit -> Heap a -> PriorityQueue a
forall a. Deficit -> Heap a -> PriorityQueue a
PriorityQueue Deficit
baseDeficit Heap a
queue')
    (Nothing, _)         -> (Maybe a
forall a. Maybe a
Nothing, PriorityQueue a
q)
  where
    f :: Maybe (a, a) -> (Maybe a, Maybe a)
f Nothing      = (Maybe a
forall a. Maybe a
Nothing, Maybe a
forall a. Maybe a
Nothing)
    f (Just (_,v :: a
v)) = (a -> Maybe a
forall a. a -> Maybe a
Just a
v,  Maybe a
forall a. Maybe a
Nothing)