program/ArrayQueue.hs

   1 module ArrayQueue (
   2         ArrayQueue, newArrayQueue,
   3         aqEnqueue, aqIsEmpty, aqDequeue
   4 ) where
   5 import Control.Monad.ST
   6 import Data.Array.ST
   7 import MonadStuff
   8
   9 data ArrayQueue s = ArrayQueue {
  10         aqArr :: STUArray s Int Int,
  11         aqHeadI :: Int, -- Also used as an "end of queue" sentinel
  12         aqTailI :: Int  -- Element value can be aqHeadI; also used as a "not queued" sentinel
  13 }
  14
  15 newArrayQueue :: (Int, Int) -> ST s (ArrayQueue s)
  16 newArrayQueue (lo, hi) = do
  17         let headI = lo - 1
  18         let tailI = lo - 2
  19         arr <- newArray (tailI, hi) tailI
  20         writeArray arr headI headI -- queue is empty
  21         writeArray arr tailI headI -- tail is head
  22         return $ ArrayQueue arr headI tailI
  23
  24 aqEnqueue :: ArrayQueue s -> Int -> ST s Bool -- Was it added?
  25 aqEnqueue (ArrayQueue arr headI tailI) newI = do
  26         newIval <- readArray arr newI
  27         if newIval == tailI
  28                 then do
  29                         lst <- readArray arr tailI
  30                         writeArray arr lst newI -- Append newI.
  31                         writeArray arr newI headI -- newI is now the tail.
  32                         writeArray arr tailI newI -- The tail is now at newI.
  33                         return True
  34                 else return False -- Already on queue.
  35
  36 aqIsEmpty :: ArrayQueue s -> ST s Bool
  37 aqIsEmpty (ArrayQueue arr headI tailI) = do
  38         first <- readArray arr headI
  39         return (first == headI)
  40
  41 aqDequeue :: ArrayQueue s -> ST s (Maybe Int)
  42 aqDequeue (ArrayQueue arr headI tailI) = do
  43         first <- readArray arr headI
  44         if first == headI
  45                 then return Nothing
  46                 else do
  47                         next <- readArray arr first
  48                         writeArray arr headI next
  49                         if next == headI
  50                                 then writeArray arr tailI headI -- Emptied queue.
  51                                 else nop
  52                         writeArray arr first tailI -- No longer on queue.
  53                         return $ Just first