content: -- SPDX-FileCopyrightText: Copyright Preetham Gujjula
  -- SPDX-License-Identifier: BSD-3-Clause
  
  module Data.PQueue.Prio.Min.Mutable
    ( -- * Documentation
      MMinPQueue,
  
      -- * Construction
      empty,
      singleton,
      insert,
  
      -- * Query
      deleteMin,
  
      -- * List conversion
      fromAscList,
    )
  where
  
  import Control.Monad (when)
  import Control.Monad.ST (ST)
  import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)
  import qualified Data.Vector as Vector
  import Data.Vector.Mutable (MVector)
  import qualified Data.Vector.Mutable as MVector
  
  data MMinPQueue s k a
    = MMinPQueue
    { vectorRef :: STRef s (MVector s (k, a)),
      sizeRef :: STRef s Int
    }
  
  empty :: ST s (MMinPQueue s k a)
  empty = do
    vec <- MVector.unsafeNew 1
    vecRef <- newSTRef vec
    sizeRef <- newSTRef 0
    pure $
      MMinPQueue
        { vectorRef = vecRef,
          sizeRef = sizeRef
        }
  
  singleton :: k -> a -> ST s (MMinPQueue s k a)
  singleton k v = do
    vec <- MVector.replicate 1 (k, v)
    vecRef <- newSTRef vec
    sizeRef <- newSTRef 1
    pure $
      MMinPQueue
        { vectorRef = vecRef,
          sizeRef = sizeRef
        }
  
  resizeIfNeeded :: MMinPQueue s k a -> ST s ()
  resizeIfNeeded pqueue = do
    vec <- readSTRef (vectorRef pqueue)
    size <- readSTRef (sizeRef pqueue)
    when (MVector.length vec == size) $ do
      vec' <- MVector.unsafeGrow vec size
      writeSTRef (vectorRef pqueue) vec'
  
  insert :: (Ord k) => k -> a -> MMinPQueue s k a -> ST s ()
  insert k v pqueue = do
    resizeIfNeeded pqueue
  
    n <- readSTRef (sizeRef pqueue)
    vec <- readSTRef (vectorRef pqueue)
    MVector.unsafeWrite vec n (k, v)
    writeSTRef (sizeRef pqueue) (n + 1)
  
    bubbleUp vec n
  
  bubbleUp :: (Ord k) => MVector s (k, a) -> Int -> ST s ()
  bubbleUp vector = go
    where
      go i =
        if i == 0
          then pure ()
          else do
            let parentIndex = (i - 1) `quot` 2
            parentVal <- MVector.unsafeRead vector parentIndex
            currentVal <- MVector.unsafeRead vector i
            when (fst parentVal > fst currentVal) $ do
              MVector.unsafeSwap vector parentIndex i
              go parentIndex
  
  deleteMin :: (Ord k) => MMinPQueue s k a -> ST s (Maybe (k, a))
  deleteMin pqueue = do
    n <- readSTRef (sizeRef pqueue)
    if n == 0
      then pure Nothing
      else do
        vector <- readSTRef (vectorRef pqueue)
        minNode <- MVector.unsafeRead vector 0
        MVector.unsafeSwap vector 0 (n - 1)
        writeSTRef (sizeRef pqueue) (n - 1)
        bubbleDown vector 0 (n - 1)
        pure (Just minNode)
  
  bubbleDown :: (Ord k) => MVector s (k, a) -> Int -> Int -> ST s ()
  bubbleDown vector i n = go i
    where
      go j = do
        case compare n (2 * j + 2) of
          LT -> pure ()
          EQ -> do
            x <- MVector.unsafeRead vector j
            x' <- MVector.unsafeRead vector (2 * j + 1)
            let j' = 2 * j + 1
            when (fst x > fst x') $ do
              MVector.unsafeWrite vector j' x
              MVector.unsafeWrite vector j x'
              go j'
          GT -> do
            x <- MVector.unsafeRead vector j
            xleft <- MVector.unsafeRead vector (2 * j + 1)
            xright <- MVector.unsafeRead vector (2 * j + 2)
            let (x', j') =
                  if fst xleft <= fst xright
                    then (xleft, 2 * j + 1)
                    else (xright, 2 * j + 2)
            when (fst x > fst x') $ do
              MVector.unsafeWrite vector j' x
              MVector.unsafeWrite vector j x'
              go j'
  
  fromAscList :: forall s k a. [(k, a)] -> ST s (MMinPQueue s k a)
  fromAscList xs = do
    vector <- Vector.thaw (Vector.fromList xs)
    vectorRef <- newSTRef vector
    sizeRef <- newSTRef (length xs)
    pure $
      MMinPQueue
        { vectorRef = vectorRef,
          sizeRef = sizeRef
        }
  