{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

-- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\).
--
-- ==== __Example__
-- Create an `IntMap` with capacity \(10\):
--
-- >>> import AtCoder.Extra.IntMap qualified as IM
-- >>> im <- IM.new @_ @Int 10
--
-- `insert`, `delete`, `lookup` and other functions are available:
--
-- >>> IM.insert im 0 100
-- >>> IM.insert im 9 101
-- >>> IM.delete im 0
-- True
--
-- >>> IM.size im
-- 1
--
-- >>> IM.lookup im 9
-- Just 101
--
-- >>> IM.lookup im 1
-- Nothing
--
-- >>> IM.lookupGT im 5
-- Just (9,101)
--
-- @since 1.1.0.0
module AtCoder.Extra.IntMap
  ( -- * IntMap
    IntMap,

    -- * Constructors
    new,
    build,

    -- * Metadata
    capacity,
    size,
    null,

    -- * Lookups
    lookup,
    member,
    notMember,

    -- ** Compartive lookups
    lookupGE,
    lookupGT,
    lookupLE,
    lookupLT,

    -- ** Max/Min lookups
    lookupMin,
    lookupMax,

    -- * Modifications

    -- ** Insertions
    insert,
    insertWith,

    -- ** Updates
    modify,
    modifyM,

    -- ** Deletions
    delete,
    delete_,
    deleteMin,
    deleteMax,

    -- * Conversions
    keys,
    elems,
    assocs,
  )
where

import AtCoder.Extra.IntSet qualified as IS
import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Maybe (fromJust)
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)
import Prelude hiding (lookup, null)

-- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\).
--
-- @since 1.1.0.0
data IntMap s a = IntMap
  { forall s a. IntMap s a -> IntSet s
setIM :: !(IS.IntSet s),
    forall s a. IntMap s a -> MVector s a
valIM :: !(VUM.MVector s a)
  }

-- | \(O(n)\) Creates an `IntMap` for an interval \([0, n)\).
--
-- @since 1.1.0.0
{-# INLINEABLE new #-}
new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntMap (PrimState m) a)
new :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (IntMap (PrimState m) a)
new Int
cap = ST (PrimState m) (IntMap (PrimState m) a)
-> m (IntMap (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (IntMap (PrimState m) a)
 -> m (IntMap (PrimState m) a))
-> ST (PrimState m) (IntMap (PrimState m) a)
-> m (IntMap (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int -> ST (PrimState m) (IntMap (PrimState m) a)
forall a s. Unbox a => Int -> ST s (IntMap s a)
newST Int
cap

-- | \(O(n + m \log n)\) Creates an `IntMap` for an interval \([0, n)\) with initial values.
--
-- @since 1.1.0.0
{-# INLINEABLE build #-}
build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (IntMap (PrimState m) a)
build :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> Vector (Int, a) -> m (IntMap (PrimState m) a)
build Int
cap Vector (Int, a)
xs = ST (PrimState m) (IntMap (PrimState m) a)
-> m (IntMap (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (IntMap (PrimState m) a)
 -> m (IntMap (PrimState m) a))
-> ST (PrimState m) (IntMap (PrimState m) a)
-> m (IntMap (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int -> Vector (Int, a) -> ST (PrimState m) (IntMap (PrimState m) a)
forall a s. Unbox a => Int -> Vector (Int, a) -> ST s (IntMap s a)
buildST Int
cap Vector (Int, a)
xs

-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the map.
--
-- @since 1.1.0.0
{-# INLINEABLE capacity #-}
capacity :: IntMap s a -> Int
capacity :: forall s a. IntMap s a -> Int
capacity = IntSet s -> Int
forall s. IntSet s -> Int
IS.capacity (IntSet s -> Int) -> (IntMap s a -> IntSet s) -> IntMap s a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap s a -> IntSet s
forall s a. IntMap s a -> IntSet s
setIM

-- | \(O(1)\) Returns the number of elements in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE size #-}
size :: (PrimMonad m) => IntMap (PrimState m) a -> m Int
size :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> m Int
size = IntSet (PrimState m) -> m Int
forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Int
IS.size (IntSet (PrimState m) -> m Int)
-> (IntMap (PrimState m) a -> IntSet (PrimState m))
-> IntMap (PrimState m) a
-> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> IntSet (PrimState m)
forall s a. IntMap s a -> IntSet s
setIM

-- | \(O(1)\) Returns whether the map is empty.
--
-- @since 1.1.0.0
{-# INLINEABLE null #-}
null :: (PrimMonad m) => IntMap (PrimState m) a -> m Bool
null :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> m Bool
null = IntSet (PrimState m) -> m Bool
forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Bool
IS.null (IntSet (PrimState m) -> m Bool)
-> (IntMap (PrimState m) a -> IntSet (PrimState m))
-> IntMap (PrimState m) a
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> IntSet (PrimState m)
forall s a. IntMap s a -> IntSet s
setIM

-- | \(O(\log n)\) Looks up the value for a key.
--
-- @since 1.1.0.0
{-# INLINEABLE lookup #-}
lookup :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe a)
lookup :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> Int -> m (Maybe a)
lookup IntMap (PrimState m) a
im Int
k = ST (PrimState m) (Maybe a) -> m (Maybe a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe a) -> m (Maybe a))
-> ST (PrimState m) (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> Int -> ST (PrimState m) (Maybe a)
forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe a)
lookupST IntMap (PrimState m) a
im Int
k

-- | \(O(\log n)\) Tests whether a key \(k\) is in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE member #-}
member :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool
member :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> Int -> m Bool
member = IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
IS.member (IntSet (PrimState m) -> Int -> m Bool)
-> (IntMap (PrimState m) a -> IntSet (PrimState m))
-> IntMap (PrimState m) a
-> Int
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> IntSet (PrimState m)
forall s a. IntMap s a -> IntSet s
setIM

-- | \(O(\log n)\) Tests whether a key \(k\) is not in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE notMember #-}
notMember :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool
notMember :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> Int -> m Bool
notMember = IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
IS.notMember (IntSet (PrimState m) -> Int -> m Bool)
-> (IntMap (PrimState m) a -> IntSet (PrimState m))
-> IntMap (PrimState m) a
-> Int
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> IntSet (PrimState m)
forall s a. IntMap s a -> IntSet s
setIM

-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest key \(k\) such that \(k \ge k_0\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupGE #-}
lookupGE :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupGE :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupGE IntMap (PrimState m) a
im Int
k = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> Int -> ST (PrimState m) (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST IntMap (PrimState m) a
im Int
k

-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest \(k\) such that \(k \gt k_0\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupGT #-}
lookupGT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupGT :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupGT IntMap (PrimState m) a
is Int
k = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> Int -> ST (PrimState m) (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST IntMap (PrimState m) a
is (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \le k_0\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupLE #-}
lookupLE :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupLE :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupLE IntMap (PrimState m) a
im Int
k = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> Int -> ST (PrimState m) (Maybe (Int, a))
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupLEST IntMap (PrimState m) a
im Int
k

-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \lt k_0\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupLT #-}
lookupLT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupLT :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))
lookupLT IntMap (PrimState m) a
im Int
k = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> Int -> ST (PrimState m) (Maybe (Int, a))
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupLEST IntMap (PrimState m) a
im (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the minimum key \(k\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupMin #-}
lookupMin :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))
lookupMin :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> m (Maybe (Int, a))
lookupMin IntMap (PrimState m) a
im = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> ST (PrimState m) (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> ST s (Maybe (Int, a))
lookupMinST IntMap (PrimState m) a
im

-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the maximum key \(k\).
--
-- @since 1.1.0.0
{-# INLINEABLE lookupMax #-}
lookupMax :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))
lookupMax :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> m (Maybe (Int, a))
lookupMax IntMap (PrimState m) a
im = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> ST (PrimState m) (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> ST s (Maybe (Int, a))
lookupMaxST IntMap (PrimState m) a
im

-- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already
-- exists, it is overwritten.
--
-- @since 1.1.0.0
{-# INLINEABLE insert #-}
insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> a -> m ()
insert :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> Int -> a -> m ()
insert IntMap (PrimState m) a
im Int
k a
v = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> Int -> a -> ST (PrimState m) ()
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> a -> ST s ()
insertST IntMap (PrimState m) a
im Int
k a
v

-- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already
-- exists, it overwritten with \(f(v_{\mathrm{new}}, v_{\mathrm{old}})\).
--
-- @since 1.1.0.0
{-# INLINEABLE insertWith #-}
insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()
insertWith :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()
insertWith IntMap (PrimState m) a
im a -> a -> a
f Int
k a
v = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a
-> (a -> a -> a) -> Int -> a -> ST (PrimState m) ()
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> (a -> a -> a) -> Int -> a -> ST s ()
insertWithST IntMap (PrimState m) a
im a -> a -> a
f Int
k a
v

-- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already
-- does not exist, nothing is performed.
--
-- @since 1.1.0.0
{-# INLINEABLE modify #-}
modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a) -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> (a -> a) -> Int -> m ()
modify IntMap (PrimState m) a
im a -> a
f Int
k = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> (a -> a) -> Int -> ST (PrimState m) ()
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> (a -> a) -> Int -> ST s ()
modifyST IntMap (PrimState m) a
im a -> a
f Int
k

-- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already
-- does not exist, nothing is performed.
--
-- @since 1.1.0.0
{-# INLINEABLE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM IntMap {MVector (PrimState m) a
IntSet (PrimState m)
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet (PrimState m)
valIM :: MVector (PrimState m) a
..} a -> m a
f Int
k = do
  Bool
b <- IntSet (PrimState m) -> Int -> m Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
IS.member IntSet (PrimState m)
setIM Int
k
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector (PrimState m) a
valIM a -> m a
f Int
k

-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no
-- such key exists. Returns whether the key existed.
--
-- @since 1.1.0.0
{-# INLINEABLE delete #-}
delete :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool
delete :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> Int -> m Bool
delete IntMap (PrimState m) a
im = ST (PrimState m) Bool -> m Bool
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Bool -> m Bool)
-> (Int -> ST (PrimState m) Bool) -> Int -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> Int -> ST (PrimState m) Bool
forall s a. IntMap s a -> Int -> ST s Bool
deleteST IntMap (PrimState m) a
im

-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no
-- such key exists.
--
-- @since 1.1.0.0
{-# INLINEABLE delete_ #-}
delete_ :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m ()
delete_ :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> Int -> m ()
delete_ IntMap (PrimState m) a
im = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ())
-> (Int -> ST (PrimState m) ()) -> Int -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> Int -> ST (PrimState m) ()
forall s a. IntMap s a -> Int -> ST s ()
deleteST_ IntMap (PrimState m) a
im

-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the minimum key \(k\) in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE deleteMin #-}
deleteMin :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))
deleteMin :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> m (Maybe (Int, a))
deleteMin IntMap (PrimState m) a
is = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> ST (PrimState m) (Maybe (Int, a))
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> ST s (Maybe (Int, a))
deleteMinST IntMap (PrimState m) a
is

-- | \(O(\log n)\) Deletes the \((k, v)\) pair with maximum key \(k\) in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE deleteMax #-}
deleteMax :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))
deleteMax :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> m (Maybe (Int, a))
deleteMax IntMap (PrimState m) a
is = ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a)))
-> ST (PrimState m) (Maybe (Int, a)) -> m (Maybe (Int, a))
forall a b. (a -> b) -> a -> b
$ IntMap (PrimState m) a -> ST (PrimState m) (Maybe (Int, a))
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> ST s (Maybe (Int, a))
deleteMaxST IntMap (PrimState m) a
is

-- | \(O(n \log n)\) Enumerates the keys in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE keys #-}
keys :: (PrimMonad m) => IntMap (PrimState m) a -> m (VU.Vector Int)
keys :: forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> m (Vector Int)
keys = ST (PrimState m) (Vector Int) -> m (Vector Int)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector Int) -> m (Vector Int))
-> (IntMap (PrimState m) a -> ST (PrimState m) (Vector Int))
-> IntMap (PrimState m) a
-> m (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> ST (PrimState m) (Vector Int)
forall s a. IntMap s a -> ST s (Vector Int)
keysST

-- | \(O(n \log n)\) Enumerates the elements (values) in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE elems #-}
elems :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector a)
elems :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> m (Vector a)
elems = ST (PrimState m) (Vector a) -> m (Vector a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector a) -> m (Vector a))
-> (IntMap (PrimState m) a -> ST (PrimState m) (Vector a))
-> IntMap (PrimState m) a
-> m (Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> ST (PrimState m) (Vector a)
forall a s. Unbox a => IntMap s a -> ST s (Vector a)
elemsST

-- | \(O(n \log n)\) Enumerates the key-value pairs in the map.
--
-- @since 1.1.0.0
{-# INLINEABLE assocs #-}
assocs :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector (Int, a))
assocs :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
IntMap (PrimState m) a -> m (Vector (Int, a))
assocs = ST (PrimState m) (Vector (Int, a)) -> m (Vector (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector (Int, a)) -> m (Vector (Int, a)))
-> (IntMap (PrimState m) a -> ST (PrimState m) (Vector (Int, a)))
-> IntMap (PrimState m) a
-> m (Vector (Int, a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (PrimState m) a -> ST (PrimState m) (Vector (Int, a))
forall a s. Unbox a => IntMap s a -> ST s (Vector (Int, a))
assocsST

-- -------------------------------------------------------------------------------------------------
-- Internal
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE newST #-}
newST :: (VU.Unbox a) => Int -> ST s (IntMap s a)
newST :: forall a s. Unbox a => Int -> ST s (IntMap s a)
newST Int
cap = do
  IntSet s
setIM <- Int -> ST s (IntSet (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (IntSet (PrimState m))
IS.new Int
cap
  MVector s a
valIM <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
cap
  IntMap s a -> ST s (IntMap s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntMap {MVector s a
IntSet s
setIM :: IntSet s
valIM :: MVector s a
setIM :: IntSet s
valIM :: MVector s a
..}

{-# INLINEABLE buildST #-}
buildST :: (VU.Unbox a) => Int -> VU.Vector (Int, a) -> ST s (IntMap s a)
buildST :: forall a s. Unbox a => Int -> Vector (Int, a) -> ST s (IntMap s a)
buildST Int
cap Vector (Int, a)
xs = do
  IntMap s a
im <- Int -> ST s (IntMap (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (IntMap (PrimState m) a)
new Int
cap
  Vector (Int, a) -> ((Int, a) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, a)
xs (((Int, a) -> ST s ()) -> ST s ())
-> ((Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
k, !a
v) -> do
    IntMap s a -> Int -> a -> ST s ()
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> a -> ST s ()
insertST IntMap s a
im Int
k a
v
  IntMap s a -> ST s (IntMap s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntMap s a
im

-- | \(O(\log n)\) Looks up the value for a key.
--
-- @since 1.1.0.0
{-# INLINEABLE lookupST #-}
lookupST :: (VU.Unbox a) => IntMap s a -> Int -> ST s (Maybe a)
lookupST :: forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe a)
lookupST im :: IntMap s a
im@IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} Int
k = do
  IntMap (PrimState (ST s)) a -> Int -> ST s Bool
forall (m :: * -> *) a.
PrimMonad m =>
IntMap (PrimState m) a -> Int -> m Bool
member IntMap s a
IntMap (PrimState (ST s)) a
im Int
k ST s Bool -> (Bool -> ST s (Maybe a)) -> ST s (Maybe a)
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Bool
True -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> ST s a -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
valIM Int
k
    Bool
False -> Maybe a -> ST s (Maybe a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

{-# INLINEABLE lookupGEST #-}
lookupGEST :: (VU.Unbox a) => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST :: forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} Int
k = do
  IntSet (PrimState (ST s)) -> Int -> ST s (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
IS.lookupGE IntSet s
IntSet (PrimState (ST s))
setIM Int
k ST s (Maybe Int)
-> (Maybe Int -> ST s (Maybe (Int, a))) -> ST s (Maybe (Int, a))
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just Int
i -> (Int, a) -> Maybe (Int, a)
forall a. a -> Maybe a
Just ((Int, a) -> Maybe (Int, a))
-> (a -> (Int, a)) -> a -> Maybe (Int, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
i,) (a -> Maybe (Int, a)) -> ST s a -> ST s (Maybe (Int, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
valIM Int
i
    Maybe Int
Nothing -> Maybe (Int, a) -> ST s (Maybe (Int, a))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, a)
forall a. Maybe a
Nothing

{-# INLINEABLE lookupLEST #-}
lookupLEST :: (HasCallStack, VU.Unbox a) => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupLEST :: forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupLEST IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} Int
k = do
  IntSet (PrimState (ST s)) -> Int -> ST s (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m (Maybe Int)
IS.lookupLE IntSet s
IntSet (PrimState (ST s))
setIM Int
k ST s (Maybe Int)
-> (Maybe Int -> ST s (Maybe (Int, a))) -> ST s (Maybe (Int, a))
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just Int
i -> (Int, a) -> Maybe (Int, a)
forall a. a -> Maybe a
Just ((Int, a) -> Maybe (Int, a))
-> (a -> (Int, a)) -> a -> Maybe (Int, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
i,) (a -> Maybe (Int, a)) -> ST s a -> ST s (Maybe (Int, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
valIM Int
i
    Maybe Int
Nothing -> Maybe (Int, a) -> ST s (Maybe (Int, a))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, a)
forall a. Maybe a
Nothing

{-# INLINEABLE lookupMinST #-}
lookupMinST :: (VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))
lookupMinST :: forall a s. Unbox a => IntMap s a -> ST s (Maybe (Int, a))
lookupMinST IntMap s a
is = IntMap s a -> Int -> ST s (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST IntMap s a
is Int
0

{-# INLINEABLE lookupMaxST #-}
lookupMaxST :: (VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))
lookupMaxST :: forall a s. Unbox a => IntMap s a -> ST s (Maybe (Int, a))
lookupMaxST IntMap s a
im = IntMap s a -> Int -> ST s (Maybe (Int, a))
forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupLEST IntMap s a
im (IntSet s -> Int
forall s. IntSet s -> Int
IS.capacity (IntMap s a -> IntSet s
forall s a. IntMap s a -> IntSet s
setIM IntMap s a
im) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

{-# INLINEABLE insertST #-}
insertST :: (HasCallStack, VU.Unbox a) => IntMap s a -> Int -> a -> ST s ()
insertST :: forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> Int -> a -> ST s ()
insertST IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} Int
k a
v = do
  IntSet (PrimState (ST s)) -> Int -> ST s ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
IntSet (PrimState m) -> Int -> m ()
IS.insert IntSet s
IntSet (PrimState (ST s))
setIM Int
k
  MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
valIM Int
k a
v

{-# INLINEABLE insertWithST #-}
insertWithST :: (HasCallStack, VU.Unbox a) => IntMap s a -> (a -> a -> a) -> Int -> a -> ST s ()
insertWithST :: forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> (a -> a -> a) -> Int -> a -> ST s ()
insertWithST IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} a -> a -> a
f Int
k a
v = do
  Bool
b <- IntSet (PrimState (ST s)) -> Int -> ST s Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
IS.member IntSet s
IntSet (PrimState (ST s))
setIM Int
k
  if Bool
b
    then do
      MVector (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s a
MVector (PrimState (ST s)) a
valIM (a -> a -> a
f a
v) Int
k
    else do
      IntSet (PrimState (ST s)) -> Int -> ST s ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
IntSet (PrimState m) -> Int -> m ()
IS.insert IntSet s
IntSet (PrimState (ST s))
setIM Int
k
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
valIM Int
k a
v

{-# INLINEABLE modifyST #-}
modifyST :: (HasCallStack, VU.Unbox a) => IntMap s a -> (a -> a) -> Int -> ST s ()
modifyST :: forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> (a -> a) -> Int -> ST s ()
modifyST IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} a -> a
f Int
k = do
  Bool
b <- IntSet (PrimState (ST s)) -> Int -> ST s Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
IS.member IntSet s
IntSet (PrimState (ST s))
setIM Int
k
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s a
MVector (PrimState (ST s)) a
valIM a -> a
f Int
k

{-# INLINEABLE deleteST #-}
deleteST :: IntMap s a -> Int -> ST s Bool
deleteST :: forall s a. IntMap s a -> Int -> ST s Bool
deleteST IntMap s a
im = IntSet (PrimState (ST s)) -> Int -> ST s Bool
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m Bool
IS.delete (IntMap s a -> IntSet s
forall s a. IntMap s a -> IntSet s
setIM IntMap s a
im)

{-# INLINEABLE deleteST_ #-}
deleteST_ :: IntMap s a -> Int -> ST s ()
deleteST_ :: forall s a. IntMap s a -> Int -> ST s ()
deleteST_ IntMap s a
im = IntSet (PrimState (ST s)) -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> Int -> m ()
IS.delete_ (IntMap s a -> IntSet s
forall s a. IntMap s a -> IntSet s
setIM IntMap s a
im)

{-# INLINEABLE deleteMinST #-}
deleteMinST :: (HasCallStack, VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))
deleteMinST :: forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> ST s (Maybe (Int, a))
deleteMinST IntMap s a
is = do
  IntMap s a -> ST s (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> ST s (Maybe (Int, a))
lookupMinST IntMap s a
is
    ST s (Maybe (Int, a))
-> (Maybe (Int, a) -> ST s (Maybe (Int, a)))
-> ST s (Maybe (Int, a))
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ((Int, a) -> ST s (Int, a))
-> Maybe (Int, a) -> ST s (Maybe (Int, a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Maybe a -> m (Maybe b)
mapM
      ( \(!Int
key, !a
val) -> do
          IntMap s a -> Int -> ST s ()
forall s a. IntMap s a -> Int -> ST s ()
deleteST_ IntMap s a
is Int
key
          (Int, a) -> ST s (Int, a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
key, a
val)
      )

{-# INLINEABLE deleteMaxST #-}
deleteMaxST :: (HasCallStack, VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))
deleteMaxST :: forall a s.
(HasCallStack, Unbox a) =>
IntMap s a -> ST s (Maybe (Int, a))
deleteMaxST IntMap s a
is = do
  IntMap s a -> ST s (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> ST s (Maybe (Int, a))
lookupMaxST IntMap s a
is
    ST s (Maybe (Int, a))
-> (Maybe (Int, a) -> ST s (Maybe (Int, a)))
-> ST s (Maybe (Int, a))
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ((Int, a) -> ST s (Int, a))
-> Maybe (Int, a) -> ST s (Maybe (Int, a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Maybe a -> m (Maybe b)
mapM
      ( \(!Int
k, !a
v) -> do
          IntMap s a -> Int -> ST s ()
forall s a. IntMap s a -> Int -> ST s ()
deleteST_ IntMap s a
is Int
k
          (Int, a) -> ST s (Int, a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
k, a
v)
      )

{-# INLINEABLE keysST #-}
keysST :: IntMap s a -> ST s (VU.Vector Int)
keysST :: forall s a. IntMap s a -> ST s (Vector Int)
keysST = IntSet s -> ST s (Vector Int)
IntSet (PrimState (ST s)) -> ST s (Vector Int)
forall (m :: * -> *).
PrimMonad m =>
IntSet (PrimState m) -> m (Vector Int)
IS.keys (IntSet s -> ST s (Vector Int))
-> (IntMap s a -> IntSet s) -> IntMap s a -> ST s (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap s a -> IntSet s
forall s a. IntMap s a -> IntSet s
setIM

{-# INLINEABLE elemsST #-}
elemsST :: (VU.Unbox a) => IntMap s a -> ST s (VU.Vector a)
elemsST :: forall a s. Unbox a => IntMap s a -> ST s (Vector a)
elemsST im :: IntMap s a
im@IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} = do
  Int
n <- IntSet (PrimState (ST s)) -> ST s Int
forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Int
IS.size IntSet s
IntSet (PrimState (ST s))
setIM
  Int -> (Int -> ST s (a, Int)) -> Int -> ST s (Vector a)
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Int -> (b -> m (a, b)) -> b -> m (Vector a)
VU.unfoldrExactNM
    Int
n
    ( \Int
i -> do
        (!Int
i', !a
x') <- Maybe (Int, a) -> (Int, a)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Int, a) -> (Int, a))
-> ST s (Maybe (Int, a)) -> ST s (Int, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap s a -> Int -> ST s (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST IntMap s a
im (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        (a, Int) -> ST s (a, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x', Int
i')
    )
    (-Int
1)

{-# INLINEABLE assocsST #-}
assocsST :: (VU.Unbox a) => IntMap s a -> ST s (VU.Vector (Int, a))
assocsST :: forall a s. Unbox a => IntMap s a -> ST s (Vector (Int, a))
assocsST im :: IntMap s a
im@IntMap {MVector s a
IntSet s
setIM :: forall s a. IntMap s a -> IntSet s
valIM :: forall s a. IntMap s a -> MVector s a
setIM :: IntSet s
valIM :: MVector s a
..} = do
  Int
n <- IntSet (PrimState (ST s)) -> ST s Int
forall (m :: * -> *). PrimMonad m => IntSet (PrimState m) -> m Int
IS.size IntSet s
IntSet (PrimState (ST s))
setIM
  Int
-> (Int -> ST s ((Int, a), Int)) -> Int -> ST s (Vector (Int, a))
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Int -> (b -> m (a, b)) -> b -> m (Vector a)
VU.unfoldrExactNM
    Int
n
    ( \Int
i -> do
        (!Int
i', !a
x') <- Maybe (Int, a) -> (Int, a)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Int, a) -> (Int, a))
-> ST s (Maybe (Int, a)) -> ST s (Int, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap s a -> Int -> ST s (Maybe (Int, a))
forall a s. Unbox a => IntMap s a -> Int -> ST s (Maybe (Int, a))
lookupGEST IntMap s a
im (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        ((Int, a), Int) -> ST s ((Int, a), Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int
i', a
x'), Int
i')
    )
    (-Int
1)