{-# 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
    {S a -> Int
capacity :: {-# UNPACK #-} !Int -- ^ Number of entries in values, initially 0
    ,S a -> Int
used :: {-# UNPACK #-} !Int -- ^ Capacity that has been used, assuming no gaps from index 0, initially 0
    ,S a -> MutableArray RealWorld (Maybe a)
values :: {-# UNPACK #-} !(MutableArray RealWorld (Maybe a))
    }


empty :: IO (Ids a)
empty :: IO (Ids a)
empty = do
    let capacity :: Int
capacity = 0
    let used :: Int
used = 0
    MutableArray RealWorld (Maybe a)
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
    IORef (S a) -> Ids a
forall a. IORef (S a) -> Ids a
Ids (IORef (S a) -> Ids a) -> IO (IORef (S a)) -> IO (Ids a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> S a -> IO (IORef (S a))
forall a. a -> IO (IORef a)
newIORef $WS :: forall a. Int -> Int -> MutableArray RealWorld (Maybe a) -> S a
S{..}


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


size :: Ids a -> IO Int
size :: Ids a -> IO Int
size (Ids ref :: 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 -> Int -> IO Int
go !Int
acc i :: Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< 0 = Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
acc
            | Bool
otherwise = do
                Maybe a
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 Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
v then Int -> Int -> IO Int
go (Int
accInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) else Int -> Int -> IO Int
go Int
acc (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
    Int -> Int -> IO Int
go 0 (Int
usedInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)


toMap :: Ids a -> IO (Map.HashMap Id a)
toMap :: Ids a -> IO (HashMap Id a)
toMap ids :: Ids a
ids = do
    HashMap Id a
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
    HashMap Id a -> IO (HashMap Id a)
forall (m :: * -> *) a. Monad m => a -> m a
return (HashMap Id a -> IO (HashMap Id a))
-> HashMap Id a -> IO (HashMap Id a)
forall a b. (a -> b) -> a -> b
$! HashMap Id a
mp

forWithKeyM_ :: Ids a -> (Id -> a -> IO ()) -> IO ()
forWithKeyM_ :: Ids a -> (Id -> a -> IO ()) -> IO ()
forWithKeyM_ (Ids ref :: IORef (S a)
ref) f :: 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 -> IO ()
go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
              | Bool
otherwise = do
                Maybe a
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
                Maybe a -> (a -> IO ()) -> IO ()
forall (m :: * -> *) a.
Applicative m =>
Maybe a -> (a -> m ()) -> m ()
whenJust Maybe a
v ((a -> IO ()) -> IO ()) -> (a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Id -> a -> IO ()
f (Id -> a -> IO ()) -> Id -> a -> IO ()
forall a b. (a -> b) -> a -> b
$ 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
                Int -> IO ()
go (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1
    Int -> IO ()
go 0

for :: Ids a -> (a -> b) -> IO (Ids b)
for :: Ids a -> (a -> b) -> IO (Ids b)
for (Ids ref :: IORef (S a)
ref) f :: a -> b
f = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    MutableArray RealWorld (Maybe b)
values2 <- Int -> Maybe b -> IO (MutableArray (PrimState IO) (Maybe b))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
capacity Maybe b
forall a. Maybe a
Nothing
    let go :: Int -> IO ()
go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
              | Bool
otherwise = do
                Maybe a
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
                Maybe a -> (a -> IO ()) -> IO ()
forall (m :: * -> *) a.
Applicative m =>
Maybe a -> (a -> m ()) -> m ()
whenJust Maybe a
v ((a -> IO ()) -> IO ()) -> (a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \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
                Int -> IO ()
go (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1
    Int -> IO ()
go 0
    IORef (S b) -> Ids b
forall a. IORef (S a) -> Ids a
Ids (IORef (S b) -> Ids b) -> IO (IORef (S b)) -> IO (Ids b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> S b -> IO (IORef (S b))
forall a. a -> IO (IORef a)
newIORef (Int -> Int -> MutableArray RealWorld (Maybe b) -> S b
forall a. Int -> Int -> MutableArray RealWorld (Maybe a) -> S a
S Int
capacity Int
used MutableArray RealWorld (Maybe b)
values2)


toListUnsafe :: Ids a -> IO [(Id, a)]
toListUnsafe :: Ids a -> IO [(Id, a)]
toListUnsafe (Ids ref :: 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 -> [(Id, a)]
index r :: State# RealWorld
r i :: Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used = []
        index r :: State# RealWorld
r i :: Int
i | IO 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
            (# r :: State# RealWorld
r, Nothing #) -> State# RealWorld -> Int -> [(Id, a)]
index State# RealWorld
r (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
            (# r :: State# RealWorld
r, Just v :: 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
+1)

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


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


null :: Ids a -> IO Bool
null :: Ids a -> IO Bool
null ids :: Ids a
ids = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 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 :: Ids a -> Id -> a -> IO ()
insert (Ids ref :: IORef (S a)
ref) (Id i :: Word32
i) v :: a
v = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let ii :: Int
ii = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i
    if Int
ii Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
capacity then do
        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
ii (Maybe a -> IO ()) -> Maybe a -> IO ()
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
v
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ii Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
used) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef (S a) -> S a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef' IORef (S a)
ref $WS :: forall a. Int -> Int -> MutableArray RealWorld (Maybe a) -> S a
S{used :: Int
used=Int
iiInt -> Int -> Int
forall a. Num a => a -> a -> a
+1,..}
     else do
        Int
c2 <- Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
capacity Int -> Int -> Int
forall a. Num a => a -> a -> a
* 2) (Int
ii Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 10000)
        MutableArray RealWorld (Maybe a)
v2 <- Int -> Maybe a -> IO (MutableArray (PrimState IO) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
c2 Maybe a
forall a. Maybe a
Nothing
        MutableArray (PrimState IO) (Maybe a)
-> Int
-> MutableArray (PrimState IO) (Maybe a)
-> Int
-> Int
-> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
copyMutableArray MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
v2 0 MutableArray RealWorld (Maybe a)
MutableArray (PrimState IO) (Maybe a)
values 0 Int
capacity
        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)
v2 Int
ii (Maybe a -> IO ()) -> Maybe a -> IO ()
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
v
        IORef (S a) -> S a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef' IORef (S a)
ref (S a -> IO ()) -> S a -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> MutableArray RealWorld (Maybe a) -> S a
forall a. Int -> Int -> MutableArray RealWorld (Maybe a) -> S a
S Int
c2 (Int
iiInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) MutableArray RealWorld (Maybe a)
v2

lookup :: Ids a -> Id -> IO (Maybe a)
lookup :: Ids a -> Id -> IO (Maybe a)
lookup (Ids ref :: IORef (S a)
ref) (Id i :: Word32
i) = do
    S{..} <- IORef (S a) -> IO (S a)
forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let ii :: Int
ii = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i
    if Int
ii Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
used then
        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
ii
     else
        Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing