{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving #-}

module Store(

    ElemIndex(),

    MStore(),

    IStore(),

    newStore,

    addElem,

    removeElem,

    readElem,

    writeElem,

    modifyElem,

    elemExists,

    firstIndex,

    indicesM,

    withIStore,

    withIStore2,

    (!),

    indices

    ) where

import qualified Data.Array.IArray as IA

import qualified Data.Array.IO as IOA

import qualified Data.IntSet as IntSet

import Data.IORef

import Control.Monad

import Control.DeepSeq

newtype ElemIndex = ElemIndex Int

    deriving (Eq, Show, Read, Ord, NFData)

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

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

firstIndex :: ElemIndex

firstIndex = ElemIndex 0

-- MStore code

initialSize :: Int

initialSize = 16

growFunc :: Int -> Int

growFunc a = a * 3 `div` 2

truncFunc :: Int -> Int

truncFunc a | a > growFunc initialSize = (a `div` 2)

            | otherwise = a

newStore :: IO (MStore e)

newStore = do

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

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

    return (MStore new)

growStore :: MStore e -> IO ()

growStore (MStore ref) = do

    (busyElems, freeElems, arr) <- readIORef ref

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

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

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

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

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

growIfNeeded :: MStore e -> IO ()

growIfNeeded m@(MStore ref) = do

    (_, freeElems, _) <- readIORef ref

    when (IntSet.null freeElems) $ growStore m

truncateIfNeeded :: MStore e -> IO ()

truncateIfNeeded (MStore ref) = do

    (busyElems, _, arr) <- readIORef ref

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

    let newM' = truncFunc (m' + 1) - 1

    when (newM' < m' && (not $ IntSet.null busyElems) && IntSet.findMax busyElems <= newM') $ do

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

        sequence_ [IOA.readArray arr i >>= IOA.writeArray newArr i | i <- IntSet.toList busyElems]

        writeIORef ref (busyElems, IntSet.fromAscList [0..newM'] `IntSet.difference` busyElems, newArr)

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

addElem m@(MStore ref) element = do

    growIfNeeded m

    (busyElems, freeElems, arr) <- readIORef ref

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

    IOA.writeArray arr n element

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

    return $ ElemIndex n

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

removeElem m@(MStore ref) (ElemIndex n) = do

    (busyElems, freeElems, arr) <- readIORef ref

    IOA.writeArray arr n (error $ "Store: no element " ++ show n)

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

    truncateIfNeeded m

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

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

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

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

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

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

    (_, _, arr) <- readIORef ref

    IOA.readArray arr n >>= IOA.writeArray arr n . f

elemExists :: MStore e -> ElemIndex -> IO Bool

elemExists (MStore ref) (ElemIndex n) = do

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

    return $ n `IntSet.notMember` free

indicesM :: MStore e -> IO [ElemIndex]

indicesM (MStore ref) = do

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

    return $ map ElemIndex $ IntSet.toList busy

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

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

m2i (MStore ref) = do

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

    c <- IOA.unsafeFreeze c'

    return $ IStore (a, c)

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

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

    (b, e, _) <- readIORef ref

    a <- IOA.unsafeThaw arr

    writeIORef ref (b, e, a)

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

withIStore m f = do

    i <- m2i m

    let res = f i

    res `seq` i2m m i

    return res

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

withIStore2 m1 m2 f = do

    i1 <- m2i m1

    i2 <- m2i m2

    let res = f i1 i2

    res `seq` i2m m1 i1

    i2m m2 i2

    return res

-- IStore code

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

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

indices :: IStore e -> [ElemIndex]

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