{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeSynonymInstances #-}
#ifdef __GLASGOW_HASKELL__
#define UNBOXED_COMPARISON_ARGUMENTS
#endif
#ifdef UNBOXED_COMPARISON_ARGUMENTS
{-# LANGUAGE MagicHash #-}
#endif
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Internal.Data.IndexedPriorityQueue
-- 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.IndexedPriorityQueue
  (
  -- * PriorityQueue type
    PriorityQueue
  , Value
  , Index

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

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

  -- * Misc operations
  , resizeHeapCapacity
  , resizeTableCapacity
  ) where

import Control.Loop
import Control.Monad
import qualified Data.Array.IO as A
import Data.Queue.Classes
import qualified ToySolver.Internal.Data.Vec as Vec
#ifdef UNBOXED_COMPARISON_ARGUMENTS
import GHC.Exts
#endif

type Index = Int
type Value = Int

-- | Priority queue implemented as array-based binary heap.
data PriorityQueue
  = PriorityQueue
#ifdef UNBOXED_COMPARISON_ARGUMENTS
  { PriorityQueue -> Int# -> Int# -> IO Bool
lt#  :: !(Int# -> Int# -> IO Bool)
#else
  { lt   :: !(Value -> Value -> IO Bool)
#endif
  , PriorityQueue -> UVec Value
heap :: !(Vec.UVec Value)
  , PriorityQueue -> UVec Value
table  :: !(Vec.UVec Index)
  }

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

#ifdef UNBOXED_COMPARISON_ARGUMENTS

{-# INLINE newPriorityQueueBy #-}
-- | Build a priority queue with a given /less than/ operator.
newPriorityQueueBy :: (Value -> Value -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy :: (Value -> Value -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy Value -> Value -> IO Bool
cmp = (Int# -> Int# -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy# Int# -> Int# -> IO Bool
cmp#
  where
    cmp# :: Int# -> Int# -> IO Bool
cmp# Int#
a Int#
b = Value -> Value -> IO Bool
cmp (Int# -> Value
I# Int#
a) (Int# -> Value
I# Int#
b)

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

{-# INLINE lt #-}
lt :: PriorityQueue -> Value -> Value -> IO Bool
lt :: PriorityQueue -> Value -> Value -> IO Bool
lt PriorityQueue
q (I# Int#
a) (I# Int#
b) = PriorityQueue -> Int# -> Int# -> IO Bool
lt# PriorityQueue
q Int#
a Int#
b

#else

-- | Build a priority queue with a given /less than/ operator.
newPriorityQueueBy :: (Value -> Value -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy cmp = do
  vec <- Vec.new
  idx <- Vec.new
  return $ PriorityQueue{ lt = cmp, heap = vec, table = idx }

#endif

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

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

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

instance NewFifo PriorityQueue IO where
  newFifo :: IO PriorityQueue
newFifo = IO PriorityQueue
newPriorityQueue

instance Enqueue PriorityQueue IO Value where
  enqueue :: PriorityQueue -> Value -> IO ()
enqueue PriorityQueue
q Value
val = do
    Bool
m <- PriorityQueue -> Value -> IO Bool
member PriorityQueue
q Value
val
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
m forall a b. (a -> b) -> a -> b
$ do
      Value
n <- forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
heap PriorityQueue
q)
      forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> e -> IO ()
Vec.push (PriorityQueue -> UVec Value
heap PriorityQueue
q) Value
val
      forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO ()
Vec.growTo (PriorityQueue -> UVec Value
table PriorityQueue
q) (Value
valforall a. Num a => a -> a -> a
+Value
1)
      forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> e -> IO ()
Vec.unsafeWrite (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
val Value
n
      PriorityQueue -> Value -> IO ()
up PriorityQueue
q Value
n

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

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

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

member :: PriorityQueue -> Value -> IO Bool
member :: PriorityQueue -> Value -> IO Bool
member PriorityQueue
q Value
v = do
  Value
n <- forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
table PriorityQueue
q)
  if Value
n forall a. Ord a => a -> a -> Bool
<= Value
v then
    forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  else do
    Value
i <- forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO e
Vec.unsafeRead (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
v
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Value
i forall a. Eq a => a -> a -> Bool
/= -Value
1

update :: PriorityQueue -> Value -> IO ()
update :: PriorityQueue -> Value -> IO ()
update PriorityQueue
q Value
v = do
  Value
i <- forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO e
Vec.unsafeRead (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
v
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Value
i forall a. Eq a => a -> a -> Bool
== -Value
1) forall a b. (a -> b) -> a -> b
$ do
    PriorityQueue -> Value -> IO ()
up PriorityQueue
q Value
i
    PriorityQueue -> Value -> IO ()
down PriorityQueue
q Value
i

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

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

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

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

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

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

-- | Pre-allocate internal buffer for @[0..n-1]@ values.
resizeTableCapacity :: PriorityQueue -> Int -> IO ()
resizeTableCapacity :: PriorityQueue -> Value -> IO ()
resizeTableCapacity PriorityQueue
q Value
capa = forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO ()
Vec.resizeCapacity (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
capa

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

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

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

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

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

{-
checkHeapProperty :: String -> PriorityQueue -> 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

  idx <- readIORef (table q)
  forM_ [0..n-1] $ \i -> do
    v <- A.readArray arr i
    i' <- A.readArray idx v
    when (i /= i') $ error $ str ++ ": invalid index " ++ show (i,v,i')
-}