1 module Store(

     2     ElemIndex(),

     3     MStore(),

     4     IStore(),

     5     newStore,

     6     addElem,

     7     removeElem,

     8     readElem,

     9     writeElem,

    10     modifyElem,

    11     firstIndex,

    12     withIStore,

    13     withIStore2,

    14     (!),

    15     indices

    16     ) where

    17 

    18 import qualified Data.Array.IArray as IA

    19 import qualified Data.Array.IO as IOA

    20 import qualified Data.IntSet as IntSet

    21 import Data.IORef

    22 import Control.Monad

    23 

    24 

    25 newtype ElemIndex = ElemIndex Int

    26     deriving (Eq)

    27 newtype MStore e = MStore (IORef (IntSet.IntSet, IntSet.IntSet, IOA.IOArray Int e))

    28 newtype IStore e = IStore (IntSet.IntSet, IA.Array Int e)

    29 

    30 instance Show ElemIndex where

    31     show (ElemIndex i) = 'i' : show i

    32 

    33 firstIndex :: ElemIndex

    34 firstIndex = ElemIndex 0

    35 

    36 -- MStore code

    37 initialSize :: Int

    38 initialSize = 10

    39 

    40 

    41 growFunc :: Int -> Int

    42 growFunc a = a * 3 `div` 2

    43 

    44 

    45 newStore :: IO (MStore e)

    46 newStore = do

    47     newar <- IOA.newArray_ (0, initialSize - 1)

    48     new <- newIORef (IntSet.empty, IntSet.fromAscList [0..initialSize - 1], newar)

    49     return (MStore new)

    50 

    51 

    52 growStore :: MStore e -> IO ()

    53 growStore (MStore ref) = do

    54     (busyElems, freeElems, arr) <- readIORef ref

    55     (_, m') <- IOA.getBounds arr

    56     let newM' = growFunc (m' + 1) - 1

    57     newArr <- IOA.newArray_ (0, newM')

    58     sequence_ [IOA.readArray arr i >>= IOA.writeArray newArr i | i <- [0..m']]

    59     writeIORef ref (busyElems, freeElems `IntSet.union` (IntSet.fromAscList [m'+1..newM']), newArr)

    60 

    61 

    62 growIfNeeded :: MStore e -> IO ()

    63 growIfNeeded m@(MStore ref) = do

    64     (_, freeElems, _) <- readIORef ref

    65     when (IntSet.null freeElems) $ growStore m

    66 

    67 

    68 addElem :: MStore e -> e -> IO ElemIndex

    69 addElem m@(MStore ref) element = do

    70     growIfNeeded m

    71     (busyElems, freeElems, arr) <- readIORef ref

    72     let (n, freeElems') = IntSet.deleteFindMin freeElems

    73     IOA.writeArray arr n element

    74     writeIORef ref (IntSet.insert n busyElems, freeElems', arr)

    75     return $ ElemIndex n

    76 

    77 

    78 removeElem :: MStore e -> ElemIndex -> IO ()

    79 removeElem (MStore ref) (ElemIndex n) = do

    80     (busyElems, freeElems, arr) <- readIORef ref

    81     IOA.writeArray arr n undefined

    82     writeIORef ref (IntSet.delete n busyElems, IntSet.insert n freeElems, arr)

    83 

    84 

    85 readElem :: MStore e -> ElemIndex -> IO e

    86 readElem (MStore ref) (ElemIndex n) = readIORef ref >>= \(_, _, arr) -> IOA.readArray arr n

    87 

    88 

    89 writeElem :: MStore e -> ElemIndex -> e -> IO ()

    90 writeElem (MStore ref) (ElemIndex n) el = readIORef ref >>= \(_, _, arr) -> IOA.writeArray arr n el

    91 

    92 

    93 modifyElem :: MStore e -> (e -> e) -> ElemIndex -> IO ()

    94 modifyElem (MStore ref) f (ElemIndex n) = do

    95     (_, _, arr) <- readIORef ref

    96     IOA.readArray arr n >>= (IOA.writeArray arr n) . f

    97 

    98 

    99 -- A way to use see MStore elements in pure code via IStore

   100 m2i :: MStore e -> IO (IStore e)

   101 m2i (MStore ref) = do

   102     (a, _, c') <- readIORef ref 

   103     c <- IOA.unsafeFreeze c'

   104     return $ IStore (a, c)

   105 

   106 withIStore :: MStore e -> (IStore e -> a) -> IO a

   107 withIStore m f = liftM f (m2i m)

   108 

   109 

   110 withIStore2 :: MStore e1 -> MStore e2 -> (IStore e1 -> IStore e2 -> a) -> IO a

   111 withIStore2 m1 m2 f = do

   112     i1 <- m2i m1

   113     i2 <- m2i m2

   114     return $ f i1 i2

   115 

   116 

   117 -- IStore code

   118 (!) :: IStore e -> ElemIndex -> e

   119 (!) (IStore (_, arr)) (ElemIndex i) = (IA.!) arr i

   120 

   121 indices :: IStore e -> [ElemIndex]

   122 indices (IStore (busy, _)) = map ElemIndex $ IntSet.toList busy