{-# 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
    MutableArray RealWorld (Maybe a)
values <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
capacity forall a. Maybe a
Nothing
    forall a. IORef (S a) -> Ids a
Ids forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (IORef a)
newIORef S{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
..}

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

sizeUpperBound :: Ids a -> IO Int
sizeUpperBound :: forall a. Ids a -> IO Int
sizeUpperBound (Ids IORef (S a)
ref) = do
    S{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: forall a. S a -> MutableArray RealWorld (Maybe a)
used :: forall a. S a -> Int
capacity :: forall a. S a -> Int
..} <- forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
used


size :: Ids a -> IO Int
size :: forall a. Ids a -> IO Int
size (Ids IORef (S a)
ref) = do
    S{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: forall a. S a -> MutableArray RealWorld (Maybe a)
used :: forall a. S a -> Int
capacity :: forall a. S a -> Int
..} <- forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let go :: Int -> Int -> IO Int
go !Int
acc Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
acc
            | Bool
otherwise = do
                Maybe a
v <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
values Int
i
                if forall a. Maybe a -> Bool
isJust Maybe a
v then Int -> Int -> IO Int
go (Int
accforall a. Num a => a -> a -> a
+Int
1) (Int
iforall a. Num a => a -> a -> a
-Int
1) else Int -> Int -> IO Int
go Int
acc (Int
iforall a. Num a => a -> a -> a
-Int
1)
    Int -> Int -> IO Int
go Int
0 (Int
usedforall a. Num a => a -> a -> a
-Int
1)


toMap :: Ids a -> IO (Map.HashMap Id a)
toMap :: forall a. Ids a -> IO (HashMap Id a)
toMap Ids a
ids = do
    HashMap Id a
mp <- forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Ids a -> IO [(Id, a)]
toListUnsafe Ids a
ids
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! HashMap Id a
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{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: forall a. S a -> MutableArray RealWorld (Maybe a)
used :: forall a. S a -> Int
capacity :: forall a. S a -> Int
..} <- forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let go :: Int -> IO ()
go !Int
i | Int
i forall a. Ord a => a -> a -> Bool
>= Int
used = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | Bool
otherwise = do
                Maybe a
v <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
values Int
i
                forall (m :: * -> *) a.
Applicative m =>
Maybe a -> (a -> m ()) -> m ()
whenJust Maybe a
v forall a b. (a -> b) -> a -> b
$ Id -> a -> IO ()
f forall a b. (a -> b) -> a -> b
$ Word32 -> Id
Id forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                Int -> IO ()
go forall a b. (a -> b) -> a -> b
$ Int
iforall a. Num a => a -> a -> a
+Int
1
    Int -> IO ()
go Int
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{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: forall a. S a -> MutableArray RealWorld (Maybe a)
used :: forall a. S a -> Int
capacity :: forall a. S a -> Int
..} <- forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    MutableArray RealWorld (Maybe b)
values2 <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
capacity forall a. Maybe a
Nothing
    let go :: Int -> IO ()
go !Int
i | Int
i forall a. Ord a => a -> a -> Bool
>= Int
used = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | Bool
otherwise = do
                Maybe a
v <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
values Int
i
                forall (m :: * -> *) a.
Applicative m =>
Maybe a -> (a -> m ()) -> m ()
whenJust Maybe a
v forall a b. (a -> b) -> a -> b
$ \a
v -> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld (Maybe b)
values2 Int
i forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a -> b
f a
v
                Int -> IO ()
go forall a b. (a -> b) -> a -> b
$ Int
iforall a. Num a => a -> a -> a
+Int
1
    Int -> IO ()
go Int
0
    forall a. IORef (S a) -> Ids a
Ids forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (IORef a)
newIORef (forall a. Int -> Int -> MutableArray RealWorld (Maybe a) -> S a
S Int
capacity Int
used MutableArray RealWorld (Maybe b)
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{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: forall a. S a -> MutableArray RealWorld (Maybe a)
used :: forall a. S a -> Int
capacity :: forall a. S a -> Int
..} <- forall a. IORef a -> IO a
readIORef IORef (S a)
ref
    let go :: Int -> IO ()
go !Int
i | Int
i forall a. Ord a => a -> a -> Bool
>= Int
used = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | Bool
otherwise = do
                Maybe a
v <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (Maybe a)
values Int
i
                forall (m :: * -> *) a.
Applicative m =>
Maybe a -> (a -> m ()) -> m ()
whenJust Maybe a
v forall a b. (a -> b) -> a -> b
$ \a
v -> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld (Maybe a)
values Int
i forall a b. (a -> b) -> a -> b
$! forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! a -> a
f a
v
                Int -> IO ()
go forall a b. (a -> b) -> a -> b
$ Int
iforall a. Num a => a -> a -> a
+Int
1
    Int -> IO ()
go Int
0


toListUnsafe :: Ids a -> IO [(Id, a)]
toListUnsafe :: forall a. Ids a -> IO [(Id, a)]
toListUnsafe (Ids IORef (S a)
ref) = do
    S{Int
MutableArray RealWorld (Maybe a)
values :: MutableArray RealWorld (Maybe a)
used :: Int
capacity :: Int
values :: forall a. S a -> MutableArray RealWorld (Maybe a)
used :: forall a. S a -> Int
capacity :: forall a. S a -> Int
..} <- 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 State# RealWorld
_ Int
i | Int
i forall a. Ord a => a -> a -> Bool
>= Int
used = []
        index State# RealWorld
r Int
i | IO State# RealWorld -> (# State# RealWorld, Maybe a #)
io <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld (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
iforall a. Num a => a -> a -> a
+Int
1)
            (# State# RealWorld
r, Just a
v  #) -> (Word32 -> Id
Id forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i, a
v) forall a. a -> [a] -> [a]
: State# RealWorld -> Int -> [(Id, a)]
index State# RealWorld
r (Int
iforall a. Num a => a -> a -> a
+Int
1)

    forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \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
    [(Id, a)]
xs <- forall a. Ids a -> IO [(Id, a)]
toListUnsafe Ids a
ids
    let demand :: [a] -> ()
demand (a
_:[a]
xs) = [a] -> ()
demand [a]
xs
        demand [] = ()
    forall a. a -> IO a
evaluate forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> ()
demand [(Id, a)]
xs
    forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Id, a)]
xs

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