{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Internal.Data.PriorityQueue
-- Copyright   :  (c) Masahiro Sakai 2012
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- Priority queue implemented as array-based binary heap.
--
-----------------------------------------------------------------------------
module ToySolver.Internal.Data.PriorityQueue
  (
  -- * PriorityQueue type
    PriorityQueue
  , Index

  -- * Constructors
  , newPriorityQueue
  , newPriorityQueueBy
  , NewFifo (..)

  -- * Operators
  , getElems
  , clear
  , clone
  , Enqueue (..)
  , Dequeue (..)
  , QueueSize (..)
  , rebuild
  , getHeapArray
  , getHeapVec

  -- * Misc operations
  , resizeHeapCapacity
  ) where

import Control.Loop
import qualified Data.Array.IO as A
import Data.Queue.Classes
import qualified ToySolver.Internal.Data.Vec as Vec

type Index = Int

-- | Priority queue implemented as array-based binary heap.
data PriorityQueue a
  = PriorityQueue
  { forall a. PriorityQueue a -> a -> a -> IO Bool
lt   :: !(a -> a -> IO Bool)
  , forall a. PriorityQueue a -> Vec a
heap :: !(Vec.Vec a)
  }

-- | Build a priority queue with default ordering ('(<)' of 'Ord' class)
newPriorityQueue :: Ord a => IO (PriorityQueue a)
newPriorityQueue :: forall a. Ord a => IO (PriorityQueue a)
newPriorityQueue = (a -> a -> IO Bool) -> IO (PriorityQueue a)
forall a. (a -> a -> IO Bool) -> IO (PriorityQueue a)
newPriorityQueueBy (\a
a a
b -> Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
b))

-- | Build a priority queue with a given /less than/ operator.
newPriorityQueueBy :: (a -> a -> IO Bool) -> IO (PriorityQueue a)
newPriorityQueueBy :: forall a. (a -> a -> IO Bool) -> IO (PriorityQueue a)
newPriorityQueueBy a -> a -> IO Bool
cmp = do
  GenericVec IOArray a
vec <- IO (GenericVec IOArray a)
forall (a :: * -> * -> *) e. MArray a e IO => IO (GenericVec a e)
Vec.new
  PriorityQueue a -> IO (PriorityQueue a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue a -> IO (PriorityQueue a))
-> PriorityQueue a -> IO (PriorityQueue a)
forall a b. (a -> b) -> a -> b
$ PriorityQueue{ lt :: a -> a -> IO Bool
lt = a -> a -> IO Bool
cmp, heap :: GenericVec IOArray a
heap = GenericVec IOArray a
vec }

-- | Return a list of all the elements of a priority queue. (not sorted)
getElems :: PriorityQueue a -> IO [a]
getElems :: forall a. PriorityQueue a -> IO [a]
getElems PriorityQueue a
q = GenericVec IOArray a -> IO [a]
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO [e]
Vec.getElems (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)

-- | Remove all elements from a priority queue.
clear :: PriorityQueue a -> IO ()
clear :: forall a. PriorityQueue a -> IO ()
clear PriorityQueue a
q = GenericVec IOArray a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO ()
Vec.clear (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)

-- | Create a copy of a priority queue.
clone :: PriorityQueue a -> IO (PriorityQueue a)
clone :: forall a. PriorityQueue a -> IO (PriorityQueue a)
clone PriorityQueue a
q = do
  GenericVec IOArray a
h2 <- GenericVec IOArray a -> IO (GenericVec IOArray a)
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO (GenericVec a e)
Vec.clone (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
  PriorityQueue a -> IO (PriorityQueue a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue a -> IO (PriorityQueue a))
-> PriorityQueue a -> IO (PriorityQueue a)
forall a b. (a -> b) -> a -> b
$ PriorityQueue{ lt :: a -> a -> IO Bool
lt = PriorityQueue a -> a -> a -> IO Bool
forall a. PriorityQueue a -> a -> a -> IO Bool
lt PriorityQueue a
q, heap :: GenericVec IOArray a
heap = GenericVec IOArray a
h2 }

instance Ord a => NewFifo (PriorityQueue a) IO where
  newFifo :: IO (PriorityQueue a)
newFifo = IO (PriorityQueue a)
forall a. Ord a => IO (PriorityQueue a)
newPriorityQueue

instance Enqueue (PriorityQueue a) IO a where
  enqueue :: PriorityQueue a -> a -> IO ()
enqueue PriorityQueue a
q a
val = do
    Index
n <- GenericVec IOArray a -> IO Index
forall (a :: * -> * -> *) e. GenericVec a e -> IO Index
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
    GenericVec IOArray a -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> e -> IO ()
Vec.push (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) a
val
    PriorityQueue a -> Index -> IO ()
forall a. PriorityQueue a -> Index -> IO ()
up PriorityQueue a
q Index
n

instance Dequeue (PriorityQueue a) IO a where
  dequeue :: PriorityQueue a -> IO (Maybe a)
dequeue PriorityQueue a
q = do
    Index
n <- GenericVec IOArray a -> IO Index
forall (a :: * -> * -> *) e. GenericVec a e -> IO Index
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
    case Index
n of
      Index
0 ->
        Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
      Index
_ -> do
        a
val <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
0
        if Index
n Index -> Index -> Bool
forall a. Eq a => a -> a -> Bool
== Index
1 then do
          GenericVec IOArray a -> Index -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO ()
Vec.resize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) (Index
nIndex -> Index -> Index
forall a. Num a => a -> a -> a
-Index
1)
        else do
          a
val1 <- GenericVec IOArray a -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO e
Vec.unsafePop (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
          GenericVec IOArray a -> Index -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
0 a
val1
          PriorityQueue a -> Index -> IO ()
forall a. PriorityQueue a -> Index -> IO ()
down PriorityQueue a
q Index
0
        Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
val)

  dequeueBatch :: PriorityQueue a -> IO [a]
dequeueBatch PriorityQueue a
q = [a] -> IO [a]
go []
    where
      go :: [a] -> IO [a]
      go :: [a] -> IO [a]
go [a]
xs = do
        Maybe a
r <- PriorityQueue a -> IO (Maybe a)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
dequeue PriorityQueue a
q
        case Maybe a
r of
          Maybe a
Nothing -> [a] -> IO [a]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
xs)
          Just a
x -> [a] -> IO [a]
go (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)

instance QueueSize (PriorityQueue a) IO where
  queueSize :: PriorityQueue a -> IO Index
queueSize PriorityQueue a
q = GenericVec IOArray a -> IO Index
forall (a :: * -> * -> *) e. GenericVec a e -> IO Index
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)

up :: PriorityQueue a -> Index -> IO ()
up :: forall a. PriorityQueue a -> Index -> IO ()
up PriorityQueue a
q !Index
i = do
  a
val <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
i
  let loop :: Index -> IO Index
loop Index
0 = Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
0
      loop Index
j = do
        let p :: Index
p = Index -> Index
parent Index
j
        a
val_p <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
p
        Bool
b <- PriorityQueue a -> a -> a -> IO Bool
forall a. PriorityQueue a -> a -> a -> IO Bool
lt PriorityQueue a
q a
val a
val_p
        if Bool
b
          then do
            GenericVec IOArray a -> Index -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
j a
val_p
            Index -> IO Index
loop Index
p
          else Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
j
  Index
j <- Index -> IO Index
loop Index
i
  GenericVec IOArray a -> Index -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
j a
val

down :: PriorityQueue a -> Index -> IO ()
down :: forall a. PriorityQueue a -> Index -> IO ()
down PriorityQueue a
q !Index
i = do
  Index
n <- GenericVec IOArray a -> IO Index
forall (a :: * -> * -> *) e. GenericVec a e -> IO Index
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
  a
val <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
i
  let loop :: Index -> IO Index
loop !Index
j = do
        let !l :: Index
l = Index -> Index
left Index
j
            !r :: Index
r = Index -> Index
right Index
j
        if Index
l Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
>= Index
n
         then Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
j
         else do
           Index
child <- do
             if Index
r Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
>= Index
n
              then Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
l
              else do
                a
val_l <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
l
                a
val_r <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
r
                Bool
b <- PriorityQueue a -> a -> a -> IO Bool
forall a. PriorityQueue a -> a -> a -> IO Bool
lt PriorityQueue a
q a
val_r a
val_l
                if Bool
b
                  then Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
r
                  else Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
l
           a
val_child <- GenericVec IOArray a -> Index -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
child
           Bool
b <- PriorityQueue a -> a -> a -> IO Bool
forall a. PriorityQueue a -> a -> a -> IO Bool
lt PriorityQueue a
q a
val_child a
val
           if Bool -> Bool
not Bool
b
             then Index -> IO Index
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Index
j
             else do
               GenericVec IOArray a -> Index -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
j a
val_child
               Index -> IO Index
loop Index
child
  Index
j <- Index -> IO Index
loop Index
i
  GenericVec IOArray a -> Index -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
j a
val

rebuild :: PriorityQueue a -> IO ()
rebuild :: forall a. PriorityQueue a -> IO ()
rebuild PriorityQueue a
q = do
  Index
n <- GenericVec IOArray a -> IO Index
forall (a :: * -> * -> *) e. GenericVec a e -> IO Index
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
  Index
-> (Index -> Bool) -> (Index -> Index) -> (Index -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
a -> (a -> Bool) -> (a -> a) -> (a -> m ()) -> m ()
forLoop Index
0 (Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
<Index
n) (Index -> Index -> Index
forall a. Num a => a -> a -> a
+Index
1) ((Index -> IO ()) -> IO ()) -> (Index -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Index
i -> do
    PriorityQueue a -> Index -> IO ()
forall a. PriorityQueue a -> Index -> IO ()
up PriorityQueue a
q Index
i

-- | Get the internal representation of a given priority queue.
getHeapArray :: PriorityQueue a -> IO (A.IOArray Index a)
getHeapArray :: forall a. PriorityQueue a -> IO (IOArray Index a)
getHeapArray PriorityQueue a
q = GenericVec IOArray a -> IO (IOArray Index a)
forall (a :: * -> * -> *) e. GenericVec a e -> IO (a Index e)
Vec.getArray (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)

-- | Get the internal representation of a given priority queue.
getHeapVec :: PriorityQueue a -> IO (Vec.Vec a)
getHeapVec :: forall a. PriorityQueue a -> IO (Vec a)
getHeapVec PriorityQueue a
q = Vec a -> IO (Vec a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue a -> Vec a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)

-- | Pre-allocate internal buffer for @n@ elements.
resizeHeapCapacity :: PriorityQueue a -> Int -> IO ()
resizeHeapCapacity :: forall a. PriorityQueue a -> Index -> IO ()
resizeHeapCapacity PriorityQueue a
q Index
capa = GenericVec IOArray a -> Index -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Index -> IO ()
Vec.resizeCapacity (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Index
capa

{--------------------------------------------------------------------
  Index "traversal" functions
--------------------------------------------------------------------}

{-# INLINE left #-}
left :: Index -> Index
left :: Index -> Index
left Index
i = Index
iIndex -> Index -> Index
forall a. Num a => a -> a -> a
*Index
2 Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Index
1

{-# INLINE right #-}
right :: Index -> Index
right :: Index -> Index
right Index
i = (Index
iIndex -> Index -> Index
forall a. Num a => a -> a -> a
+Index
1)Index -> Index -> Index
forall a. Num a => a -> a -> a
*Index
2;

{-# INLINE parent #-}
parent :: Index -> Index
parent :: Index -> Index
parent Index
i = (Index
iIndex -> Index -> Index
forall a. Num a => a -> a -> a
-Index
1) Index -> Index -> Index
forall a. Integral a => a -> a -> a
`div` Index
2

{--------------------------------------------------------------------
  test
--------------------------------------------------------------------}

{-
checkHeapProperty :: String -> PriorityQueue a -> IO ()
checkHeapProperty str q = do
  (n,arr) <- readIORef (heap q)
  let go i = do
        val <- A.readArray arr i
        forM_ [left i, right i] $ \j ->
          when (j < n) $ do
            val2 <- A.readArray arr j
            b <- lt q val2 val
            when b $ do
              error (str ++ ": invalid heap " ++ show j)
            go j
  when (n > 0) $ go 0
-}