DynArray.hs

module Util.DynArray (DynArray,new,newMatrix,newLinearArray,read
                     ,write,writes
                     ,resizeTo,updateIO,update,bounds,index,indexBy
                     ,elements,showDynArray
                     ,deleteInLinearArray
                     ,deleteColumnsInMatrix,deleteRowsInMatrix)
    where

import Prelude hiding (init,read)
import Control.Applicative ((<$>))
import Control.Monad (forM_,when,foldM)
import Control.Exception (assert)
import Data.Array.IO (IOArray,readArray,writeArray,getBounds,getElems)
import qualified Data.Array.IO as ArrayIO
import Data.Ix (Ix,range,inRange)
import Data.IORef
import Data.List (sort)

data DynArray i e = DynArray { ref :: IORef (IOArray i e)
                             , grow :: (i,i) -> i -> (i,i)
                             , init :: e
                             }

defaultGrowingSize = 10

showDynArray :: (Ix i,Show e) => DynArray i e -> IO String
showDynArray dynArray = show <$> (readIORef (ref dynArray) >>= getElems)

new :: Ix i => (i,i) -> ((i,i) -> i -> (i,i)) -> e -> IO (DynArray i e)
new ix grow e = do
  array <- ArrayIO.newArray ix e
  ref <- newIORef array
  return $ DynArray ref grow e

newMatrix :: (Int,Int) -> e -> IO (DynArray (Int,Int) e)
newMatrix (initRows,initColumns) = 
    let grow (_,(r,c)) (i,j) = ( (0,0)
                               , ( max r $ i + defaultGrowingSize - 1
                                 , max c $ j + defaultGrowingSize - 1))
    in
      new ((0,0),(initRows-1,initColumns-1)) grow

newLinearArray :: Int -> e -> IO (DynArray Int e)
newLinearArray numElements =
    let grow _ i = (0,i + defaultGrowingSize - 1)
    in
      new (0,numElements-1) grow

read :: Ix i => i -> DynArray i e -> IO e
read i dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  if inRange bounds i then readArray array i
                      else return $ init dynArray

write :: Ix i => i -> e -> DynArray i e -> IO ()
write i e dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  if inRange bounds i 
   then writeArray array i e
   else do
     array' <- ArrayIO.newArray (grow dynArray bounds i) $ init dynArray
     copyUnsafe array array'
     writeArray array' i e
     writeIORef (ref dynArray) array'

writes :: (Ix i,Enum i) => i -> [e] -> DynArray i e -> IO ()
writes i es dynArray = 
    forM_ (zip [i..] es) $ \(pos,e) -> write pos e dynArray

copyUnsafe :: Ix i => IOArray i e -> IOArray i e -> IO ()
copyUnsafe from to = do
  r <- getBounds from
  forM_ (range r) $ \i -> readArray from i >>= writeArray to i

resizeTo :: Ix i => (i,i) -> DynArray i e -> IO ()
resizeTo range dynArray = 
    ArrayIO.newArray range (init dynArray) >>= writeIORef (ref dynArray)

updateIO :: Ix i => i -> (e -> IO e) -> DynArray i e -> IO ()
updateIO i f dynArray = 
    read i dynArray >>= f >>= \e -> write i e dynArray

update ::  Ix i => i -> (e -> e) -> DynArray i e -> IO ()
update i f = updateIO i (return . f)
  
bounds :: Ix i => DynArray i e -> IO (i,i)
bounds dynArray = readIORef (ref dynArray) >>= getBounds

indexBy :: (Enum i, Ix i) => (e -> Bool) -> DynArray i e -> IO (Maybe i)
indexBy equals dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  let indexRec i = 
          if inRange bounds i 
          then do value <- readArray array i
                  if equals value then return $ Just i
                                  else indexRec $ succ i
          else return Nothing
  indexRec $ fst bounds

index :: (Enum i, Ix i, Eq e) => e -> DynArray i e -> IO (Maybe i)
index e = indexBy ((==) e)

elements :: Ix i => DynArray i e -> IO [e]
elements dynArray = readIORef (ref dynArray) >>= getElems

deleteInLinearArray :: [Int] -> DynArray Int e -> IO ()
deleteInLinearArray is dynArray = do
  assert (is == sort is) $ return ()
  array <- readIORef $ ref dynArray
  (minB,maxB) <- getBounds array

  let deleteElement (x:xs,deleted) index | x == index = 
          return (xs,deleted+1)

      deleteElement (xs,deleted) index = do
        when (deleted > 0) $ 
          readArray array index >>= writeArray array (index - deleted)
        return (xs,deleted)

  (_,deleted) <- foldM deleteElement (is,0) [minB..maxB] 
  forM_ [maxB-deleted+1..maxB] $ \x -> writeArray array x $ init dynArray

deleteColumnsInMatrix :: [Int] -> DynArray (Int,Int) e -> IO ()
deleteColumnsInMatrix columns dynArray = do
  assert (columns == sort columns) $ return ()
  array <- readIORef $ ref dynArray
  ((minRow,minCol),(maxRow,maxCol)) <- getBounds array
                                       
  let copyColumn from to =
          forM_ [minRow..maxRow] $ \r -> 
              readArray array (r,from) >>= writeArray array (r,to)

      deleteColumn (col:cols,deleted) index | col == index = 
          return (cols,deleted+1)

      deleteColumn (cols,deleted) index = do
        when (deleted > 0) $ copyColumn index $ index - deleted
        return (cols,deleted)

      resetColumnToInit col = 
          forM_ [minRow..maxRow] $ \r -> 
              writeArray array (r,col) $ init dynArray

  (_,deleted) <- foldM deleteColumn (columns,0) [minCol..maxCol] 
  forM_ [maxCol-deleted+1..maxCol] resetColumnToInit

deleteRowsInMatrix :: [Int] -> DynArray (Int,Int) e -> IO ()
deleteRowsInMatrix rowsToDelete dynArray = do
  assert (rowsToDelete == sort rowsToDelete) $ return ()
  array <- readIORef $ ref dynArray
  ((minRow,minCol),(maxRow,maxCol)) <- getBounds array
                                       
  let copyRow from to =
          forM_ [minCol..maxCol] $ \c -> 
              readArray array (from,c) >>= writeArray array (to,c)

      deleteRow (row:rows,deleted) index | row == index = 
          return (rows,deleted+1)

      deleteRow (rows,deleted) index = do
        when (deleted > 0) $ copyRow index $ index - deleted
        return (rows,deleted)

      resetRowToInit row = 
          forM_ [minCol..maxCol] $ \c -> 
              writeArray array (row,c) $ init dynArray

  (_,deleted) <- foldM deleteRow (rowsToDelete,0) [minRow..maxRow] 
  forM_ [maxRow-deleted+1..maxRow] resetRowToInit