{-# LANGUAGE RecordWildCards, BangPatterns, GADTs, UnboxedTuples #-}
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
,S a -> Int
used :: {-# UNPACK #-} !Int
,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
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