1 module Store(

     2     ElemIndex(),

     3     MStore(),

     4     IStore(),

     5     newStore,

     6     addElem,

     7     removeElem,

     8     readElem,

     9     writeElem,

    10     modifyElem,

    11     elemExists,

    12     firstIndex,

    13     indicesM,

    14     withIStore,

    15     withIStore2,

    16     (!),

    17     indices

    18     ) where

    19 

    20 import qualified Data.Array.IArray as IA

    21 import qualified Data.Array.IO as IOA

    22 import qualified Data.IntSet as IntSet

    23 import Data.IORef

    24 import Control.Monad

    25 

    26 

    27 newtype ElemIndex = ElemIndex Int

    28     deriving (Eq, Show, Read, Ord)

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

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

    31 

    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 (error $ "Store: no element " ++ show n)

    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 elemExists :: MStore e -> ElemIndex -> IO Bool

    99 elemExists (MStore ref) (ElemIndex n) = do

   100     (_, free, _) <- readIORef ref

   101     return $ n `IntSet.notMember` free

   102 

   103 indicesM :: MStore e -> IO [ElemIndex]

   104 indicesM (MStore ref) = do

   105     (busy, _, _) <- readIORef ref

   106     return $ map ElemIndex $ IntSet.toList busy

   107 

   108 

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

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

   111 m2i (MStore ref) = do

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

   113     c <- IOA.unsafeFreeze c'

   114     return $ IStore (a, c)

   115 

   116 i2m :: (MStore e) -> IStore e -> IO ()

   117 i2m (MStore ref) (IStore (_, arr)) = do

   118     (b, e, _) <- readIORef ref

   119     a <- IOA.unsafeThaw arr

   120     writeIORef ref (b, e, a)

   121 

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

   123 withIStore m f = do

   124     i <- m2i m

   125     let res = f i

   126     res `seq` i2m m i

   127     return res

   128 

   129 

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

   131 withIStore2 m1 m2 f = do

   132     i1 <- m2i m1

   133     i2 <- m2i m2

   134     let res = f i1 i2

   135     res `seq` i2m m1 i1

   136     i2m m2 i2

   137     return res

   138 

   139 

   140 -- IStore code

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

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

   143 

   144 indices :: IStore e -> [ElemIndex]

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