{-|
Module      : Z.Data.Vector.FlatIntMap
Description : Fast int map based on sorted vector
Copyright   : (c) Dong Han, 2017-2019
              (c) Tao He, 2018-2019
License     : BSD
Maintainer  : winterland1989@gmail.com
Stability   : experimental
Portability : non-portable

This module provides a simple int key value map based on sorted vector and binary search. It's particularly
suitable for small sized key value collections such as deserializing intermediate representation.
But can also used in various place where insertion and deletion is rare but require fast lookup.

-}

module Z.Data.Vector.FlatIntMap
  ( -- * FlatIntMap backed by sorted vector
    FlatIntMap, sortedKeyValues, size, null, empty, map', imap'
  , pack, packN, packR, packRN
  , unpack, unpackR, packVector, packVectorR
  , lookup
  , delete
  , insert
  , adjust'
  , merge, mergeWithKey'
    -- * fold and traverse
  , foldrWithKey, foldrWithKey', foldlWithKey, foldlWithKey', traverseWithKey
    -- * binary & linear search on vectors
  , linearSearch, linearSearchR, binarySearch
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.ST
import qualified Data.Foldable              as Foldable
import qualified Data.Traversable           as Traversable
import qualified Data.Semigroup             as Semigroup
import qualified Data.Monoid                as Monoid
import qualified Data.Primitive.SmallArray  as A
import qualified Z.Data.Vector.Base         as V
import qualified Z.Data.Vector.Sort         as V
import qualified Z.Data.Text.Print          as T
import           Data.Function              (on)
import           Data.Bits                  (unsafeShiftR)
import           Data.Data
import           Prelude hiding (lookup, null)
import           Test.QuickCheck.Arbitrary (Arbitrary(..), CoArbitrary(..))

--------------------------------------------------------------------------------

newtype FlatIntMap v = FlatIntMap { FlatIntMap v -> Vector (IPair v)
sortedKeyValues :: V.Vector (V.IPair v) }
    deriving (Int -> FlatIntMap v -> ShowS
[FlatIntMap v] -> ShowS
FlatIntMap v -> String
(Int -> FlatIntMap v -> ShowS)
-> (FlatIntMap v -> String)
-> ([FlatIntMap v] -> ShowS)
-> Show (FlatIntMap v)
forall v. Show v => Int -> FlatIntMap v -> ShowS
forall v. Show v => [FlatIntMap v] -> ShowS
forall v. Show v => FlatIntMap v -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FlatIntMap v] -> ShowS
$cshowList :: forall v. Show v => [FlatIntMap v] -> ShowS
show :: FlatIntMap v -> String
$cshow :: forall v. Show v => FlatIntMap v -> String
showsPrec :: Int -> FlatIntMap v -> ShowS
$cshowsPrec :: forall v. Show v => Int -> FlatIntMap v -> ShowS
Show, FlatIntMap v -> FlatIntMap v -> Bool
(FlatIntMap v -> FlatIntMap v -> Bool)
-> (FlatIntMap v -> FlatIntMap v -> Bool) -> Eq (FlatIntMap v)
forall v. Eq v => FlatIntMap v -> FlatIntMap v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FlatIntMap v -> FlatIntMap v -> Bool
$c/= :: forall v. Eq v => FlatIntMap v -> FlatIntMap v -> Bool
== :: FlatIntMap v -> FlatIntMap v -> Bool
$c== :: forall v. Eq v => FlatIntMap v -> FlatIntMap v -> Bool
Eq, Eq (FlatIntMap v)
Eq (FlatIntMap v)
-> (FlatIntMap v -> FlatIntMap v -> Ordering)
-> (FlatIntMap v -> FlatIntMap v -> Bool)
-> (FlatIntMap v -> FlatIntMap v -> Bool)
-> (FlatIntMap v -> FlatIntMap v -> Bool)
-> (FlatIntMap v -> FlatIntMap v -> Bool)
-> (FlatIntMap v -> FlatIntMap v -> FlatIntMap v)
-> (FlatIntMap v -> FlatIntMap v -> FlatIntMap v)
-> Ord (FlatIntMap v)
FlatIntMap v -> FlatIntMap v -> Bool
FlatIntMap v -> FlatIntMap v -> Ordering
FlatIntMap v -> FlatIntMap v -> FlatIntMap v
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall v. Ord v => Eq (FlatIntMap v)
forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Bool
forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Ordering
forall v. Ord v => FlatIntMap v -> FlatIntMap v -> FlatIntMap v
min :: FlatIntMap v -> FlatIntMap v -> FlatIntMap v
$cmin :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> FlatIntMap v
max :: FlatIntMap v -> FlatIntMap v -> FlatIntMap v
$cmax :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> FlatIntMap v
>= :: FlatIntMap v -> FlatIntMap v -> Bool
$c>= :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Bool
> :: FlatIntMap v -> FlatIntMap v -> Bool
$c> :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Bool
<= :: FlatIntMap v -> FlatIntMap v -> Bool
$c<= :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Bool
< :: FlatIntMap v -> FlatIntMap v -> Bool
$c< :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Bool
compare :: FlatIntMap v -> FlatIntMap v -> Ordering
$ccompare :: forall v. Ord v => FlatIntMap v -> FlatIntMap v -> Ordering
$cp1Ord :: forall v. Ord v => Eq (FlatIntMap v)
Ord, Typeable)

instance T.Print v => T.Print (FlatIntMap v) where
    {-# INLINE toUTF8BuilderP #-}
    toUTF8BuilderP :: Int -> FlatIntMap v -> Builder ()
toUTF8BuilderP Int
p (FlatIntMap Vector (IPair v)
vec) = Bool -> Builder () -> Builder ()
T.parenWhen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (Builder () -> Builder ()) -> Builder () -> Builder ()
forall a b. (a -> b) -> a -> b
$ do
        Builder ()
"FlatIntMap{"
        Builder ()
-> (IPair v -> Builder ()) -> Vector (IPair v) -> Builder ()
forall (v :: * -> *) a.
Vec v a =>
Builder () -> (a -> Builder ()) -> v a -> Builder ()
T.intercalateVec Builder ()
T.comma (\ (V.IPair Int
i v
v) ->
            Int -> Builder ()
forall a. Print a => a -> Builder ()
T.toUTF8Builder Int
i Builder () -> Builder () -> Builder ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Char -> Builder ()
T.char7 Char
':' Builder () -> Builder () -> Builder ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> v -> Builder ()
forall a. Print a => a -> Builder ()
T.toUTF8Builder v
v) Vector (IPair v)
vec
        Char -> Builder ()
T.char7 Char
'}'

instance (Arbitrary v) => Arbitrary (FlatIntMap v) where
    arbitrary :: Gen (FlatIntMap v)
arbitrary = [IPair v] -> FlatIntMap v
forall v. [IPair v] -> FlatIntMap v
pack ([IPair v] -> FlatIntMap v) -> Gen [IPair v] -> Gen (FlatIntMap v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [IPair v]
forall a. Arbitrary a => Gen a
arbitrary
    shrink :: FlatIntMap v -> [FlatIntMap v]
shrink FlatIntMap v
v = [IPair v] -> FlatIntMap v
forall v. [IPair v] -> FlatIntMap v
pack ([IPair v] -> FlatIntMap v) -> [[IPair v]] -> [FlatIntMap v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [IPair v] -> [[IPair v]]
forall a. Arbitrary a => a -> [a]
shrink (FlatIntMap v -> [IPair v]
forall v. FlatIntMap v -> [IPair v]
unpack FlatIntMap v
v)

instance (CoArbitrary v) => CoArbitrary (FlatIntMap v) where
    coarbitrary :: FlatIntMap v -> Gen b -> Gen b
coarbitrary = [IPair v] -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary ([IPair v] -> Gen b -> Gen b)
-> (FlatIntMap v -> [IPair v]) -> FlatIntMap v -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatIntMap v -> [IPair v]
forall v. FlatIntMap v -> [IPair v]
unpack

instance Semigroup.Semigroup (FlatIntMap v) where
    {-# INLINE (<>) #-}
    <> :: FlatIntMap v -> FlatIntMap v -> FlatIntMap v
(<>) = FlatIntMap v -> FlatIntMap v -> FlatIntMap v
forall v. FlatIntMap v -> FlatIntMap v -> FlatIntMap v
merge

instance Monoid.Monoid (FlatIntMap v) where
    {-# INLINE mappend #-}
    mappend :: FlatIntMap v -> FlatIntMap v -> FlatIntMap v
mappend = FlatIntMap v -> FlatIntMap v -> FlatIntMap v
forall v. FlatIntMap v -> FlatIntMap v -> FlatIntMap v
merge
    {-# INLINE mempty #-}
    mempty :: FlatIntMap v
mempty = FlatIntMap v
forall v. FlatIntMap v
empty

instance NFData v => NFData (FlatIntMap v) where
    {-# INLINE rnf #-}
    rnf :: FlatIntMap v -> ()
rnf (FlatIntMap Vector (IPair v)
ivs) = Vector (IPair v) -> ()
forall a. NFData a => a -> ()
rnf Vector (IPair v)
ivs

instance Functor (FlatIntMap) where
    {-# INLINE fmap #-}
    fmap :: (a -> b) -> FlatIntMap a -> FlatIntMap b
fmap a -> b
f (FlatIntMap Vector (IPair a)
vs) = Vector (IPair b) -> FlatIntMap b
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair a -> IPair b) -> Vector (IPair a) -> Vector (IPair b)
forall (u :: * -> *) (v :: * -> *) a b.
(Vec u a, Vec v b) =>
(a -> b) -> u a -> v b
V.map' ((a -> b) -> IPair a -> IPair b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) Vector (IPair a)
vs)

instance Foldable.Foldable FlatIntMap where
    {-# INLINE foldr' #-}
    foldr' :: (a -> b -> b) -> b -> FlatIntMap a -> b
foldr' a -> b -> b
f = (Int -> a -> b -> b) -> b -> FlatIntMap a -> b
forall v a. (Int -> v -> a -> a) -> a -> FlatIntMap v -> a
foldrWithKey' ((a -> b -> b) -> Int -> a -> b -> b
forall a b. a -> b -> a
const a -> b -> b
f)
    {-# INLINE foldr #-}
    foldr :: (a -> b -> b) -> b -> FlatIntMap a -> b
foldr a -> b -> b
f = (Int -> a -> b -> b) -> b -> FlatIntMap a -> b
forall v a. (Int -> v -> a -> a) -> a -> FlatIntMap v -> a
foldrWithKey ((a -> b -> b) -> Int -> a -> b -> b
forall a b. a -> b -> a
const a -> b -> b
f)
    {-# INLINE foldl' #-}
    foldl' :: (b -> a -> b) -> b -> FlatIntMap a -> b
foldl' b -> a -> b
f = (b -> Int -> a -> b) -> b -> FlatIntMap a -> b
forall a v. (a -> Int -> v -> a) -> a -> FlatIntMap v -> a
foldlWithKey' (\ b
a Int
_ a
v -> b -> a -> b
f b
a a
v)
    {-# INLINE foldl #-}
    foldl :: (b -> a -> b) -> b -> FlatIntMap a -> b
foldl b -> a -> b
f = (b -> Int -> a -> b) -> b -> FlatIntMap a -> b
forall a v. (a -> Int -> v -> a) -> a -> FlatIntMap v -> a
foldlWithKey (\ b
a Int
_ a
v -> b -> a -> b
f b
a a
v)
    {-# INLINE toList #-}
    toList :: FlatIntMap a -> [a]
toList = (IPair a -> a) -> [IPair a] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap IPair a -> a
forall a. IPair a -> a
V.isnd ([IPair a] -> [a])
-> (FlatIntMap a -> [IPair a]) -> FlatIntMap a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatIntMap a -> [IPair a]
forall v. FlatIntMap v -> [IPair v]
unpack
    {-# INLINE null #-}
    null :: FlatIntMap a -> Bool
null (FlatIntMap Vector (IPair a)
vs) = Vector (IPair a) -> Bool
forall (v :: * -> *) a. Vec v a => v a -> Bool
V.null Vector (IPair a)
vs
    {-# INLINE length #-}
    length :: FlatIntMap a -> Int
length (FlatIntMap Vector (IPair a)
vs) = Vector (IPair a) -> Int
forall (v :: * -> *) a. Vec v a => v a -> Int
V.length Vector (IPair a)
vs
    {-# INLINE elem #-}
    elem :: a -> FlatIntMap a -> Bool
elem a
a (FlatIntMap Vector (IPair a)
vs) = a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem a
a ((IPair a -> a) -> [IPair a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map IPair a -> a
forall a. IPair a -> a
V.isnd ([IPair a] -> [a]) -> [IPair a] -> [a]
forall a b. (a -> b) -> a -> b
$ Vector (IPair a) -> [IPair a]
forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpack Vector (IPair a)
vs)

instance Traversable.Traversable FlatIntMap where
    {-# INLINE traverse #-}
    traverse :: (a -> f b) -> FlatIntMap a -> f (FlatIntMap b)
traverse a -> f b
f = (Int -> a -> f b) -> FlatIntMap a -> f (FlatIntMap b)
forall (t :: * -> *) a b.
Applicative t =>
(Int -> a -> t b) -> FlatIntMap a -> t (FlatIntMap b)
traverseWithKey ((a -> f b) -> Int -> a -> f b
forall a b. a -> b -> a
const a -> f b
f)

size :: FlatIntMap v -> Int
{-# INLINE size #-}
size :: FlatIntMap v -> Int
size = Vector (IPair v) -> Int
forall (v :: * -> *) a. Vec v a => v a -> Int
V.length (Vector (IPair v) -> Int)
-> (FlatIntMap v -> Vector (IPair v)) -> FlatIntMap v -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatIntMap v -> Vector (IPair v)
forall v. FlatIntMap v -> Vector (IPair v)
sortedKeyValues

null :: FlatIntMap v -> Bool
{-# INLINE null #-}
null :: FlatIntMap v -> Bool
null = Vector (IPair v) -> Bool
forall (v :: * -> *) a. Vec v a => v a -> Bool
V.null (Vector (IPair v) -> Bool)
-> (FlatIntMap v -> Vector (IPair v)) -> FlatIntMap v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatIntMap v -> Vector (IPair v)
forall v. FlatIntMap v -> Vector (IPair v)
sortedKeyValues

map' :: (v -> v') -> FlatIntMap v -> FlatIntMap v'
{-# INLINE map' #-}
map' :: (v -> v') -> FlatIntMap v -> FlatIntMap v'
map' v -> v'
f (FlatIntMap Vector (IPair v)
vs) = Vector (IPair v') -> FlatIntMap v'
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v') -> Vector (IPair v) -> Vector (IPair v')
forall (u :: * -> *) (v :: * -> *) a b.
(Vec u a, Vec v b) =>
(a -> b) -> u a -> v b
V.map' ((v -> v') -> IPair v -> IPair v'
forall a b. (a -> b) -> IPair a -> IPair b
V.mapIPair' v -> v'
f) Vector (IPair v)
vs)

imap' :: (Int -> v -> v') -> FlatIntMap v -> FlatIntMap v'
{-# INLINE imap' #-}
imap' :: (Int -> v -> v') -> FlatIntMap v -> FlatIntMap v'
imap' Int -> v -> v'
f (FlatIntMap Vector (IPair v)
vs) = Vector (IPair v') -> FlatIntMap v'
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((Int -> IPair v -> IPair v')
-> Vector (IPair v) -> Vector (IPair v')
forall (u :: * -> *) (v :: * -> *) a b.
(Vec u a, Vec v b) =>
(Int -> a -> b) -> u a -> v b
V.imap' (\ Int
i -> (v -> v') -> IPair v -> IPair v'
forall a b. (a -> b) -> IPair a -> IPair b
V.mapIPair' (Int -> v -> v'
f Int
i)) Vector (IPair v)
vs)

-- | /O(1)/ empty flat map.
empty :: FlatIntMap v
{-# INLINE empty #-}
empty :: FlatIntMap v
empty = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap Vector (IPair v)
forall (v :: * -> *) a. Vec v a => v a
V.empty

-- | /O(N*logN)/ Pack list of key values, on key duplication prefer left one.
pack :: [V.IPair v] -> FlatIntMap v
{-# INLINE pack #-}
pack :: [IPair v] -> FlatIntMap v
pack [IPair v]
kvs = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v -> Bool)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool)
-> (IPair v -> Int) -> IPair v -> IPair v -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ((IPair v -> IPair v -> Ordering)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (IPair v -> Int) -> IPair v -> IPair v -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ([IPair v] -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => [a] -> v a
V.pack [IPair v]
kvs)))

-- | /O(N*logN)/ Pack list of key values with suggested size, on key duplication prefer left one.
packN :: Int -> [V.IPair v] -> FlatIntMap v
{-# INLINE packN #-}
packN :: Int -> [IPair v] -> FlatIntMap v
packN Int
n [IPair v]
kvs = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v -> Bool)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool)
-> (IPair v -> Int) -> IPair v -> IPair v -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ((IPair v -> IPair v -> Ordering)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (IPair v -> Int) -> IPair v -> IPair v -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) (Int -> [IPair v] -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => Int -> [a] -> v a
V.packN Int
n [IPair v]
kvs)))

-- | /O(N*logN)/ Pack list of key values, on key duplication prefer right one.
packR :: [V.IPair v] -> FlatIntMap v
{-# INLINE packR #-}
packR :: [IPair v] -> FlatIntMap v
packR [IPair v]
kvs = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v -> Bool)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool)
-> (IPair v -> Int) -> IPair v -> IPair v -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ((IPair v -> IPair v -> Ordering)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (IPair v -> Int) -> IPair v -> IPair v -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ([IPair v] -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => [a] -> v a
V.pack [IPair v]
kvs)))

-- | /O(N*logN)/ Pack list of key values with suggested size, on key duplication prefer right one.
packRN :: Int -> [V.IPair v] -> FlatIntMap v
{-# INLINE packRN #-}
packRN :: Int -> [IPair v] -> FlatIntMap v
packRN Int
n [IPair v]
kvs = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v -> Bool)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool)
-> (IPair v -> Int) -> IPair v -> IPair v -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ((IPair v -> IPair v -> Ordering)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (IPair v -> Int) -> IPair v -> IPair v -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) (Int -> [IPair v] -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => Int -> [a] -> v a
V.packN Int
n [IPair v]
kvs)))

-- | /O(N)/ Unpack key value pairs to a list sorted by keys in ascending order.
--
-- This function works with @foldr/build@ fusion in base.
unpack :: FlatIntMap v -> [V.IPair v]
{-# INLINE unpack #-}
unpack :: FlatIntMap v -> [IPair v]
unpack = Vector (IPair v) -> [IPair v]
forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpack (Vector (IPair v) -> [IPair v])
-> (FlatIntMap v -> Vector (IPair v)) -> FlatIntMap v -> [IPair v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatIntMap v -> Vector (IPair v)
forall v. FlatIntMap v -> Vector (IPair v)
sortedKeyValues

-- | /O(N)/ Unpack key value pairs to a list sorted by keys in descending order.
--
-- This function works with @foldr/build@ fusion in base.
unpackR :: FlatIntMap v -> [V.IPair v]
{-# INLINE unpackR #-}
unpackR :: FlatIntMap v -> [IPair v]
unpackR = Vector (IPair v) -> [IPair v]
forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpackR (Vector (IPair v) -> [IPair v])
-> (FlatIntMap v -> Vector (IPair v)) -> FlatIntMap v -> [IPair v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FlatIntMap v -> Vector (IPair v)
forall v. FlatIntMap v -> Vector (IPair v)
sortedKeyValues

-- | /O(N*logN)/ Pack vector of key values, on key duplication prefer left one.
packVector :: V.Vector (V.IPair v) -> FlatIntMap v
{-# INLINE packVector #-}
packVector :: Vector (IPair v) -> FlatIntMap v
packVector Vector (IPair v)
kvs = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v -> Bool)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool)
-> (IPair v -> Int) -> IPair v -> IPair v -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ((IPair v -> IPair v -> Ordering)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (IPair v -> Int) -> IPair v -> IPair v -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) Vector (IPair v)
kvs))

-- | /O(N*logN)/ Pack vector of key values, on key duplication prefer right one.
packVectorR :: V.Vector (V.IPair v) -> FlatIntMap v
{-# INLINE packVectorR #-}
packVectorR :: Vector (IPair v) -> FlatIntMap v
packVectorR Vector (IPair v)
kvs = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap ((IPair v -> IPair v -> Bool)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool)
-> (IPair v -> Int) -> IPair v -> IPair v -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) ((IPair v -> IPair v -> Ordering)
-> Vector (IPair v) -> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (IPair v -> Int) -> IPair v -> IPair v -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` IPair v -> Int
forall a. IPair a -> Int
V.ifst) Vector (IPair v)
kvs))

-- | /O(logN)/ Binary search on flat map.
lookup :: Int -> FlatIntMap v -> Maybe v
{-# INLINABLE lookup #-}
lookup :: Int -> FlatIntMap v -> Maybe v
lookup Int
_ (FlatIntMap (V.Vector SmallArray (IPair v)
_ Int
_ Int
0)) = Maybe v
forall a. Maybe a
Nothing
lookup Int
k' (FlatIntMap (V.Vector SmallArray (IPair v)
arr Int
s0 Int
l)) = Int -> Int -> Maybe v
go Int
s0 (Int
s0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Int -> Maybe v
go !Int
s !Int
e
        | Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
e =
            case SmallArray (IPair v)
arr SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
s of (V.IPair Int
k v
v) | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k'  -> v -> Maybe v
forall a. a -> Maybe a
Just v
v
                                                            | Bool
otherwise -> Maybe v
forall a. Maybe a
Nothing
        | Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>  Int
e = Maybe v
forall a. Maybe a
Nothing
        | Bool
otherwise =
            let mid :: Int
mid = (Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
e) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
                (V.IPair Int
k v
v)  = SmallArray (IPair v)
arr SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
mid
            in case Int
k' Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
k of Ordering
LT -> Int -> Int -> Maybe v
go Int
s (Int
midInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                                      Ordering
GT -> Int -> Int -> Maybe v
go (Int
midInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
e
                                      Ordering
_  -> v -> Maybe v
forall a. a -> Maybe a
Just v
v

-- | /O(N)/ Insert new key value into map, replace old one if key exists.
insert :: Int -> v -> FlatIntMap v -> FlatIntMap v
{-# INLINE insert #-}
insert :: Int -> v -> FlatIntMap v -> FlatIntMap v
insert Int
k v
v (FlatIntMap vec :: Vector (IPair v)
vec@(V.Vector SmallArray (IPair v)
arr Int
s Int
l)) =
    case Vector (IPair v) -> Int -> Either Int Int
forall v. Vector (IPair v) -> Int -> Either Int Int
binarySearch Vector (IPair v)
vec Int
k of
        Left Int
i -> Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (Int
-> (forall s. MArr (IArray Vector) s (IPair v) -> ST s ())
-> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
Int -> (forall s. MArr (IArray v) s a -> ST s ()) -> v a
V.create (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (\ MArr (IArray Vector) s (IPair v)
marr -> do
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
iInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr Int
0 SmallArray (IPair v)
arr Int
s Int
i
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr Int
i (Int -> v -> IPair v
forall a. Int -> a -> IPair a
V.IPair Int
k v
v)
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
iInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
l) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallArray (IPair v)
arr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
s) (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)))
        Right Int
i -> Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (SmallArray (IPair v) -> Int -> Int -> Vector (IPair v)
forall a. SmallArray a -> Int -> Int -> Vector a
V.Vector ((forall s. ST s (SmallArray (IPair v))) -> SmallArray (IPair v)
forall a. (forall s. ST s a) -> a
runST (do
            let arr' :: SmallArray (IPair v)
arr' = SmallArray (IPair v) -> Int -> Int -> SmallArray (IPair v)
forall a. SmallArray a -> Int -> Int -> SmallArray a
A.cloneSmallArray SmallArray (IPair v)
arr Int
s Int
l
            SmallMutableArray s (IPair v)
marr <- SmallArray (IPair v)
-> ST s (SmallMutableArray (PrimState (ST s)) (IPair v))
forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> m (SmallMutableArray (PrimState m) a)
A.unsafeThawSmallArray SmallArray (IPair v)
arr'
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
i (Int -> v -> IPair v
forall a. Int -> a -> IPair a
V.IPair Int
k v
v)
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> ST s (SmallArray (IPair v))
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
A.unsafeFreezeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr)) Int
0 Int
l)

-- | /O(N)/ Delete a key value pair by key.
delete :: Int -> FlatIntMap v -> FlatIntMap v
{-# INLINE delete #-}
delete :: Int -> FlatIntMap v -> FlatIntMap v
delete Int
k m :: FlatIntMap v
m@(FlatIntMap vec :: Vector (IPair v)
vec@(V.Vector SmallArray (IPair v)
arr Int
s Int
l)) =
    case Vector (IPair v) -> Int -> Either Int Int
forall v. Vector (IPair v) -> Int -> Either Int Int
binarySearch Vector (IPair v)
vec Int
k of
        Left Int
_  -> FlatIntMap v
m
        Right Int
i -> Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (Vector (IPair v) -> FlatIntMap v)
-> Vector (IPair v) -> FlatIntMap v
forall a b. (a -> b) -> a -> b
$ Int
-> (forall s. MArr (IArray Vector) s (IPair v) -> ST s ())
-> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
Int -> (forall s. MArr (IArray v) s a -> ST s ()) -> v a
V.create (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (\ MArr (IArray Vector) s (IPair v)
marr -> do
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
iInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr Int
0 SmallArray (IPair v)
arr Int
s Int
i
            let i' :: Int
i' = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i'Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
l) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr Int
i SmallArray (IPair v)
arr (Int
i'Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
s) (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i'))

-- | /O(N)/ Modify a value by key.
--
-- The value is evaluated to WHNF before writing into map.
adjust' :: (v -> v) -> Int -> FlatIntMap v -> FlatIntMap v
{-# INLINE adjust' #-}
adjust' :: (v -> v) -> Int -> FlatIntMap v -> FlatIntMap v
adjust' v -> v
f Int
k m :: FlatIntMap v
m@(FlatIntMap vec :: Vector (IPair v)
vec@(V.Vector SmallArray (IPair v)
arr Int
s Int
l)) =
    case Vector (IPair v) -> Int -> Either Int Int
forall v. Vector (IPair v) -> Int -> Either Int Int
binarySearch Vector (IPair v)
vec Int
k of
        Left Int
_  -> FlatIntMap v
m
        Right Int
i -> Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (Vector (IPair v) -> FlatIntMap v)
-> Vector (IPair v) -> FlatIntMap v
forall a b. (a -> b) -> a -> b
$ Int
-> (forall s. MArr (IArray Vector) s (IPair v) -> ST s ())
-> Vector (IPair v)
forall (v :: * -> *) a.
Vec v a =>
Int -> (forall s. MArr (IArray v) s a -> ST s ()) -> v a
V.create Int
l (\ MArr (IArray Vector) s (IPair v)
marr -> do
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr Int
0 SmallArray (IPair v)
arr Int
s Int
l
            let !v' :: v
v' = v -> v
f (IPair v -> v
forall a. IPair a -> a
V.isnd (SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
A.indexSmallArray SmallArray (IPair v)
arr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
s)))
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray (PrimState (ST s)) (IPair v)
MArr (IArray Vector) s (IPair v)
marr Int
i (Int -> v -> IPair v
forall a. Int -> a -> IPair a
V.IPair Int
k v
v'))

-- | /O(n+m)/ Merge two 'FlatIntMap', prefer right value on key duplication.
merge :: forall v. FlatIntMap v -> FlatIntMap v -> FlatIntMap v
{-# INLINE merge #-}
merge :: FlatIntMap v -> FlatIntMap v -> FlatIntMap v
merge fmL :: FlatIntMap v
fmL@(FlatIntMap (V.Vector SmallArray (IPair v)
arrL Int
sL Int
lL)) fmR :: FlatIntMap v
fmR@(FlatIntMap (V.Vector SmallArray (IPair v)
arrR Int
sR Int
lR))
    | FlatIntMap v -> Bool
forall a. FlatIntMap a -> Bool
null FlatIntMap v
fmL = FlatIntMap v
fmR
    | FlatIntMap v -> Bool
forall a. FlatIntMap a -> Bool
null FlatIntMap v
fmR = FlatIntMap v
fmL
    | Bool
otherwise = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (Int
-> (forall s. MArr (IArray Vector) s (IPair v) -> ST s Int)
-> Vector (IPair v)
forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
Int -> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
V.createN (Int
lLInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lR) (Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go Int
sL Int
sR Int
0))
  where
    endL :: Int
endL = Int
sL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lL
    endR :: Int
endR = Int
sR Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lR
    go :: Int -> Int -> Int -> A.SmallMutableArray s (V.IPair v) -> ST s Int
    go :: Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go !Int
i !Int
j !Int
k SmallMutableArray s (IPair v)
marr
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
endL = do
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k SmallArray (IPair v)
arrR Int
j (Int
lRInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
j)
            Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lRInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
j
        | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
endR = do
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k SmallArray (IPair v)
arrL Int
i (Int
lLInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
            Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lLInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i
        | Bool
otherwise = do
            kvL :: IPair v
kvL@(V.IPair Int
kL v
_) <- SmallArray (IPair v)
arrL SmallArray (IPair v) -> Int -> ST s (IPair v)
forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
i
            kvR :: IPair v
kvR@(V.IPair Int
kR v
_) <- SmallArray (IPair v)
arrR SmallArray (IPair v) -> Int -> ST s (IPair v)
forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
j
            case Int
kL Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
kR of Ordering
LT -> do SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k IPair v
kvL
                                             Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (IPair v)
marr
                                    Ordering
EQ -> do SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k IPair v
kvR
                                             Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (IPair v)
marr
                                    Ordering
_  -> do SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k IPair v
kvR
                                             Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go Int
i (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (IPair v)
marr

-- | /O(n+m)/ Merge two 'FlatIntMap' with a merge function.
mergeWithKey' :: forall v. (Int -> v -> v -> v) -> FlatIntMap v -> FlatIntMap v -> FlatIntMap v
{-# INLINABLE mergeWithKey' #-}
mergeWithKey' :: (Int -> v -> v -> v)
-> FlatIntMap v -> FlatIntMap v -> FlatIntMap v
mergeWithKey' Int -> v -> v -> v
f fmL :: FlatIntMap v
fmL@(FlatIntMap (V.Vector SmallArray (IPair v)
arrL Int
sL Int
lL)) fmR :: FlatIntMap v
fmR@(FlatIntMap (V.Vector SmallArray (IPair v)
arrR Int
sR Int
lR))
    | FlatIntMap v -> Bool
forall a. FlatIntMap a -> Bool
null FlatIntMap v
fmL = FlatIntMap v
fmR
    | FlatIntMap v -> Bool
forall a. FlatIntMap a -> Bool
null FlatIntMap v
fmR = FlatIntMap v
fmL
    | Bool
otherwise = Vector (IPair v) -> FlatIntMap v
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (Int
-> (forall s. MArr (IArray Vector) s (IPair v) -> ST s Int)
-> Vector (IPair v)
forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
Int -> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
V.createN (Int
lLInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lR) (Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go Int
sL Int
sR Int
0))
  where
    endL :: Int
endL = Int
sL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lL
    endR :: Int
endR = Int
sR Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lR
    go :: Int -> Int -> Int -> A.SmallMutableArray s (V.IPair v) -> ST s Int
    go :: Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go !Int
i !Int
j !Int
k SmallMutableArray s (IPair v)
marr
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
endL = do
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k SmallArray (IPair v)
arrR Int
j (Int
lRInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
j)
            Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lRInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
j
        | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
endR = do
            SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> SmallArray (IPair v) -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k SmallArray (IPair v)
arrL Int
i (Int
lLInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
            Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lLInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i
        | Bool
otherwise = do
            kvL :: IPair v
kvL@(V.IPair Int
kL v
vL) <- SmallArray (IPair v)
arrL SmallArray (IPair v) -> Int -> ST s (IPair v)
forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
i
            kvR :: IPair v
kvR@(V.IPair Int
kR v
vR) <- SmallArray (IPair v)
arrR SmallArray (IPair v) -> Int -> ST s (IPair v)
forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
j
            case Int
kL Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
kR of Ordering
LT -> do SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k IPair v
kvL
                                             Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (IPair v)
marr
                                    Ordering
EQ -> do let !v' :: v
v' = Int -> v -> v -> v
f Int
kL v
vL v
vR
                                             SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k (Int -> v -> IPair v
forall a. Int -> a -> IPair a
V.IPair Int
kL v
v')
                                             Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (IPair v)
marr
                                    Ordering
_  -> do SmallMutableArray (PrimState (ST s)) (IPair v)
-> Int -> IPair v -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (IPair v)
SmallMutableArray (PrimState (ST s)) (IPair v)
marr Int
k IPair v
kvR
                                             Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
forall s.
Int -> Int -> Int -> SmallMutableArray s (IPair v) -> ST s Int
go Int
i (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (IPair v)
marr

-- | /O(n)/ Reduce this map by applying a binary operator to all
-- elements, using the given starting value (typically the
-- right-identity of the operator).
--
-- During folding k is in descending order.
foldrWithKey :: (Int -> v -> a -> a) -> a -> FlatIntMap v -> a
{-# INLINE foldrWithKey #-}
foldrWithKey :: (Int -> v -> a -> a) -> a -> FlatIntMap v -> a
foldrWithKey Int -> v -> a -> a
f a
a (FlatIntMap Vector (IPair v)
vs) = (IPair v -> a -> a) -> a -> Vector (IPair v) -> a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (V.IPair Int
k v
v) a
a' -> Int -> v -> a -> a
f Int
k v
v a
a') a
a Vector (IPair v)
vs

-- | /O(n)/ Reduce this map by applying a binary operator to all
-- elements, using the given starting value (typically the
-- right-identity of the operator).
--
-- During folding Int is in ascending order.
foldlWithKey :: (a -> Int -> v -> a) -> a -> FlatIntMap v -> a
{-# INLINE foldlWithKey #-}
foldlWithKey :: (a -> Int -> v -> a) -> a -> FlatIntMap v -> a
foldlWithKey a -> Int -> v -> a
f a
a (FlatIntMap Vector (IPair v)
vs) = (a -> IPair v -> a) -> a -> Vector (IPair v) -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ a
a' (V.IPair Int
k v
v) -> a -> Int -> v -> a
f a
a' Int
k v
v) a
a Vector (IPair v)
vs

-- | /O(n)/ Reduce this map by applying a binary operator to all
-- elements, using the given starting value (typically the
-- right-identity of the operator).
--
-- During folding Int is in descending order.
foldrWithKey' :: (Int -> v -> a -> a) -> a -> FlatIntMap v -> a
{-# INLINE foldrWithKey' #-}
foldrWithKey' :: (Int -> v -> a -> a) -> a -> FlatIntMap v -> a
foldrWithKey' Int -> v -> a -> a
f a
a (FlatIntMap Vector (IPair v)
vs) = (IPair v -> a -> a) -> a -> Vector (IPair v) -> a
forall (v :: * -> *) a b. Vec v a => (a -> b -> b) -> b -> v a -> b
V.foldr' (\ (V.IPair Int
k v
v) -> Int -> v -> a -> a
f Int
k v
v) a
a Vector (IPair v)
vs

-- | /O(n)/ Reduce this map by applying a binary operator to all
-- elements, using the given starting value (typically the
-- right-identity of the operator).
--
-- During folding Int is in ascending order.
foldlWithKey' :: (a -> Int -> v -> a) -> a -> FlatIntMap v -> a
{-# INLINE foldlWithKey' #-}
foldlWithKey' :: (a -> Int -> v -> a) -> a -> FlatIntMap v -> a
foldlWithKey' a -> Int -> v -> a
f a
a (FlatIntMap Vector (IPair v)
vs) = (a -> IPair v -> a) -> a -> Vector (IPair v) -> a
forall (v :: * -> *) a b. Vec v a => (b -> a -> b) -> b -> v a -> b
V.foldl' (\ a
a' (V.IPair Int
k v
v) -> a -> Int -> v -> a
f a
a' Int
k v
v) a
a Vector (IPair v)
vs

-- | /O(n)/.
--
-- @'traverseWithKey' f s == 'pack' \<$\> 'traverse' (\(k, v) -> (,) k \<$\> f k v) ('unpack' m)@
-- That is, behaves exactly like a regular 'traverse' except that the traversing
-- function also has access to the key associated with a value.
traverseWithKey :: Applicative t => (Int -> a -> t b) -> FlatIntMap a -> t (FlatIntMap b)
{-# INLINE traverseWithKey #-}
traverseWithKey :: (Int -> a -> t b) -> FlatIntMap a -> t (FlatIntMap b)
traverseWithKey Int -> a -> t b
f (FlatIntMap Vector (IPair a)
vs) = Vector (IPair b) -> FlatIntMap b
forall v. Vector (IPair v) -> FlatIntMap v
FlatIntMap (Vector (IPair b) -> FlatIntMap b)
-> t (Vector (IPair b)) -> t (FlatIntMap b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (IPair a -> t (IPair b))
-> Vector (IPair a) -> t (Vector (IPair b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\ (V.IPair Int
k a
v) -> Int -> b -> IPair b
forall a. Int -> a -> IPair a
V.IPair Int
k (b -> IPair b) -> t b -> t (IPair b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> t b
f Int
k a
v) Vector (IPair a)
vs

--------------------------------------------------------------------------------

-- | Find the key's index in the vector slice, if key exists return 'Right',
-- otherwise 'Left', i.e. the insert index
--
-- This function only works on ascending sorted vectors.
binarySearch :: V.Vector (V.IPair v) -> Int -> Either Int Int
{-# INLINABLE binarySearch #-}
binarySearch :: Vector (IPair v) -> Int -> Either Int Int
binarySearch (V.Vector SmallArray (IPair v)
_ Int
_ Int
0) Int
_   = Int -> Either Int Int
forall a b. a -> Either a b
Left Int
0
binarySearch (V.Vector SmallArray (IPair v)
arr Int
s0 Int
l) !Int
k' = Int -> Int -> Either Int Int
go Int
s0 (Int
s0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Int -> Either Int Int
go !Int
s !Int
e
        | Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
e =
            let V.IPair Int
k v
_  = SmallArray (IPair v)
arr SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
s
            in case Int
k' Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
k of Ordering
LT -> Int -> Either Int Int
forall a b. a -> Either a b
Left Int
s
                                      Ordering
GT -> let !s' :: Int
s' = Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1 in Int -> Either Int Int
forall a b. a -> Either a b
Left Int
s'
                                      Ordering
_  -> Int -> Either Int Int
forall a b. b -> Either a b
Right Int
s
        | Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>  Int
e = Int -> Either Int Int
forall a b. a -> Either a b
Left Int
s
        | Bool
otherwise =
            let !mid :: Int
mid = (Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
e) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
                (V.IPair Int
k v
_)  = SmallArray (IPair v)
arr SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
mid
            in case Int
k' Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
k of Ordering
LT -> Int -> Int -> Either Int Int
go Int
s (Int
midInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                                      Ordering
GT -> Int -> Int -> Either Int Int
go (Int
midInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
e
                                      Ordering
_  -> Int -> Either Int Int
forall a b. b -> Either a b
Right Int
mid

--------------------------------------------------------------------------------

-- | linear scan search from left to right, return the first one if exist.
linearSearch :: V.Vector (V.IPair v) -> Int -> Maybe v
{-# INLINABLE linearSearch #-}
linearSearch :: Vector (IPair v) -> Int -> Maybe v
linearSearch (V.Vector SmallArray (IPair v)
arr Int
s Int
l) !Int
k' = Int -> Maybe v
go Int
s
  where
    !end :: Int
end = Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l
    go :: Int -> Maybe v
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end = Maybe v
forall a. Maybe a
Nothing
        | Bool
otherwise =
            let V.IPair Int
k v
v  = SmallArray (IPair v)
arr SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
i
            in if Int
k' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k then v -> Maybe v
forall a. a -> Maybe a
Just v
v else Int -> Maybe v
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

-- | linear scan search from right to left, return the first one if exist.
linearSearchR :: V.Vector (V.IPair v) -> Int -> Maybe v
{-# INLINABLE linearSearchR #-}
linearSearchR :: Vector (IPair v) -> Int -> Maybe v
linearSearchR (V.Vector SmallArray (IPair v)
arr Int
s Int
l) !Int
k' = Int -> Maybe v
go (Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Maybe v
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
s = Maybe v
forall a. Maybe a
Nothing
        | Bool
otherwise =
            let V.IPair Int
k v
v  = SmallArray (IPair v)
arr SmallArray (IPair v) -> Int -> IPair v
forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
i
            in if Int
k' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k then v -> Maybe v
forall a. a -> Maybe a
Just v
v else Int -> Maybe v
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)