    {-# 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