{-# 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
(
PriorityQueue
, Value
, Index
, newPriorityQueue
, newPriorityQueueBy
, NewFifo (..)
, getElems
, clear
, clone
, Enqueue (..)
, Dequeue (..)
, QueueSize (..)
, member
, update
, rebuild
, getHeapArray
, getHeapVec
, 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
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)
}
newPriorityQueue :: IO PriorityQueue
newPriorityQueue :: IO PriorityQueue
newPriorityQueue = (Value -> Value -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy (\Value
a Value
b -> Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Value
a Value -> Value -> Bool
forall a. Ord a => a -> a -> Bool
< Value
b))
#ifdef UNBOXED_COMPARISON_ARGUMENTS
{-# INLINE newPriorityQueueBy #-}
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)
newPriorityQueueBy# :: (Int# -> Int# -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy# :: (Int# -> Int# -> IO Bool) -> IO PriorityQueue
newPriorityQueueBy# Int# -> Int# -> IO Bool
cmp# = do
UVec Value
vec <- IO (UVec Value)
forall (a :: * -> * -> *) e. MArray a e IO => IO (GenericVec a e)
Vec.new
UVec Value
idx <- IO (UVec Value)
forall (a :: * -> * -> *) e. MArray a e IO => IO (GenericVec a e)
Vec.new
PriorityQueue -> IO PriorityQueue
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue -> IO PriorityQueue)
-> PriorityQueue -> IO PriorityQueue
forall a b. (a -> b) -> a -> b
$ PriorityQueue :: (Int# -> Int# -> IO Bool)
-> UVec Value -> UVec Value -> PriorityQueue
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
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
getElems :: PriorityQueue -> IO [Value]
getElems :: PriorityQueue -> IO [Value]
getElems PriorityQueue
q = UVec Value -> IO [Value]
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO [e]
Vec.getElems (PriorityQueue -> UVec Value
heap PriorityQueue
q)
clear :: PriorityQueue -> IO ()
clear :: PriorityQueue -> IO ()
clear PriorityQueue
q = do
UVec Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO ()
Vec.clear (PriorityQueue -> UVec Value
heap PriorityQueue
q)
UVec Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO ()
Vec.clear (PriorityQueue -> UVec Value
table PriorityQueue
q)
clone :: PriorityQueue -> IO PriorityQueue
clone :: PriorityQueue -> IO PriorityQueue
clone PriorityQueue
q = do
UVec Value
h2 <- UVec Value -> IO (UVec Value)
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 <- UVec Value -> IO (UVec Value)
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO (GenericVec a e)
Vec.clone (PriorityQueue -> UVec Value
table PriorityQueue
q)
PriorityQueue -> IO PriorityQueue
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue -> IO PriorityQueue)
-> PriorityQueue -> IO PriorityQueue
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
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
m (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Value
n <- UVec Value -> IO Value
forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
heap PriorityQueue
q)
UVec Value -> Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> e -> IO ()
Vec.push (PriorityQueue -> UVec Value
heap PriorityQueue
q) Value
val
UVec Value -> Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO ()
Vec.growTo (PriorityQueue -> UVec Value
table PriorityQueue
q) (Value
valValue -> Value -> Value
forall a. Num a => a -> a -> a
+Value
1)
UVec Value -> Value -> Value -> IO ()
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 <- UVec Value -> IO Value
forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
heap PriorityQueue
q)
case Value
n of
Value
0 ->
Maybe Value -> IO (Maybe Value)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Value
forall a. Maybe a
Nothing
Value
_ -> do
Value
val <- UVec Value -> Value -> IO Value
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO e
Vec.unsafeRead (PriorityQueue -> UVec Value
heap PriorityQueue
q) Value
0
UVec Value -> Value -> Value -> IO ()
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 Value -> Value -> Bool
forall a. Eq a => a -> a -> Bool
== Value
1 then do
UVec Value -> Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO ()
Vec.resize (PriorityQueue -> UVec Value
heap PriorityQueue
q) (Value
nValue -> Value -> Value
forall a. Num a => a -> a -> a
-Value
1)
else do
Value
val1 <- UVec Value -> IO Value
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> IO e
Vec.unsafePop (PriorityQueue -> UVec Value
heap PriorityQueue
q)
UVec Value -> Value -> Value -> IO ()
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
UVec Value -> Value -> Value -> IO ()
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
Maybe Value -> IO (Maybe Value)
forall (m :: * -> *) a. Monad m => a -> m a
return (Value -> Maybe Value
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 <- PriorityQueue -> IO (Maybe Value)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
dequeue PriorityQueue
q
case Maybe Value
r of
Maybe Value
Nothing -> [Value] -> IO [Value]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
xs)
Just Value
x -> [Value] -> IO [Value]
go (Value
xValue -> [Value] -> [Value]
forall a. a -> [a] -> [a]
:[Value]
xs)
instance QueueSize PriorityQueue IO where
queueSize :: PriorityQueue -> IO Value
queueSize PriorityQueue
q = UVec Value -> IO Value
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 <- UVec Value -> IO Value
forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
table PriorityQueue
q)
if Value
n Value -> Value -> Bool
forall a. Ord a => a -> a -> Bool
<= Value
v then
Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
else do
Value
i <- UVec Value -> Value -> IO Value
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO e
Vec.unsafeRead (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
v
Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$! Value
i Value -> Value -> Bool
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 <- UVec Value -> Value -> IO Value
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO e
Vec.unsafeRead (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
v
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Value
i Value -> Value -> Bool
forall a. Eq a => a -> a -> Bool
== -Value
1) (IO () -> IO ()) -> IO () -> IO ()
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 <- UVec Value -> Value -> IO Value
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 = Value -> IO Value
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 <- UVec Value -> Value -> IO Value
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
UVec Value -> Value -> Value -> IO ()
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
UVec Value -> Value -> Value -> IO ()
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 Value -> IO Value
forall (m :: * -> *) a. Monad m => a -> m a
return Value
j
Value
j <- Value -> IO Value
loop Value
i
UVec Value -> Value -> Value -> IO ()
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
UVec Value -> Value -> Value -> IO ()
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 <- UVec Value -> IO Value
forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
heap PriorityQueue
q)
Value
val <- UVec Value -> Value -> IO Value
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 Value -> Value -> Bool
forall a. Ord a => a -> a -> Bool
>= Value
n
then Value -> IO Value
forall (m :: * -> *) a. Monad m => a -> m a
return Value
j
else do
Value
child <- do
if Value
r Value -> Value -> Bool
forall a. Ord a => a -> a -> Bool
>= Value
n
then Value -> IO Value
forall (m :: * -> *) a. Monad m => a -> m a
return Value
l
else do
Value
val_l <- UVec Value -> Value -> IO Value
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 <- UVec Value -> Value -> IO Value
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 Value -> IO Value
forall (m :: * -> *) a. Monad m => a -> m a
return Value
r
else Value -> IO Value
forall (m :: * -> *) a. Monad m => a -> m a
return Value
l
Value
val_child <- UVec Value -> Value -> IO Value
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 Value -> IO Value
forall (m :: * -> *) a. Monad m => a -> m a
return Value
j
else do
UVec Value -> Value -> Value -> IO ()
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
UVec Value -> Value -> Value -> IO ()
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
UVec Value -> Value -> Value -> IO ()
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
UVec Value -> Value -> Value -> IO ()
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 <- UVec Value -> IO Value
forall (a :: * -> * -> *) e. GenericVec a e -> IO Value
Vec.getSize (PriorityQueue -> UVec Value
heap PriorityQueue
q)
Value
-> (Value -> Bool) -> (Value -> Value) -> (Value -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
a -> (a -> Bool) -> (a -> a) -> (a -> m ()) -> m ()
forLoop Value
0 (Value -> Value -> Bool
forall a. Ord a => a -> a -> Bool
<Value
n) (Value -> Value -> Value
forall a. Num a => a -> a -> a
+Value
1) ((Value -> IO ()) -> IO ()) -> (Value -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Value
i -> do
PriorityQueue -> Value -> IO ()
up PriorityQueue
q Value
i
getHeapArray :: PriorityQueue -> IO (A.IOUArray Index Value)
getHeapArray :: PriorityQueue -> IO (IOUArray Value Value)
getHeapArray PriorityQueue
q = UVec Value -> IO (IOUArray Value Value)
forall (a :: * -> * -> *) e. GenericVec a e -> IO (a Value e)
Vec.getArray (PriorityQueue -> UVec Value
heap PriorityQueue
q)
getHeapVec :: PriorityQueue -> IO (Vec.UVec Value)
getHeapVec :: PriorityQueue -> IO (UVec Value)
getHeapVec PriorityQueue
q = UVec Value -> IO (UVec Value)
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue -> UVec Value
heap PriorityQueue
q)
resizeHeapCapacity :: PriorityQueue -> Int -> IO ()
resizeHeapCapacity :: PriorityQueue -> Value -> IO ()
resizeHeapCapacity PriorityQueue
q Value
capa = UVec Value -> Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO ()
Vec.resizeCapacity (PriorityQueue -> UVec Value
heap PriorityQueue
q) Value
capa
resizeTableCapacity :: PriorityQueue -> Int -> IO ()
resizeTableCapacity :: PriorityQueue -> Value -> IO ()
resizeTableCapacity PriorityQueue
q Value
capa = UVec Value -> Value -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Value -> IO ()
Vec.resizeCapacity (PriorityQueue -> UVec Value
table PriorityQueue
q) Value
capa
{-# INLINE left #-}
left :: Index -> Index
left :: Value -> Value
left Value
i = Value
iValue -> Value -> Value
forall a. Num a => a -> a -> a
*Value
2 Value -> Value -> Value
forall a. Num a => a -> a -> a
+ Value
1
{-# INLINE right #-}
right :: Index -> Index
right :: Value -> Value
right Value
i = (Value
iValue -> Value -> Value
forall a. Num a => a -> a -> a
+Value
1)Value -> Value -> Value
forall a. Num a => a -> a -> a
*Value
2;
{-# INLINE parent #-}
parent :: Index -> Index
parent :: Value -> Value
parent Value
i = (Value
iValue -> Value -> Value
forall a. Num a => a -> a -> a
-Value
1) Value -> Value -> Value
forall a. Integral a => a -> a -> a
`div` Value
2