{-|
Module      : Z.Data.Vector.FlatMap
Description : Fast 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 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.FlatMap
  ( -- * FlatMap backed by sorted vector
    FlatMap, sortedKeyValues, size, null, empty, map', kmap'
  , pack, packN, packR, packRN
  , unpack, unpackR, packVector, packVectorR
  , lookup
  , delete
  , insert
  , adjust'
  , merge, mergeWithKey'
    -- * fold and traverse
  , foldrWithKey, foldrWithKey', foldlWithKey, foldlWithKey', traverseWithKey
    -- * binary search on vectors
  , binarySearch
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.ST
import qualified Data.Primitive.SmallArray  as A
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 Z.Data.Vector.Base         as V
import qualified Z.Data.Vector.Extra        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 FlatMap k v = FlatMap { forall k v. FlatMap k v -> Vector (k, v)
sortedKeyValues :: V.Vector (k, v) }
    deriving (Int -> FlatMap k v -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> FlatMap k v -> ShowS
forall k v. (Show k, Show v) => [FlatMap k v] -> ShowS
forall k v. (Show k, Show v) => FlatMap k v -> String
showList :: [FlatMap k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [FlatMap k v] -> ShowS
show :: FlatMap k v -> String
$cshow :: forall k v. (Show k, Show v) => FlatMap k v -> String
showsPrec :: Int -> FlatMap k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> FlatMap k v -> ShowS
Show, FlatMap k v -> FlatMap k v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => FlatMap k v -> FlatMap k v -> Bool
/= :: FlatMap k v -> FlatMap k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => FlatMap k v -> FlatMap k v -> Bool
== :: FlatMap k v -> FlatMap k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => FlatMap k v -> FlatMap k v -> Bool
Eq, FlatMap k v -> FlatMap k v -> Bool
FlatMap k v -> FlatMap k v -> Ordering
FlatMap k v -> FlatMap k v -> FlatMap k 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 {k} {v}. (Ord k, Ord v) => Eq (FlatMap k v)
forall k v. (Ord k, Ord v) => FlatMap k v -> FlatMap k v -> Bool
forall k v.
(Ord k, Ord v) =>
FlatMap k v -> FlatMap k v -> Ordering
forall k v.
(Ord k, Ord v) =>
FlatMap k v -> FlatMap k v -> FlatMap k v
min :: FlatMap k v -> FlatMap k v -> FlatMap k v
$cmin :: forall k v.
(Ord k, Ord v) =>
FlatMap k v -> FlatMap k v -> FlatMap k v
max :: FlatMap k v -> FlatMap k v -> FlatMap k v
$cmax :: forall k v.
(Ord k, Ord v) =>
FlatMap k v -> FlatMap k v -> FlatMap k v
>= :: FlatMap k v -> FlatMap k v -> Bool
$c>= :: forall k v. (Ord k, Ord v) => FlatMap k v -> FlatMap k v -> Bool
> :: FlatMap k v -> FlatMap k v -> Bool
$c> :: forall k v. (Ord k, Ord v) => FlatMap k v -> FlatMap k v -> Bool
<= :: FlatMap k v -> FlatMap k v -> Bool
$c<= :: forall k v. (Ord k, Ord v) => FlatMap k v -> FlatMap k v -> Bool
< :: FlatMap k v -> FlatMap k v -> Bool
$c< :: forall k v. (Ord k, Ord v) => FlatMap k v -> FlatMap k v -> Bool
compare :: FlatMap k v -> FlatMap k v -> Ordering
$ccompare :: forall k v.
(Ord k, Ord v) =>
FlatMap k v -> FlatMap k v -> Ordering
Ord, Typeable)

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

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

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

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

instance Ord k => Monoid.Monoid (FlatMap k v) where
    {-# INLINE mappend #-}
    mappend :: FlatMap k v -> FlatMap k v -> FlatMap k v
mappend = forall a. Semigroup a => a -> a -> a
(<>)
    {-# INLINE mempty #-}
    mempty :: FlatMap k v
mempty = forall k v. FlatMap k v
empty

instance (NFData k, NFData v) => NFData (FlatMap k v) where
    {-# INLINE rnf #-}
    rnf :: FlatMap k v -> ()
rnf (FlatMap Vector (k, v)
kvs) = forall a. NFData a => a -> ()
rnf Vector (k, v)
kvs

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

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

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

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

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

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

kmap' :: (k -> v -> v') -> FlatMap k v -> FlatMap k v'
{-# INLINE kmap' #-}
kmap' :: forall k v v'. (k -> v -> v') -> FlatMap k v -> FlatMap k v'
kmap' k -> v -> v'
f (FlatMap Vector (k, v)
vs) = forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (u :: * -> *) (v :: * -> *) a b.
(Vec u a, Vec v b) =>
(a -> b) -> u a -> v b
V.map' (\ (k
k, v
v) -> (k
k, k -> v -> v'
f k
k v
v)) Vector (k, v)
vs)

-- | /O(1)/ empty flat map.
empty :: FlatMap k v
{-# NOINLINE empty #-}
empty :: forall k v. FlatMap k v
empty = forall k v. Vector (k, v) -> FlatMap k v
FlatMap 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 :: Ord k => [(k, v)] -> FlatMap k v
{-# INLINABLE pack #-}
pack :: forall k v. Ord k => [(k, v)] -> FlatMap k v
pack [(k, v)]
kvs = forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) (forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (forall a. Ord a => a -> a -> Ordering
compare forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) (forall (v :: * -> *) a. Vec v a => [a] -> v a
V.pack [(k, v)]
kvs)))

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

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

-- | /O(N*logN)/ Pack list of key values with suggested size, on key duplication prefer right one.
packRN :: Ord k => Int -> [(k, v)] -> FlatMap k v
{-# INLINABLE packRN #-}
packRN :: forall k v. Ord k => Int -> [(k, v)] -> FlatMap k v
packRN Int
n [(k, v)]
kvs = forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) (forall (v :: * -> *) a.
Vec v a =>
(a -> a -> Ordering) -> v a -> v a
V.mergeSortBy (forall a. Ord a => a -> a -> Ordering
compare forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) (forall (v :: * -> *) a. Vec v a => Int -> [a] -> v a
V.packN Int
n [(k, 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 :: FlatMap k v -> [(k, v)]
{-# INLINE unpack #-}
unpack :: forall k v. FlatMap k v -> [(k, v)]
unpack = forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. FlatMap k v -> Vector (k, 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 :: FlatMap k v -> [(k, v)]
{-# INLINE unpackR #-}
unpackR :: forall k v. FlatMap k v -> [(k, v)]
unpackR = forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpackR forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. FlatMap k v -> Vector (k, v)
sortedKeyValues

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

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

-- | /O(logN)/ Binary search on flat map.
lookup :: Ord k => k -> FlatMap k v -> Maybe v
{-# INLINABLE lookup #-}
lookup :: forall k v. Ord k => k -> FlatMap k v -> Maybe v
lookup k
_  (FlatMap (V.Vector SmallArray (k, v)
_ Int
_ Int
0)) = forall a. Maybe a
Nothing
lookup k
k' (FlatMap (V.Vector SmallArray (k, v)
arr Int
s Int
l)) = Int -> Int -> Maybe v
go Int
s (Int
sforall a. Num a => a -> a -> a
+Int
lforall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Int -> Maybe v
go !Int
i !Int
j
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
j =
            case SmallArray (k, v)
arr forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
i of (k
k, v
v)  | k
k forall a. Eq a => a -> a -> Bool
== k
k'  -> forall a. a -> Maybe a
Just v
v
                                                      | Bool
otherwise -> forall a. Maybe a
Nothing
        | Int
i forall a. Ord a => a -> a -> Bool
>  Int
j = forall a. Maybe a
Nothing
        | Bool
otherwise =
            let mid :: Int
mid = (Int
iforall a. Num a => a -> a -> a
+Int
j) forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
                (k
k, v
v)  = SmallArray (k, v)
arr forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
mid
            in case k
k' forall a. Ord a => a -> a -> Ordering
`compare` k
k of Ordering
LT -> Int -> Int -> Maybe v
go Int
i (Int
midforall a. Num a => a -> a -> a
-Int
1)
                                      Ordering
GT -> Int -> Int -> Maybe v
go (Int
midforall a. Num a => a -> a -> a
+Int
1) Int
j
                                      Ordering
_  -> forall a. a -> Maybe a
Just v
v

-- | /O(N)/ Insert new key value into map, replace old one if key exists.
insert :: Ord k => k -> v -> FlatMap k v -> FlatMap k v
{-# INLINABLE insert #-}
insert :: forall k v. Ord k => k -> v -> FlatMap k v -> FlatMap k v
insert k
k v
v (FlatMap Vector (k, v)
vec) =
    case forall k v. Ord k => Vector (k, v) -> k -> Either Int Int
binarySearch Vector (k, v)
vec k
k of
        Left Int
i -> forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
v a -> Int -> a -> v a
V.unsafeInsertIndex Vector (k, v)
vec Int
i (k
k, v
v))
        Right Int
i -> forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
v a -> Int -> (a -> a) -> v a
V.unsafeModifyIndex Vector (k, v)
vec Int
i (forall a b. a -> b -> a
const (k
k, v
v)))

-- | /O(N)/ Delete a key value pair by key.
delete :: Ord k => k -> FlatMap k v -> FlatMap k v
{-# INLINABLE delete #-}
delete :: forall k v. Ord k => k -> FlatMap k v -> FlatMap k v
delete k
k m :: FlatMap k v
m@(FlatMap Vector (k, v)
vec) =
    case forall k v. Ord k => Vector (k, v) -> k -> Either Int Int
binarySearch Vector (k, v)
vec k
k of
        Left Int
_ -> FlatMap k v
m
        Right Int
i -> forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
v a -> Int -> v a
V.unsafeDeleteIndex Vector (k, v)
vec Int
i)

-- | /O(N)/ Modify a value by key.
--
-- The value is evaluated to WHNF before writing into map.
adjust' :: Ord k => (v -> v) -> k -> FlatMap k v -> FlatMap k v
{-# INLINABLE adjust' #-}
adjust' :: forall k v. Ord k => (v -> v) -> k -> FlatMap k v -> FlatMap k v
adjust' v -> v
f k
k m :: FlatMap k v
m@(FlatMap Vector (k, v)
vec) =
    case forall k v. Ord k => Vector (k, v) -> k -> Either Int Int
binarySearch Vector (k, v)
vec k
k of
        Left Int
_ -> FlatMap k v
m
        Right Int
i -> forall k v. Vector (k, v) -> FlatMap k v
FlatMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
v a -> Int -> (a -> a) -> v a
V.unsafeModifyIndex Vector (k, v)
vec Int
i forall a b. (a -> b) -> a -> b
$
            \ (k
k', v
v) -> let !v' :: v
v' = v -> v
f v
v in (k
k', v
v')

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

-- | /O(n+m)/ Merge two 'FlatMap' with a merge function.
mergeWithKey' :: forall k v. Ord k => (k -> v -> v -> v) -> FlatMap k v -> FlatMap k v -> FlatMap k v
{-# INLINABLE mergeWithKey' #-}
mergeWithKey' :: forall k v.
Ord k =>
(k -> v -> v -> v) -> FlatMap k v -> FlatMap k v -> FlatMap k v
mergeWithKey' k -> v -> v -> v
f fmL :: FlatMap k v
fmL@(FlatMap (V.Vector SmallArray (k, v)
arrL Int
sL Int
lL)) fmR :: FlatMap k v
fmR@(FlatMap (V.Vector SmallArray (k, v)
arrR Int
sR Int
lR))
    | forall k a. FlatMap k a -> Bool
null FlatMap k v
fmL = FlatMap k v
fmR
    | forall k a. FlatMap k a -> Bool
null FlatMap k v
fmR = FlatMap k v
fmL
    | Bool
otherwise = forall k v. Vector (k, v) -> FlatMap k v
FlatMap (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
Int -> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
V.createN (Int
lLforall a. Num a => a -> a -> a
+Int
lR) (forall s.
Int -> Int -> Int -> SmallMutableArray s (k, v) -> ST s Int
go Int
sL Int
sR Int
0))
  where
    endL :: Int
endL = Int
sL forall a. Num a => a -> a -> a
+ Int
lL
    endR :: Int
endR = Int
sR forall a. Num a => a -> a -> a
+ Int
lR
    go :: Int -> Int -> Int -> A.SmallMutableArray s (k, v) -> ST s Int
    go :: forall s.
Int -> Int -> Int -> SmallMutableArray s (k, v) -> ST s Int
go !Int
i !Int
j !Int
k SmallMutableArray s (k, v)
marr
        | Int
i forall a. Ord a => a -> a -> Bool
>= Int
endL = do
            forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s (k, v)
marr Int
k SmallArray (k, v)
arrR Int
j (Int
lRforall a. Num a => a -> a -> a
-Int
j)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Int
kforall a. Num a => a -> a -> a
+Int
lRforall a. Num a => a -> a -> a
-Int
j
        | Int
j forall a. Ord a => a -> a -> Bool
>= Int
endR = do
            forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s (k, v)
marr Int
k SmallArray (k, v)
arrL Int
i (Int
lLforall a. Num a => a -> a -> a
-Int
i)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Int
kforall a. Num a => a -> a -> a
+Int
lLforall a. Num a => a -> a -> a
-Int
i
        | Bool
otherwise = do
            kvL :: (k, v)
kvL@(k
kL, v
vL) <- SmallArray (k, v)
arrL forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
i
            kvR :: (k, v)
kvR@(k
kR, v
vR) <- SmallArray (k, v)
arrR forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
j
            case k
kL forall a. Ord a => a -> a -> Ordering
`compare` k
kR of Ordering
LT -> do forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (k, v)
marr Int
k (k, v)
kvL
                                             forall s.
Int -> Int -> Int -> SmallMutableArray s (k, v) -> ST s Int
go (Int
iforall a. Num a => a -> a -> a
+Int
1) Int
j (Int
kforall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (k, v)
marr
                                    Ordering
EQ -> do let !v' :: v
v' = k -> v -> v -> v
f k
kL v
vL v
vR
                                             forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (k, v)
marr Int
k (k
kL, v
v')
                                             forall s.
Int -> Int -> Int -> SmallMutableArray s (k, v) -> ST s Int
go (Int
iforall a. Num a => a -> a -> a
+Int
1) (Int
jforall a. Num a => a -> a -> a
+Int
1) (Int
kforall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (k, v)
marr
                                    Ordering
_  -> do forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s (k, v)
marr Int
k (k, v)
kvR
                                             forall s.
Int -> Int -> Int -> SmallMutableArray s (k, v) -> ST s Int
go Int
i (Int
jforall a. Num a => a -> a -> a
+Int
1) (Int
kforall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s (k, 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 :: (k -> v -> a -> a) -> a -> FlatMap k v -> a
{-# INLINE foldrWithKey #-}
foldrWithKey :: forall k v a. (k -> v -> a -> a) -> a -> FlatMap k v -> a
foldrWithKey k -> v -> a -> a
f a
a (FlatMap Vector (k, v)
vs) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> v -> a -> a
f) a
a Vector (k, 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 k is in ascending order.
foldlWithKey :: (a -> k -> v -> a) -> a -> FlatMap k v -> a
{-# INLINE foldlWithKey #-}
foldlWithKey :: forall a k v. (a -> k -> v -> a) -> a -> FlatMap k v -> a
foldlWithKey a -> k -> v -> a
f a
a (FlatMap Vector (k, v)
vs) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ a
a' (k
k,v
v) -> a -> k -> v -> a
f a
a' k
k v
v) a
a Vector (k, 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 k is in descending order.
foldrWithKey' :: (k -> v -> a -> a) -> a -> FlatMap k v -> a
{-# INLINE foldrWithKey' #-}
foldrWithKey' :: forall k v a. (k -> v -> a -> a) -> a -> FlatMap k v -> a
foldrWithKey' k -> v -> a -> a
f a
a (FlatMap Vector (k, v)
vs) = forall (v :: * -> *) a b. Vec v a => (a -> b -> b) -> b -> v a -> b
V.foldr' (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> v -> a -> a
f) a
a Vector (k, 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 k is in ascending order.
foldlWithKey' :: (a -> k -> v -> a) -> a -> FlatMap k v -> a
{-# INLINE foldlWithKey' #-}
foldlWithKey' :: forall a k v. (a -> k -> v -> a) -> a -> FlatMap k v -> a
foldlWithKey' a -> k -> v -> a
f a
a (FlatMap Vector (k, v)
vs) = forall (v :: * -> *) a b. Vec v a => (b -> a -> b) -> b -> v a -> b
V.foldl' (\ a
a' (k
k,v
v) -> a -> k -> v -> a
f a
a' k
k v
v) a
a Vector (k, 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 => (k -> a -> t b) -> FlatMap k a -> t (FlatMap k b)
{-# INLINE traverseWithKey #-}
traverseWithKey :: forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> FlatMap k a -> t (FlatMap k b)
traverseWithKey k -> a -> t b
f (FlatMap Vector (k, a)
vs) = forall k v. Vector (k, v) -> FlatMap k v
FlatMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> *) a (u :: * -> *) b (f :: * -> *).
(Vec v a, Vec u b, Applicative f) =>
(a -> f b) -> v a -> f (u b)
V.traverse (\ (k
k,a
v) -> (k
k,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> a -> t b
f k
k a
v) Vector (k, 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 :: Ord k => V.Vector (k, v) -> k -> Either Int Int
{-# INLINABLE binarySearch #-}
binarySearch :: forall k v. Ord k => Vector (k, v) -> k -> Either Int Int
binarySearch (V.Vector SmallArray (k, v)
_ Int
_ Int
0) k
_   = forall a b. a -> Either a b
Left Int
0
binarySearch (V.Vector SmallArray (k, v)
arr Int
s Int
l) !k
k' = Int -> Int -> Either Int Int
go Int
s (Int
sforall a. Num a => a -> a -> a
+Int
lforall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Int -> Either Int Int
go !Int
i !Int
j
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
j =
            let (k
k, v
_)  = SmallArray (k, v)
arr forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
i
            in case k
k' forall a. Ord a => a -> a -> Ordering
`compare` k
k of Ordering
LT -> forall a b. a -> Either a b
Left Int
i
                                      Ordering
GT -> let !i' :: Int
i' = Int
iforall a. Num a => a -> a -> a
+Int
1 in forall a b. a -> Either a b
Left Int
i'
                                      Ordering
_  -> forall a b. b -> Either a b
Right Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
>  Int
j = forall a b. a -> Either a b
Left Int
i
        | Bool
otherwise =
            let !mid :: Int
mid = (Int
iforall a. Num a => a -> a -> a
+Int
j) forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
                (k
k, v
_)  = SmallArray (k, v)
arr forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
mid
            in case k
k' forall a. Ord a => a -> a -> Ordering
`compare` k
k of Ordering
LT -> Int -> Int -> Either Int Int
go Int
i (Int
midforall a. Num a => a -> a -> a
-Int
1)
                                      Ordering
GT -> Int -> Int -> Either Int Int
go (Int
midforall a. Num a => a -> a -> a
+Int
1) Int
j
                                      Ordering
_  -> forall a b. b -> Either a b
Right Int
mid