aivika-6.0.0: Simulation/Aivika/PriorityQueue/EventQueue.hs
{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module : Simulation.Aivika.PriorityQueue
-- Copyright : Copyright (c) 2009-2017, David Sorokin <david.sorokin@gmail.com>
-- License : BSD3
-- Maintainer : David Sorokin <david.sorokin@gmail.com>
-- Stability : experimental
-- Tested with: GHC 8.0.1
--
-- An imperative heap-based priority queue to implement the event queue.
--
module Simulation.Aivika.PriorityQueue.EventQueue
(PriorityQueue,
Priority,
queueNull,
queueCount,
newQueue,
enqueue,
dequeue,
queueFront) where
import Data.Array
import Data.Array.MArray.Safe
import Data.Array.IO.Safe
import Data.IORef
import Data.Maybe
import Control.Monad
-- | The priority value (greater is higher).
type Priority = Int
-- | The 'PriorityQueue' type represents an imperative heap-based
-- priority queue.
data PriorityQueue a =
PriorityQueue { pqKeys :: IORef (IOUArray Int Double),
pqPris :: IORef (IOUArray Int Priority),
pqVals :: IORef (IOArray Int a),
pqSize :: IORef Int }
increase :: PriorityQueue a -> Int -> IO ()
increase pq capacity =
do let keyRef = pqKeys pq
priRef = pqPris pq
valRef = pqVals pq
keys <- readIORef keyRef
pris <- readIORef priRef
vals <- readIORef valRef
(il, iu) <- getBounds keys
let len = (iu - il) + 1
capacity' | len < 64 = max capacity ((len + 1) * 2)
| otherwise = max capacity ((len `div` 2) * 3)
il' = il
iu' = il + capacity' - 1
keys' <- newArray_ (il', iu')
pris' <- newArray_ (il', iu')
vals' <- newArray_ (il', iu')
mapM_ (\i -> do { k <- readArray keys i; writeArray keys' i k }) [il..iu]
mapM_ (\i -> do { p <- readArray pris i; writeArray pris' i p }) [il..iu]
mapM_ (\i -> do { v <- readArray vals i; writeArray vals' i v }) [il..iu]
writeIORef keyRef keys'
writeIORef priRef pris'
writeIORef valRef vals'
siftUp :: IOUArray Int Double
-> IOUArray Int Priority
-> IOArray Int a
-> Int
-> Double
-> Priority
-> a
-> IO ()
siftUp keys pris vals i k p v = loop i
where loop i =
if i == 0
then do writeArray keys i k
writeArray pris i p
writeArray vals i v
else do let n = (i - 1) `div` 2
kn <- readArray keys n
pn <- readArray pris n
if gte k p kn pn -- (k, -p) >= (kn, -pn)
then do writeArray keys i k
writeArray pris i p
writeArray vals i v
else do pn <- readArray pris n
vn <- readArray vals n
writeArray keys i kn
writeArray pris i pn
writeArray vals i vn
loop n
siftDown :: IOUArray Int Double
-> IOUArray Int Priority
-> IOArray Int a
-> Int
-> Int
-> Double
-> Priority
-> a
-> IO ()
siftDown keys pris vals size i k p v = loop i
where loop i =
if i >= (size `div` 2)
then do writeArray keys i k
writeArray pris i p
writeArray vals i v
else do let n = 2 * i + 1
n' = n + 1
kn <- readArray keys n
pn <- readArray pris n
if n' >= size
then if lte k p kn pn -- (k, -p) <= (kn, -pn)
then do writeArray keys i k
writeArray pris i p
writeArray vals i v
else do pn <- readArray pris n
vn <- readArray vals n
writeArray keys i kn
writeArray pris i pn
writeArray vals i vn
loop n
else do kn' <- readArray keys n'
pn' <- readArray pris n'
-- (kn, -pn) > (kn', -pn')
let n'' = if gt kn pn kn' pn' then n' else n
kn'' = if n'' == n' then kn' else kn
pn'' = if n'' == n' then pn' else pn
if lte k p kn'' pn'' -- (k, -p) <= (kn'', -pn'')
then do writeArray keys i k
writeArray pris i p
writeArray vals i v
else do pn'' <- readArray pris n''
vn'' <- readArray vals n''
writeArray keys i kn''
writeArray pris i pn''
writeArray vals i vn''
loop n''
-- | Test whether the priority queue is empty.
queueNull :: PriorityQueue a -> IO Bool
queueNull pq =
do size <- readIORef (pqSize pq)
return $ size == 0
-- | Return the number of elements in the priority queue.
queueCount :: PriorityQueue a -> IO Int
queueCount pq = readIORef (pqSize pq)
-- | Create a new priority queue.
newQueue :: IO (PriorityQueue a)
newQueue =
do keys <- newArray_ (0, 10)
pris <- newArray_ (0, 10)
vals <- newArray_ (0, 10)
keyRef <- newIORef keys
priRef <- newIORef pris
valRef <- newIORef vals
sizeRef <- newIORef 0
return PriorityQueue { pqKeys = keyRef,
pqPris = priRef,
pqVals = valRef,
pqSize = sizeRef }
-- | Enqueue a new element with the specified priority.
enqueue :: PriorityQueue a -> Double -> Priority -> a -> IO ()
enqueue pq k p v =
do i <- readIORef (pqSize pq)
keys <- readIORef (pqKeys pq)
(il, iu) <- getBounds keys
when (i >= iu - il) $ increase pq (i + 2) -- plus one element on the end
writeIORef (pqSize pq) (i + 1)
keys <- readIORef (pqKeys pq) -- it can be another! (side-effect)
pris <- readIORef (pqPris pq)
vals <- readIORef (pqVals pq)
siftUp keys pris vals i k p v
-- | Dequeue the element with the minimal priority.
dequeue :: PriorityQueue a -> IO ()
dequeue pq =
do size <- readIORef (pqSize pq)
when (size == 0) $ error "Empty priority queue: dequeue"
let i = size - 1
writeIORef (pqSize pq) i
keys <- readIORef (pqKeys pq)
pris <- readIORef (pqPris pq)
vals <- readIORef (pqVals pq)
k <- readArray keys i
p <- readArray pris i
v <- readArray vals i
let k0 = 0.0
p0 = 0
v0 = undefined
-- k0 <- readArray keys size
-- p0 <- readArray pris size
-- v0 <- readArray vals size
writeArray keys i k0
writeArray pris i p0
writeArray vals i v0
when (i > 0) $
siftDown keys pris vals i 0 k p v
-- | Return the element with the minimal priority.
queueFront :: PriorityQueue a -> IO (Double, Priority, a)
queueFront pq =
do size <- readIORef (pqSize pq)
when (size == 0) $ error "Empty priority queue: front"
keys <- readIORef (pqKeys pq)
pris <- readIORef (pqPris pq)
vals <- readIORef (pqVals pq)
k <- readArray keys 0
p <- readArray pris 0
v <- readArray vals 0
return (k, p, v)
-- | Whether the first pair is greater than the second one.
gt :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE gt #-}
gt k1 p1 k2 p2 = (k1 > k2) || (k1 == k2 && p1 < p2)
-- | Whether the first pair is greater than or equal to the second one.
gte :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE gte #-}
gte k1 p1 k2 p2 = (k1 > k2) || (k1 == k2 && p1 <= p2)
-- | Whether the first pair is less than the second one.
lt :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE lt #-}
lt k1 p1 k2 p2 = gt k2 p2 k1 p1
-- | Whether the first pair is less than or equal to the second one.
lte :: Double -> Priority -> Double -> Priority -> Bool
{-# INLINE lte #-}
lte k1 p1 k2 p2 = gte k2 p2 k1 p1