MHeap.hs

{-# LANGUAGE BangPatterns, CPP, MagicHash, UnboxedTuples, GeneralizedNewtypeDeriving, NoImplicitPrelude #-}
module Data.MHeap1 where

import GHC.Base

data MHeap a = MHeap { mhValues   :: SmallMutableArray# RealWorld a
                     , mhCount    :: Int#
                     , mhCapacity :: Int#
                     }
-- | Creates an empty heap with capacity 0.
empty :: IO (MHeap a)
empty = IO $ \s ->
  case newSmallArray# 0# undefined s of
    (# s', marr #) -> (# s', MHeap marr 0# 0# #)

-- | Accesses the smallest element in the heap.
unsafeTop :: MHeap a -> IO a
unsafeTop mheap = IO $ \s ->
  readSmallArray# (mhValues mheap) 0# s

-- | Check if the heap needs resizing when adding more elements.
resizeIfNecessary :: MHeap a
                  -> State# RealWorld
                  -> (# State# RealWorld, MHeap a #)
resizeIfNecessary mh0 s0
  | isTrue# (mhCount mh0 ==# mhCapacity mh0) = resize mh0 s0
  | otherwise = (# s0, mh0 #)
  where
    resize mh s =
      let !newCap = let capPlusOne = mhCapacity mh +# 1#
                        capOverallocate = quotInt# (mhCapacity mh *# 3#) 2#
                    in if isTrue# (capOverallocate ># capPlusOne)
                       then capOverallocate
                       else capPlusOne
      in case newSmallArray# newCap undefined s of
        (# s', marr' #) ->
          case copySmallMutableArray# (mhValues mh) 0# marr' 0# (mhCount mh) s' of
            s'' -> (# s'', MHeap marr' (mhCount mh) newCap #)

-- | Restores the heap property from right to left.
adjustUpwards :: (a -> a -> Ordering)
              -> SmallMutableArray# RealWorld a
              -> Int#
              -> a
              -> State# RealWorld
              -> State# RealWorld
adjustUpwards cmp arr i a s | isTrue# (i ># 0#) =
                          let parent = (i -# 1#) `quotInt#` 2#
                          in case readSmallArray# arr parent s of
                            (# s', p #) -> case cmp p a of
                              GT -> case writeSmallArray# arr i p s' of
                                s'' -> adjustUpwards cmp arr parent a s''
                              _ -> writeSmallArray# arr i a s'
                        | otherwise = writeSmallArray# arr i a s

-- | Restores the heap property from left to right.
adjustDownwards :: (a -> a -> Ordering)
              -> SmallMutableArray# RealWorld a -- ^ Array in heap form
              -> Int# -- ^ the index to start heapifing
              -> Int# -- ^ the length of the array (may differ from capacity)
              -> a    -- ^ the element to insert
              -> State# RealWorld
              -> State# RealWorld
adjustDownwards cmp arr i len a s
  | isTrue# (left <# len) =
      let (# s''', i', c #) = case readSmallArray# arr left s of
            (# s', lc #) | isTrue# (right <# len) ->
                           case readSmallArray# arr right s' of
                             (# s'', rc #) -> case cmp lc rc of
                               GT -> (# s'', right, rc #)
                               _  -> (# s'', left, lc #)
                         | otherwise -> (# s', left, lc #)
      in case cmp a c of
        GT -> case writeSmallArray# arr i c s''' of
          s'''' -> adjustDownwards cmp arr i' len a s''''
        _  -> writeSmallArray# arr i a s'''
  | otherwise = writeSmallArray# arr i a s
  where
    left = 1# +# 2# *# i
    right = left +# 1#

insert :: (a -> a -> Ordering) -> a -> MHeap a -> IO (MHeap a)
insert cmp a mh =
  IO $ \s -> case resizeIfNecessary mh s of
    (# s', mh' #) -> case adjustUpwards cmp (mhValues mh') (mhCount mh') a s' of
      s'' -> (# s'', mh' { mhCount = mhCount mh' +# 1#}  #)

minView :: (a -> a -> Ordering) -> MHeap a -> IO (Maybe (a, MHeap a))
minView cmp mh
  | isTrue# (mhCount mh ># 0#) = IO $ \s ->
      case readSmallArray# (mhValues mh) 0# s of
        (# s', a #) -> case deleteAtIndex cmp (mhValues mh) 0# (mhCount mh) s' of
          s'' -> (# s'', Just (a, MHeap (mhValues mh) (mhCount mh -# 1#) (mhCapacity mh)) #)
  | otherwise = IO $ \s -> (# s, Nothing #)

deleteAtIndex :: (a -> a -> Ordering)
              -> SmallMutableArray# RealWorld a
              -> Int#
              -> Int#
              -> State# RealWorld
              -> State# RealWorld
deleteAtIndex cmp arr i len s =
  case readSmallArray# arr (len -# 1#) s of
    (# s', a #) -> case writeSmallArray# arr i a s' of
      s'' -> case readSmallArray# arr parent s'' of
        (# s''', b #) -> case cmp b a of
          GT -> adjustUpwards cmp arr i a s'''
          _ -> adjustDownwards cmp arr i (len -# 1#) a s'''
  where
    parent = (i -# 1#) `quotInt#` 2#