{-# LANGUAGE RecordWildCards, BangPatterns, GADTs, UnboxedTuples #-}

-- Note that argument order is more like IORef than Map, because its mutable
module General.Ids(
    Ids, Id,
    empty, insert, lookup,
    null, size, sizeUpperBound,
    forWithKeyM_, for,
    toList, toMap
    ) where

import Data.IORef.Extra
import Data.Primitive.Array
import Control.Exception
import General.Intern(Id(..))
import Control.Monad.Extra
import Data.Maybe
import Data.Functor
import qualified Data.HashMap.Strict as Map
import Prelude hiding (lookup, null)
import GHC.IO(IO(..))
import GHC.Exts(RealWorld)


newtype Ids a = Ids (IORef (S a))

data S a = S
    {capacity :: {-# UNPACK #-} !Int -- ^ Number of entries in values, initially 0
    ,used :: {-# UNPACK #-} !Int -- ^ Capacity that has been used, assuming no gaps from index 0, initially 0
    ,values :: {-# UNPACK #-} !(MutableArray RealWorld (Maybe a))
    }


empty :: IO (Ids a)
empty = do
    let capacity = 0
    let used = 0
    values <- newArray capacity Nothing
    Ids <$> newIORef S{..}


sizeUpperBound :: Ids a -> IO Int
sizeUpperBound (Ids ref) = do
    S{..} <- readIORef ref
    return used


size :: Ids a -> IO Int
size (Ids ref) = do
    S{..} <- readIORef ref
    let go !acc i
            | i < 0 = return acc
            | otherwise = do
                v <- readArray values i
                if isJust v then go (acc+1) (i-1) else go acc (i-1)
    go 0 (used-1)


toMap :: Ids a -> IO (Map.HashMap Id a)
toMap ids = do
    mp <- Map.fromList <$> toListUnsafe ids
    return $! mp

forWithKeyM_ :: Ids a -> (Id -> a -> IO ()) -> IO ()
forWithKeyM_ (Ids ref) f = do
    S{..} <- readIORef ref
    let go !i | i >= used = return ()
              | otherwise = do
                v <- readArray values i
                whenJust v $ f $ Id $ fromIntegral i
                go $ i+1
    go 0

for :: Ids a -> (a -> b) -> IO (Ids b)
for (Ids ref) f = do
    S{..} <- readIORef ref
    values2 <- newArray capacity Nothing
    let go !i | i >= used = return ()
              | otherwise = do
                v <- readArray values i
                whenJust v $ \v -> writeArray values2 i $ Just $ f v
                go $ i+1
    go 0
    Ids <$> newIORef (S capacity used values2)


toListUnsafe :: Ids a -> IO [(Id, a)]
toListUnsafe (Ids ref) = do
    S{..} <- readIORef ref

    -- execute in O(1) stack
    -- see https://neilmitchell.blogspot.co.uk/2015/09/making-sequencemapm-for-io-take-o1-stack.html
    let index r i | i >= used = []
        index r i | IO io <- readArray values i = case io r of
            (# r, Nothing #) -> index r (i+1)
            (# r, Just v  #) -> (Id $ fromIntegral i, v) : index r (i+1)

    IO $ \r -> (# r, index r 0 #)


toList :: Ids a -> IO [(Id, a)]
toList ids = do
    xs <- toListUnsafe ids
    let demand (x:xs) = demand xs
        demand [] = ()
    evaluate $ demand xs
    return xs


null :: Ids a -> IO Bool
null ids = (== 0) <$> sizeUpperBound ids


insert :: Ids a -> Id -> a -> IO ()
insert (Ids ref) (Id i) v = do
    S{..} <- readIORef ref
    let ii = fromIntegral i
    if ii < capacity then do
        writeArray values ii $ Just v
        when (ii >= used) $ writeIORef' ref S{used=ii+1,..}
     else do
        c2 <- return $ max (capacity * 2) (ii + 10000)
        v2 <- newArray c2 Nothing
        copyMutableArray v2 0 values 0 capacity
        writeArray v2 ii $ Just v
        writeIORef' ref $ S c2 (ii+1) v2

lookup :: Ids a -> Id -> IO (Maybe a)
lookup (Ids ref) (Id i) = do
    S{..} <- readIORef ref
    let ii = fromIntegral i
    if ii < used then
        readArray values ii
     else
        return Nothing