content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE RecordWildCards #-}
  
  module Heap (
      PriorityQueue (..),
      Precedence (..),
      newPrecedence,
      empty,
      isEmpty,
      enqueue,
      dequeue,
      delete,
  ) where
  
  #if __GLASGOW_HASKELL__ < 709
  import Control.Applicative ((<$>))
  import Data.Word (Word)
  #endif
  import Data.Array (Array, listArray, (!))
  import Data.Heap (Heap)
  import qualified Data.Heap as H
  
  ----------------------------------------------------------------
  
  type Key = Int
  type Weight = Int
  type Deficit = Word -- Deficit can be overflowed
  
  data Precedence = Precedence
      { deficit :: {-# UNPACK #-} !Deficit
      , weight :: {-# UNPACK #-} !Weight
      , -- stream dependency, used by the upper layer
        dependency :: {-# UNPACK #-} !Key
      }
      deriving (Show)
  
  -- | For test only
  newPrecedence :: Weight -> Precedence
  newPrecedence w = Precedence 0 w 0
  
  instance Eq Precedence where
      Precedence d1 _ _ == Precedence d2 _ _ = d1 == d2
  
  instance Ord Precedence where
      -- This is correct even if one of them is overflowed
      Precedence d1 _ _ < Precedence d2 _ _ = d1 /= d2 && d2 - d1 <= deficitStepsW
      Precedence d1 _ _ <= Precedence d2 _ _ = d2 - d1 <= deficitStepsW
  
  data Entry a = Entry Key Precedence a
  
  instance Eq (Entry a) where
      Entry _ p1 _ == Entry _ p2 _ = p1 == p2
  
  instance Ord (Entry a) where
      Entry _ p1 _ < Entry _ p2 _ = p1 < p2
      Entry _ p1 _ <= Entry _ p2 _ = p1 <= p2
  
  -- FIXME: The base (Word64) would be overflowed.
  --        In that case, the heap must be re-constructed.
  data PriorityQueue a = PriorityQueue
      { baseDeficit :: {-# UNPACK #-} !Deficit
      , queue :: Heap (Entry a)
      }
  
  ----------------------------------------------------------------
  
  magicDeficit :: Deficit
  magicDeficit = 0
  
  deficitSteps :: Int
  deficitSteps = 65536
  
  deficitStepsW :: Word
  deficitStepsW = fromIntegral deficitSteps
  
  deficitList :: [Deficit]
  deficitList = map calc idxs
    where
      idxs = [1 .. 256] :: [Double]
      calc w = round (fromIntegral deficitSteps / w)
  
  deficitTable :: Array Int Deficit
  deficitTable = listArray (1, 256) deficitList
  
  weightToDeficit :: Weight -> Deficit
  weightToDeficit w = deficitTable ! w
  
  ----------------------------------------------------------------
  
  empty :: PriorityQueue a
  empty = PriorityQueue 0 H.empty
  
  isEmpty :: PriorityQueue a -> Bool
  isEmpty (PriorityQueue _ h) = H.null h
  
  enqueue :: Key -> Precedence -> a -> PriorityQueue a -> PriorityQueue a
  enqueue k p v PriorityQueue{..} = PriorityQueue b queue'
    where
      !d = weightToDeficit (weight p)
      !b = if deficit p == magicDeficit then baseDeficit else deficit p
      !deficit' = max (b + d) baseDeficit
      !p' = p{deficit = deficit'}
      !queue' = H.insert (Entry k p' v) queue
  
  dequeue :: PriorityQueue a -> Maybe (Key, Precedence, a, PriorityQueue a)
  dequeue (PriorityQueue _ heap) = case H.uncons heap of
      Nothing -> Nothing
      Just (Entry k p v, heap') -> Just (k, p, v, PriorityQueue (deficit p) heap')
  
  delete :: Key -> PriorityQueue a -> (Maybe a, PriorityQueue a)
  delete k (PriorityQueue base heap) = (mv, PriorityQueue base' heap')
    where
      !(h, heap') = H.partition (\(Entry k' _ _) -> k' == k) heap
      mv = case H.viewMin h of
          Nothing -> Nothing
          Just (Entry _ _ v, _) -> Just v
      base' = case H.viewMin heap of
          Nothing -> base
          Just (Entry k' p _, _)
              | k == k' -> deficit p
              | otherwise -> base
  