{-# 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, fromList,
    null, size, sizeUpperBound,
    forWithKeyM_, forCopy, forMutate,
    toList, elems, toMap
    ) where

import Data.IORef.Extra
import Data.Primitive.Array hiding (fromList)
import Control.Exception
import General.Intern(Id(..))
import Control.Monad.Extra
import Data.List.Extra(zipFrom)
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
    {forall a. S a -> Int
capacity :: {-# UNPACK #-} !Int -- ^ Number of entries in values, initially 0
    ,forall a. S a -> Int
used :: {-# UNPACK #-} !Int -- ^ Capacity that has been used, assuming no gaps from index 0, initially 0
    ,forall a. S a -> MutableArray RealWorld (Maybe a)
values :: {-# UNPACK #-} !(MutableArray RealWorld (Maybe a))
    }


empty :: IO (Ids a)
empty :: forall a. IO (Ids a)
empty = do
    let capacity :: Int
capacity = Int
0
    let used :: Int
used = Int
0
    values <- Int -> Maybe a -> IO (MutableArray (PrimState IO) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
capacity Maybe a
forall a. Maybe a
Nothing
    Ids <$> newIORef S{..}

fromList :: [a] -> IO (Ids a)
fromList :: forall a. [a] -> IO (Ids a)
fromList [a]
xs = do
    let capacity :: Int
capacity = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
    let used :: Int
used = Int
capacity
    values <- Int -> Maybe a -> IO (MutableArray (PrimState IO) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
capacity Maybe a
forall a. Maybe a
Nothing
    forM_ (zipFrom 0 xs) $ \(Int
i, a
x) ->
        MutableArray (PrimState IO) (Maybe a) -> Int -> Maybe a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
i (Maybe a -> IO ()) -> Maybe a -> IO ()
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
x
    Ids <$> newIORef S{..}

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


size :: Ids a -> IO Int
size :: forall a. Ids a -> IO Int
size (Ids IORef (S a)
ref) = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let go !Int
acc Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
acc
            | Bool
otherwise = do
                v <- MutableArray (PrimState IO) (Maybe a) -> Int -> IO (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
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 :: forall a. Ids a -> IO (HashMap Id a)
toMap Ids a
ids = do
    mp <- [(Id, a)] -> HashMap Id a
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Id, a)] -> HashMap Id a) -> IO [(Id, a)] -> IO (HashMap Id a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ids a -> IO [(Id, a)]
forall a. Ids a -> IO [(Id, a)]
toListUnsafe Ids a
ids
    pure $! mp

forWithKeyM_ :: Ids a -> (Id -> a -> IO ()) -> IO ()
forWithKeyM_ :: forall a. Ids a -> (Id -> a -> IO ()) -> IO ()
forWithKeyM_ (Ids IORef (S a)
ref) Id -> a -> IO ()
f = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | Bool
otherwise = do
                v <- MutableArray (PrimState IO) (Maybe a) -> Int -> IO (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
i
                whenJust v $ f $ Id $ fromIntegral i
                go $ i+1
    go 0

forCopy :: Ids a -> (a -> b) -> IO (Ids b)
forCopy :: forall a b. Ids a -> (a -> b) -> IO (Ids b)
forCopy (Ids IORef (S a)
ref) a -> b
f = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    values2 <- newArray capacity Nothing
    let go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | Bool
otherwise = do
                v <- MutableArray (PrimState IO) (Maybe a) -> Int -> IO (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
i
                whenJust v $ \a
v -> MutableArray (PrimState IO) (Maybe b) -> Int -> Maybe b -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld (Maybe b)
MutableArray (PrimState IO) (Maybe b)
values2 Int
i (Maybe b -> IO ()) -> Maybe b -> IO ()
forall a b. (a -> b) -> a -> b
$ b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> b -> Maybe b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
v
                go $ i+1
    go 0
    Ids <$> newIORef (S capacity used values2)


forMutate :: Ids a -> (a -> a) -> IO ()
forMutate :: forall a. Ids a -> (a -> a) -> IO ()
forMutate (Ids IORef (S a)
ref) a -> a
f = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | Bool
otherwise = do
                v <- MutableArray (PrimState IO) (Maybe a) -> Int -> IO (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
i
                whenJust v $ \a
v -> MutableArray (PrimState IO) (Maybe a) -> Int -> Maybe a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
i (Maybe a -> IO ()) -> Maybe a -> IO ()
forall a b. (a -> b) -> a -> b
$! a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! a -> a
f a
v
                go $ i+1
    go 0


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

    -- execute in O(1) stack
    -- see https://neilmitchell.blogspot.co.uk/2015/09/making-sequencemapm-for-io-take-o1-stack.html
    let index State# RealWorld
_ Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = []
        index State# RealWorld
r Int
i | IO State# RealWorld -> (# State# RealWorld, Maybe a #)
io <- MutableArray (PrimState IO) (Maybe a) -> Int -> IO (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values Int
i = case State# RealWorld -> (# State# RealWorld, Maybe a #)
io State# RealWorld
r of
            (# State# RealWorld
r, Maybe a
Nothing #) -> State# RealWorld -> Int -> [(Id, a)]
index State# RealWorld
r (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
            (# State# RealWorld
r, Just a
v  #) -> (Word32 -> Id
Id (Word32 -> Id) -> Word32 -> Id
forall a b. (a -> b) -> a -> b
$ Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i, a
v) (Id, a) -> [(Id, a)] -> [(Id, a)]
forall a. a -> [a] -> [a]
: State# RealWorld -> Int -> [(Id, a)]
index State# RealWorld
r (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

    IO $ \State# RealWorld
r -> (# State# RealWorld
r, State# RealWorld -> Int -> [(Id, a)]
index State# RealWorld
r Int
0 #)


toList :: Ids a -> IO [(Id, a)]
toList :: forall a. Ids a -> IO [(Id, a)]
toList Ids a
ids = do
    xs <- Ids a -> IO [(Id, a)]
forall a. Ids a -> IO [(Id, a)]
toListUnsafe Ids a
ids
    let demand (a
_:[a]
xs) = [a] -> ()
demand [a]
xs
        demand [] = ()
    evaluate $ demand xs
    pure xs

elems :: Ids a -> IO [a]
elems :: forall a. Ids a -> IO [a]
elems Ids a
ids = ((Id, a) -> a) -> [(Id, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Id, a) -> a
forall a b. (a, b) -> b
snd ([(Id, a)] -> [a]) -> IO [(Id, a)] -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ids a -> IO [(Id, a)]
forall a. Ids a -> IO [(Id, a)]
toList Ids a
ids

null :: Ids a -> IO Bool
null :: forall a. Ids a -> IO Bool
null Ids a
ids = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Int -> Bool) -> IO Int -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ids a -> IO Int
forall a. Ids a -> IO Int
sizeUpperBound Ids a
ids


insert :: Ids a -> Id -> a -> IO ()
insert :: forall a. Ids a -> Id -> a -> IO ()
insert (Ids IORef (S a)
ref) (Id Word32
i) a
v = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let ii = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i
    if ii < capacity then do
        writeArray values ii $ Just v
        when (ii >= used) $ writeIORef' ref S{used=ii+1,..}
     else do
        c2<- pure $ 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 :: forall a. Ids a -> Id -> IO (Maybe a)
lookup (Ids IORef (S a)
ref) (Id Word32
i) = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let ii = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i
    if ii < used then
        readArray values ii
     else
        pure Nothing