Rewrite Bellman-Ford and min-cost flow, especially to stop the latter from crashing.

[match/match.git] / program / ArrayQueue.hs

diff --git a/program/ArrayQueue.hs b/program/ArrayQueue.hs
new file mode 100644 (file)
index 0000000..252b5fb
--- /dev/null
+++ b/program/ArrayQueue.hs
@@ -0,0 +1,53 @@
+module ArrayQueue (
+       ArrayQueue, newArrayQueue,
+       aqEnqueue, aqIsEmpty, aqDequeue
+) where
+import Control.Monad.ST
+import Data.Array.ST
+import MonadStuff
+
+data ArrayQueue s = ArrayQueue {
+       aqArr :: STUArray s Int Int,
+       aqHeadI :: Int, -- Also used as an "end of queue" sentinel
+       aqTailI :: Int  -- Element value can be aqHeadI; also used as a "not queued" sentinel
+}
+
+newArrayQueue :: (Int, Int) -> ST s (ArrayQueue s)
+newArrayQueue (lo, hi) = do
+       let headI = lo - 1
+       let tailI = lo - 2
+       arr <- newArray (tailI, hi) tailI
+       writeArray arr headI headI -- queue is empty
+       writeArray arr tailI headI -- tail is head
+       return $ ArrayQueue arr headI tailI
+
+aqEnqueue :: ArrayQueue s -> Int -> ST s Bool -- Was it added?
+aqEnqueue (ArrayQueue arr headI tailI) newI = do
+       newIval <- readArray arr newI
+       if newIval == tailI
+               then do
+                       lst <- readArray arr tailI
+                       writeArray arr lst newI -- Append newI.
+                       writeArray arr newI headI -- newI is now the tail.
+                       writeArray arr tailI newI -- The tail is now at newI.
+                       return True
+               else return False -- Already on queue.
+
+aqIsEmpty :: ArrayQueue s -> ST s Bool
+aqIsEmpty (ArrayQueue arr headI tailI) = do
+       first <- readArray arr headI
+       return (first == headI)
+
+aqDequeue :: ArrayQueue s -> ST s (Maybe Int)
+aqDequeue (ArrayQueue arr headI tailI) = do
+       first <- readArray arr headI
+       if first == headI
+               then return Nothing
+               else do
+                       next <- readArray arr first
+                       writeArray arr headI next
+                       if next == headI
+                               then writeArray arr tailI headI -- Emptied queue.
+                               else nop
+                       writeArray arr first tailI -- No longer on queue.
+                       return $ Just first