{-# LANGUAGE TypeFamilies #-}
-- |
-- Module:
--   Data.FastMutableIntMap
-- Description:
--   A mutable version of 'IntMap'
module Data.FastMutableIntMap
  ( FastMutableIntMap
  , new
  , newEmpty
  , insert
  , isEmpty
  , getFrozenAndClear
  , size
  , applyPatch
  , PatchIntMap (..)
  , traverseIntMapPatchWithKey
  , lookup
  , forIntersectionWithImmutable_
  , for_
  , patchIntMapNewElements
  , patchIntMapNewElementsMap
  , getDeletions
  , toList
  ) where

--TODO: Pure JS version
--TODO: Fast copy to FastIntMap
--TODO: Fast patch type

import Prelude hiding (lookup)

import Control.Monad.IO.Class
import Data.Foldable (traverse_)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.IORef
import Data.Patch.Class
import Data.Patch.IntMap

-- | A 'FastMutableIntMap' holds a map of values of type @a@ and allows low-overhead modifications via IO.
-- Operations on 'FastMutableIntMap' run in IO.
newtype FastMutableIntMap a = FastMutableIntMap (IORef (IntMap a))

-- | Create a new 'FastMutableIntMap' out of an 'IntMap'
new :: IntMap a -> IO (FastMutableIntMap a)
new :: IntMap a -> IO (FastMutableIntMap a)
new m :: IntMap a
m = IORef (IntMap a) -> FastMutableIntMap a
forall a. IORef (IntMap a) -> FastMutableIntMap a
FastMutableIntMap (IORef (IntMap a) -> FastMutableIntMap a)
-> IO (IORef (IntMap a)) -> IO (FastMutableIntMap a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap a -> IO (IORef (IntMap a))
forall a. a -> IO (IORef a)
newIORef IntMap a
m

-- | Create a new empty 'FastMutableIntMap'
newEmpty :: IO (FastMutableIntMap a)
newEmpty :: IO (FastMutableIntMap a)
newEmpty = IORef (IntMap a) -> FastMutableIntMap a
forall a. IORef (IntMap a) -> FastMutableIntMap a
FastMutableIntMap (IORef (IntMap a) -> FastMutableIntMap a)
-> IO (IORef (IntMap a)) -> IO (FastMutableIntMap a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap a -> IO (IORef (IntMap a))
forall a. a -> IO (IORef a)
newIORef IntMap a
forall a. IntMap a
IntMap.empty

-- | Insert an element into a 'FastMutableIntMap' at the given key
insert :: FastMutableIntMap a -> Int -> a -> IO ()
insert :: FastMutableIntMap a -> Int -> a -> IO ()
insert (FastMutableIntMap r :: IORef (IntMap a)
r) k :: Int
k v :: a
v = IORef (IntMap a) -> (IntMap a -> IntMap a) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (IntMap a)
r ((IntMap a -> IntMap a) -> IO ())
-> (IntMap a -> IntMap a) -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k a
v

-- | Attempt to lookup an element by key in a 'FastMutableIntMap'
lookup :: FastMutableIntMap a -> Int -> IO (Maybe a)
lookup :: FastMutableIntMap a -> Int -> IO (Maybe a)
lookup (FastMutableIntMap r :: IORef (IntMap a)
r) k :: Int
k = Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k (IntMap a -> Maybe a) -> IO (IntMap a) -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r

-- | Runs the provided action over the intersection of a 'FastMutableIntMap' and an 'IntMap'
forIntersectionWithImmutable_ :: MonadIO m => FastMutableIntMap a -> IntMap b -> (a -> b -> m ()) -> m ()
forIntersectionWithImmutable_ :: FastMutableIntMap a -> IntMap b -> (a -> b -> m ()) -> m ()
forIntersectionWithImmutable_ (FastMutableIntMap r :: IORef (IntMap a)
r) b :: IntMap b
b f :: a -> b -> m ()
f = do
  IntMap a
a <- IO (IntMap a) -> m (IntMap a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IntMap a) -> m (IntMap a)) -> IO (IntMap a) -> m (IntMap a)
forall a b. (a -> b) -> a -> b
$ IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r
  ((a, b) -> m ()) -> IntMap (a, b) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ((a -> b -> m ()) -> (a, b) -> m ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> m ()
f) (IntMap (a, b) -> m ()) -> IntMap (a, b) -> m ()
forall a b. (a -> b) -> a -> b
$ (a -> b -> (a, b)) -> IntMap a -> IntMap b -> IntMap (a, b)
forall a b c. (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
IntMap.intersectionWith (,) IntMap a
a IntMap b
b

-- | Runs the provided action over the values of a 'FastMutableIntMap'
for_ :: MonadIO m => FastMutableIntMap a -> (a -> m ()) -> m ()
for_ :: FastMutableIntMap a -> (a -> m ()) -> m ()
for_ (FastMutableIntMap r :: IORef (IntMap a)
r) f :: a -> m ()
f = do
  IntMap a
a <- IO (IntMap a) -> m (IntMap a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IntMap a) -> m (IntMap a)) -> IO (IntMap a) -> m (IntMap a)
forall a b. (a -> b) -> a -> b
$ IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r
  (a -> m ()) -> IntMap a -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ a -> m ()
f IntMap a
a

-- | Checks whether a 'FastMutableIntMap' is empty
isEmpty :: FastMutableIntMap a -> IO Bool
isEmpty :: FastMutableIntMap a -> IO Bool
isEmpty (FastMutableIntMap r :: IORef (IntMap a)
r) = IntMap a -> Bool
forall a. IntMap a -> Bool
IntMap.null (IntMap a -> Bool) -> IO (IntMap a) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r

-- | Retrieves the size of a 'FastMutableIntMap'
size :: FastMutableIntMap a -> IO Int
size :: FastMutableIntMap a -> IO Int
size (FastMutableIntMap r :: IORef (IntMap a)
r) = IntMap a -> Int
forall a. IntMap a -> Int
IntMap.size (IntMap a -> Int) -> IO (IntMap a) -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r

-- | Make an immutable snapshot of the datastructure and clear it
getFrozenAndClear :: FastMutableIntMap a -> IO (IntMap a)
getFrozenAndClear :: FastMutableIntMap a -> IO (IntMap a)
getFrozenAndClear (FastMutableIntMap r :: IORef (IntMap a)
r) = do
  IntMap a
result <- IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r
  IORef (IntMap a) -> IntMap a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IntMap a)
r IntMap a
forall a. IntMap a
IntMap.empty
  IntMap a -> IO (IntMap a)
forall (m :: * -> *) a. Monad m => a -> m a
return IntMap a
result

-- | Updates the value of a 'FastMutableIntMap' with the given patch (see 'Data.Patch.IntMap'),
-- and returns an 'IntMap' with the modified keys and values.
applyPatch :: FastMutableIntMap a -> PatchIntMap a -> IO (IntMap a)
applyPatch :: FastMutableIntMap a -> PatchIntMap a -> IO (IntMap a)
applyPatch (FastMutableIntMap r :: IORef (IntMap a)
r) p :: PatchIntMap a
p@(PatchIntMap m :: IntMap (Maybe a)
m) = do
  IntMap a
v <- IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r
  IORef (IntMap a) -> IntMap a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IntMap a)
r (IntMap a -> IO ()) -> IntMap a -> IO ()
forall a b. (a -> b) -> a -> b
$! PatchIntMap a
-> PatchTarget (PatchIntMap a) -> PatchTarget (PatchIntMap a)
forall p. Patch p => p -> PatchTarget p -> PatchTarget p
applyAlways PatchIntMap a
p IntMap a
PatchTarget (PatchIntMap a)
v
  IntMap a -> IO (IntMap a)
forall (m :: * -> *) a. Monad m => a -> m a
return (IntMap a -> IO (IntMap a)) -> IntMap a -> IO (IntMap a)
forall a b. (a -> b) -> a -> b
$ IntMap a -> IntMap (Maybe a) -> IntMap a
forall a b. IntMap a -> IntMap b -> IntMap a
IntMap.intersection IntMap a
v IntMap (Maybe a)
m

toList :: FastMutableIntMap a -> IO [(Int, a)]
toList :: FastMutableIntMap a -> IO [(Int, a)]
toList (FastMutableIntMap r :: IORef (IntMap a)
r) = IntMap a -> [(Int, a)]
forall a. IntMap a -> [(Int, a)]
IntMap.toList (IntMap a -> [(Int, a)]) -> IO (IntMap a) -> IO [(Int, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
r