{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
module ToySolver.Internal.Data.PriorityQueue
(
PriorityQueue
, Index
, newPriorityQueue
, newPriorityQueueBy
, NewFifo (..)
, getElems
, clear
, clone
, Enqueue (..)
, Dequeue (..)
, QueueSize (..)
, rebuild
, getHeapArray
, getHeapVec
, 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
data PriorityQueue a
= PriorityQueue
{ PriorityQueue a -> a -> a -> IO Bool
lt :: !(a -> a -> IO Bool)
, PriorityQueue a -> Vec a
heap :: !(Vec.Vec a)
}
newPriorityQueue :: Ord a => IO (PriorityQueue a)
newPriorityQueue :: 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 (m :: * -> *) a. Monad m => a -> m a
return (a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
b))
newPriorityQueueBy :: (a -> a -> IO Bool) -> IO (PriorityQueue a)
newPriorityQueueBy :: (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 (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 :: forall a. (a -> a -> IO Bool) -> Vec a -> PriorityQueue a
PriorityQueue{ lt :: a -> a -> IO Bool
lt = a -> a -> IO Bool
cmp, heap :: GenericVec IOArray a
heap = GenericVec IOArray a
vec }
getElems :: PriorityQueue a -> IO [a]
getElems :: 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)
clear :: PriorityQueue a -> IO ()
clear :: 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)
clone :: PriorityQueue a -> IO (PriorityQueue a)
clone :: 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 (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 :: forall a. (a -> a -> IO Bool) -> Vec a -> PriorityQueue a
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
Int
n <- GenericVec IOArray a -> IO Int
forall (a :: * -> * -> *) e. GenericVec a e -> IO Int
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 -> Int -> IO ()
forall a. PriorityQueue a -> Int -> IO ()
up PriorityQueue a
q Int
n
instance Dequeue (PriorityQueue a) IO a where
dequeue :: PriorityQueue a -> IO (Maybe a)
dequeue PriorityQueue a
q = do
Int
n <- GenericVec IOArray a -> IO Int
forall (a :: * -> * -> *) e. GenericVec a e -> IO Int
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
case Int
n of
Int
0 ->
Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
Int
_ -> do
a
val <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
0
if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 then do
GenericVec IOArray a -> Int -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO ()
Vec.resize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
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 -> Int -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
0 a
val1
PriorityQueue a -> Int -> IO ()
forall a. PriorityQueue a -> Int -> IO ()
down PriorityQueue a
q Int
0
Maybe a -> IO (Maybe 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 (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 Int
queueSize PriorityQueue a
q = GenericVec IOArray a -> IO Int
forall (a :: * -> * -> *) e. GenericVec a e -> IO Int
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
up :: PriorityQueue a -> Index -> IO ()
up :: PriorityQueue a -> Int -> IO ()
up PriorityQueue a
q !Int
i = do
a
val <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
i
let loop :: Int -> IO Int
loop Int
0 = Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
loop Int
j = do
let p :: Int
p = Int -> Int
parent Int
j
a
val_p <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
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 -> Int -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
j a
val_p
Int -> IO Int
loop Int
p
else Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
j
Int
j <- Int -> IO Int
loop Int
i
GenericVec IOArray a -> Int -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
j a
val
down :: PriorityQueue a -> Index -> IO ()
down :: PriorityQueue a -> Int -> IO ()
down PriorityQueue a
q !Int
i = do
Int
n <- GenericVec IOArray a -> IO Int
forall (a :: * -> * -> *) e. GenericVec a e -> IO Int
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
a
val <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
i
let loop :: Int -> IO Int
loop !Int
j = do
let !l :: Int
l = Int -> Int
left Int
j
!r :: Int
r = Int -> Int
right Int
j
if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n
then Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
j
else do
Int
child <- do
if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n
then Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
else do
a
val_l <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
l
a
val_r <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
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 Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
r
else Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
a
val_child <- GenericVec IOArray a -> Int -> IO a
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO e
Vec.unsafeRead (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
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 Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
j
else do
GenericVec IOArray a -> Int -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
j a
val_child
Int -> IO Int
loop Int
child
Int
j <- Int -> IO Int
loop Int
i
GenericVec IOArray a -> Int -> a -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> e -> IO ()
Vec.unsafeWrite (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
j a
val
rebuild :: PriorityQueue a -> IO ()
rebuild :: PriorityQueue a -> IO ()
rebuild PriorityQueue a
q = do
Int
n <- GenericVec IOArray a -> IO Int
forall (a :: * -> * -> *) e. GenericVec a e -> IO Int
Vec.getSize (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
a -> (a -> Bool) -> (a -> a) -> (a -> m ()) -> m ()
forLoop Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
n) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
PriorityQueue a -> Int -> IO ()
forall a. PriorityQueue a -> Int -> IO ()
up PriorityQueue a
q Int
i
getHeapArray :: PriorityQueue a -> IO (A.IOArray Index a)
getHeapArray :: PriorityQueue a -> IO (IOArray Int a)
getHeapArray PriorityQueue a
q = GenericVec IOArray a -> IO (IOArray Int a)
forall (a :: * -> * -> *) e. GenericVec a e -> IO (a Int e)
Vec.getArray (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
getHeapVec :: PriorityQueue a -> IO (Vec.Vec a)
getHeapVec :: PriorityQueue a -> IO (Vec a)
getHeapVec PriorityQueue a
q = Vec a -> IO (Vec a)
forall (m :: * -> *) a. Monad m => a -> m a
return (PriorityQueue a -> Vec a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q)
resizeHeapCapacity :: PriorityQueue a -> Int -> IO ()
resizeHeapCapacity :: PriorityQueue a -> Int -> IO ()
resizeHeapCapacity PriorityQueue a
q Int
capa = GenericVec IOArray a -> Int -> IO ()
forall (a :: * -> * -> *) e.
MArray a e IO =>
GenericVec a e -> Int -> IO ()
Vec.resizeCapacity (PriorityQueue a -> GenericVec IOArray a
forall a. PriorityQueue a -> Vec a
heap PriorityQueue a
q) Int
capa
{-# INLINE left #-}
left :: Index -> Index
left :: Int -> Int
left Int
i = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
{-# INLINE right #-}
right :: Index -> Index
right :: Int -> Int
right Int
i = (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2;
{-# INLINE parent #-}
parent :: Index -> Index
parent :: Int -> Int
parent Int
i = (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2