module Z.Data.Vector.FlatMap
(
FlatMap, sortedKeyValues, size, null, empty, map', kmap'
, pack, packN, packR, packRN
, unpack, unpackR, packVector, packVectorR
, lookup
, delete
, insert
, adjust'
, merge, mergeWithKey'
, foldrWithKey, foldrWithKey', foldlWithKey, foldlWithKey', traverseWithKey
, 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)
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
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)))
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)))
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)))
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)))
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
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
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))
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))
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
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)))
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)
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')
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
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
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
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
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
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
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
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