{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
module Data.IMap
    ( IMap
    , Run(..)
    , empty
    , Data.IMap.null
    , singleton
    , insert
    , delete
    , restrict
    , lookup
    , splitLE
    , intersectionWith
    , mapMaybe
    , addToKeys
    , unsafeUnion
    , fromList
    , unsafeRuns
    , unsafeToAscList
    ) where

import Data.List (foldl')
import Data.Monoid
import Data.IntMap.Strict (IntMap)
import GHC.Generics
import Control.DeepSeq
import Prelude hiding (lookup)
import qualified Data.IntMap.Strict as IM

-- | Semantically, 'IMap' and 'IntMap' are identical; but 'IMap' is more
-- efficient when large sequences of contiguous keys are mapped to the same
-- value.
newtype IMap a = IMap { forall a. IMap a -> IntMap (Run a)
_runs :: IntMap (Run a) } deriving (Int -> IMap a -> ShowS
[IMap a] -> ShowS
IMap a -> String
(Int -> IMap a -> ShowS)
-> (IMap a -> String) -> ([IMap a] -> ShowS) -> Show (IMap a)
forall a. Show a => Int -> IMap a -> ShowS
forall a. Show a => [IMap a] -> ShowS
forall a. Show a => IMap a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> IMap a -> ShowS
showsPrec :: Int -> IMap a -> ShowS
$cshow :: forall a. Show a => IMap a -> String
show :: IMap a -> String
$cshowList :: forall a. Show a => [IMap a] -> ShowS
showList :: [IMap a] -> ShowS
Show, (forall a b. (a -> b) -> IMap a -> IMap b)
-> (forall a b. a -> IMap b -> IMap a) -> Functor IMap
forall a b. a -> IMap b -> IMap a
forall a b. (a -> b) -> IMap a -> IMap b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> IMap a -> IMap b
fmap :: forall a b. (a -> b) -> IMap a -> IMap b
$c<$ :: forall a b. a -> IMap b -> IMap a
<$ :: forall a b. a -> IMap b -> IMap a
Functor, ReadPrec [IMap a]
ReadPrec (IMap a)
Int -> ReadS (IMap a)
ReadS [IMap a]
(Int -> ReadS (IMap a))
-> ReadS [IMap a]
-> ReadPrec (IMap a)
-> ReadPrec [IMap a]
-> Read (IMap a)
forall a. Read a => ReadPrec [IMap a]
forall a. Read a => ReadPrec (IMap a)
forall a. Read a => Int -> ReadS (IMap a)
forall a. Read a => ReadS [IMap a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: forall a. Read a => Int -> ReadS (IMap a)
readsPrec :: Int -> ReadS (IMap a)
$creadList :: forall a. Read a => ReadS [IMap a]
readList :: ReadS [IMap a]
$creadPrec :: forall a. Read a => ReadPrec (IMap a)
readPrec :: ReadPrec (IMap a)
$creadListPrec :: forall a. Read a => ReadPrec [IMap a]
readListPrec :: ReadPrec [IMap a]
Read, (forall x. IMap a -> Rep (IMap a) x)
-> (forall x. Rep (IMap a) x -> IMap a) -> Generic (IMap a)
forall x. Rep (IMap a) x -> IMap a
forall x. IMap a -> Rep (IMap a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (IMap a) x -> IMap a
forall a x. IMap a -> Rep (IMap a) x
$cfrom :: forall a x. IMap a -> Rep (IMap a) x
from :: forall x. IMap a -> Rep (IMap a) x
$cto :: forall a x. Rep (IMap a) x -> IMap a
to :: forall x. Rep (IMap a) x -> IMap a
Generic, IMap a -> ()
(IMap a -> ()) -> NFData (IMap a)
forall a. NFData a => IMap a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. NFData a => IMap a -> ()
rnf :: IMap a -> ()
NFData)

{-# INLINE unsafeRuns #-}
-- | This function is unsafe because 'IMap's that compare equal may split their
-- runs into different chunks; consumers must promise that they do not treat
-- run boundaries specially.
unsafeRuns :: IMap a -> IntMap (Run a)
unsafeRuns :: forall a. IMap a -> IntMap (Run a)
unsafeRuns = IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs

instance Eq a => Eq (IMap a) where
    IMap IntMap (Run a)
m == :: IMap a -> IMap a -> Bool
== IMap IntMap (Run a)
m' = [(Int, Run a)] -> [(Int, Run a)] -> Bool
forall {a}. Eq a => [(Int, Run a)] -> [(Int, Run a)] -> Bool
go (IntMap (Run a) -> [(Int, Run a)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap (Run a)
m) (IntMap (Run a) -> [(Int, Run a)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap (Run a)
m') where
        go :: [(Int, Run a)] -> [(Int, Run a)] -> Bool
go ((Int
k, Run Int
n a
a):[(Int, Run a)]
kvs) ((Int
k', Run Int
n' a
a'):[(Int, Run a)]
kvs')
            = Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k' Bool -> Bool -> Bool
&& a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a' Bool -> Bool -> Bool
&& case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
n Int
n' of
                Ordering
LT -> [(Int, Run a)] -> [(Int, Run a)] -> Bool
go [(Int, Run a)]
kvs ((Int
k'Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n, Int -> a -> Run a
forall a. Int -> a -> Run a
Run (Int
n'Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n) a
a')(Int, Run a) -> [(Int, Run a)] -> [(Int, Run a)]
forall a. a -> [a] -> [a]
:[(Int, Run a)]
kvs')
                Ordering
EQ -> [(Int, Run a)] -> [(Int, Run a)] -> Bool
go [(Int, Run a)]
kvs [(Int, Run a)]
kvs'
                Ordering
GT -> [(Int, Run a)] -> [(Int, Run a)] -> Bool
go ((Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n', Int -> a -> Run a
forall a. Int -> a -> Run a
Run (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n') a
a)(Int, Run a) -> [(Int, Run a)] -> [(Int, Run a)]
forall a. a -> [a] -> [a]
:[(Int, Run a)]
kvs) [(Int, Run a)]
kvs'
        go [] [] = Bool
True
        go [(Int, Run a)]
_ [(Int, Run a)]
_ = Bool
False

instance Ord a => Ord (IMap a) where
    compare :: IMap a -> IMap a -> Ordering
compare (IMap IntMap (Run a)
m) (IMap IntMap (Run a)
m') = [(Int, Run a)] -> [(Int, Run a)] -> Ordering
forall {a}. Ord a => [(Int, Run a)] -> [(Int, Run a)] -> Ordering
go (IntMap (Run a) -> [(Int, Run a)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap (Run a)
m) (IntMap (Run a) -> [(Int, Run a)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap (Run a)
m') where
        go :: [(Int, Run a)] -> [(Int, Run a)] -> Ordering
go [] [] = Ordering
EQ
        go [] [(Int, Run a)]
_  = Ordering
LT
        go [(Int, Run a)]
_  [] = Ordering
GT
        go ((Int
k, Run Int
n a
a):[(Int, Run a)]
kvs) ((Int
k', Run Int
n' a
a'):[(Int, Run a)]
kvs')
            = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
k Int
k' Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
a a
a' Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
n Int
n' of
                Ordering
LT -> [(Int, Run a)] -> [(Int, Run a)] -> Ordering
go [(Int, Run a)]
kvs ((Int
k'Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n, Int -> a -> Run a
forall a. Int -> a -> Run a
Run (Int
n'Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n) a
a')(Int, Run a) -> [(Int, Run a)] -> [(Int, Run a)]
forall a. a -> [a] -> [a]
:[(Int, Run a)]
kvs')
                Ordering
EQ -> [(Int, Run a)] -> [(Int, Run a)] -> Ordering
go [(Int, Run a)]
kvs [(Int, Run a)]
kvs'
                Ordering
GT -> [(Int, Run a)] -> [(Int, Run a)] -> Ordering
go ((Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n', Int -> a -> Run a
forall a. Int -> a -> Run a
Run (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n') a
a)(Int, Run a) -> [(Int, Run a)] -> [(Int, Run a)]
forall a. a -> [a] -> [a]
:[(Int, Run a)]
kvs) [(Int, Run a)]
kvs'

-- | Zippy: '(<*>)' combines values at equal keys, discarding any values whose
-- key is in only one of its two arguments.
instance Applicative IMap where
    pure :: forall a. a -> IMap a
pure a
a = IntMap (Run a) -> IMap a
forall a. IntMap (Run a) -> IMap a
IMap (IntMap (Run a) -> IMap a)
-> ([(Int, Run a)] -> IntMap (Run a)) -> [(Int, Run a)] -> IMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Run a)] -> IntMap (Run a)
forall a. [(Int, a)] -> IntMap a
IM.fromDistinctAscList ([(Int, Run a)] -> IMap a) -> [(Int, Run a)] -> IMap a
forall a b. (a -> b) -> a -> b
$
        [ (Int
forall a. Bounded a => a
minBound, Int -> a -> Run a
forall a. Int -> a -> Run a
Run Int
forall a. Bounded a => a
maxBound a
a)
        , (-Int
1, Int -> a -> Run a
forall a. Int -> a -> Run a
Run Int
forall a. Bounded a => a
maxBound a
a)
        , (Int
forall a. Bounded a => a
maxBoundInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, Int -> a -> Run a
forall a. Int -> a -> Run a
Run Int
2 a
a)
        ]
    <*> :: forall a b. IMap (a -> b) -> IMap a -> IMap b
(<*>) = ((a -> b) -> a -> b) -> IMap (a -> b) -> IMap a -> IMap b
forall a b c. (a -> b -> c) -> IMap a -> IMap b -> IMap c
intersectionWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)

-- | @Run n a@ represents @n@ copies of the value @a@.
data Run a = Run
    { forall a. Run a -> Int
len :: !Int
    , forall a. Run a -> a
val :: !a
    } deriving (Run a -> Run a -> Bool
(Run a -> Run a -> Bool) -> (Run a -> Run a -> Bool) -> Eq (Run a)
forall a. Eq a => Run a -> Run a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Run a -> Run a -> Bool
== :: Run a -> Run a -> Bool
$c/= :: forall a. Eq a => Run a -> Run a -> Bool
/= :: Run a -> Run a -> Bool
Eq, Eq (Run a)
Eq (Run a) =>
(Run a -> Run a -> Ordering)
-> (Run a -> Run a -> Bool)
-> (Run a -> Run a -> Bool)
-> (Run a -> Run a -> Bool)
-> (Run a -> Run a -> Bool)
-> (Run a -> Run a -> Run a)
-> (Run a -> Run a -> Run a)
-> Ord (Run a)
Run a -> Run a -> Bool
Run a -> Run a -> Ordering
Run a -> Run a -> Run a
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 a. Ord a => Eq (Run a)
forall a. Ord a => Run a -> Run a -> Bool
forall a. Ord a => Run a -> Run a -> Ordering
forall a. Ord a => Run a -> Run a -> Run a
$ccompare :: forall a. Ord a => Run a -> Run a -> Ordering
compare :: Run a -> Run a -> Ordering
$c< :: forall a. Ord a => Run a -> Run a -> Bool
< :: Run a -> Run a -> Bool
$c<= :: forall a. Ord a => Run a -> Run a -> Bool
<= :: Run a -> Run a -> Bool
$c> :: forall a. Ord a => Run a -> Run a -> Bool
> :: Run a -> Run a -> Bool
$c>= :: forall a. Ord a => Run a -> Run a -> Bool
>= :: Run a -> Run a -> Bool
$cmax :: forall a. Ord a => Run a -> Run a -> Run a
max :: Run a -> Run a -> Run a
$cmin :: forall a. Ord a => Run a -> Run a -> Run a
min :: Run a -> Run a -> Run a
Ord, ReadPrec [Run a]
ReadPrec (Run a)
Int -> ReadS (Run a)
ReadS [Run a]
(Int -> ReadS (Run a))
-> ReadS [Run a]
-> ReadPrec (Run a)
-> ReadPrec [Run a]
-> Read (Run a)
forall a. Read a => ReadPrec [Run a]
forall a. Read a => ReadPrec (Run a)
forall a. Read a => Int -> ReadS (Run a)
forall a. Read a => ReadS [Run a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: forall a. Read a => Int -> ReadS (Run a)
readsPrec :: Int -> ReadS (Run a)
$creadList :: forall a. Read a => ReadS [Run a]
readList :: ReadS [Run a]
$creadPrec :: forall a. Read a => ReadPrec (Run a)
readPrec :: ReadPrec (Run a)
$creadListPrec :: forall a. Read a => ReadPrec [Run a]
readListPrec :: ReadPrec [Run a]
Read, Int -> Run a -> ShowS
[Run a] -> ShowS
Run a -> String
(Int -> Run a -> ShowS)
-> (Run a -> String) -> ([Run a] -> ShowS) -> Show (Run a)
forall a. Show a => Int -> Run a -> ShowS
forall a. Show a => [Run a] -> ShowS
forall a. Show a => Run a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Run a -> ShowS
showsPrec :: Int -> Run a -> ShowS
$cshow :: forall a. Show a => Run a -> String
show :: Run a -> String
$cshowList :: forall a. Show a => [Run a] -> ShowS
showList :: [Run a] -> ShowS
Show, (forall a b. (a -> b) -> Run a -> Run b)
-> (forall a b. a -> Run b -> Run a) -> Functor Run
forall a b. a -> Run b -> Run a
forall a b. (a -> b) -> Run a -> Run b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Run a -> Run b
fmap :: forall a b. (a -> b) -> Run a -> Run b
$c<$ :: forall a b. a -> Run b -> Run a
<$ :: forall a b. a -> Run b -> Run a
Functor, (forall x. Run a -> Rep (Run a) x)
-> (forall x. Rep (Run a) x -> Run a) -> Generic (Run a)
forall x. Rep (Run a) x -> Run a
forall x. Run a -> Rep (Run a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Run a) x -> Run a
forall a x. Run a -> Rep (Run a) x
$cfrom :: forall a x. Run a -> Rep (Run a) x
from :: forall x. Run a -> Rep (Run a) x
$cto :: forall a x. Rep (Run a) x -> Run a
to :: forall x. Rep (Run a) x -> Run a
Generic, Run a -> ()
(Run a -> ()) -> NFData (Run a)
forall a. NFData a => Run a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. NFData a => Run a -> ()
rnf :: Run a -> ()
NFData)

instance Foldable    Run where foldMap :: forall m a. Monoid m => (a -> m) -> Run a -> m
foldMap a -> m
f Run a
r = a -> m
f (Run a -> a
forall a. Run a -> a
val Run a
r)
instance Traversable Run where sequenceA :: forall (f :: * -> *) a. Applicative f => Run (f a) -> f (Run a)
sequenceA (Run Int
n f a
v) = Int -> a -> Run a
forall a. Int -> a -> Run a
Run Int
n (a -> Run a) -> f a -> f (Run a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
v

empty :: IMap a
empty :: forall a. IMap a
empty = IntMap (Run a) -> IMap a
forall a. IntMap (Run a) -> IMap a
IMap IntMap (Run a)
forall a. IntMap a
IM.empty

null :: IMap a -> Bool
null :: forall a. IMap a -> Bool
null = IntMap (Run a) -> Bool
forall a. IntMap a -> Bool
IM.null (IntMap (Run a) -> Bool)
-> (IMap a -> IntMap (Run a)) -> IMap a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs

singleton :: Int -> Run a -> IMap a
singleton :: forall a. Int -> Run a -> IMap a
singleton Int
k Run a
r
    | Run a -> Int
forall a. Run a -> Int
len Run a
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 = IntMap (Run a) -> IMap a
forall a. IntMap (Run a) -> IMap a
IMap (Int -> Run a -> IntMap (Run a)
forall a. Int -> a -> IntMap a
IM.singleton Int
k Run a
r)
    | Bool
otherwise = IMap a
forall a. IMap a
empty

insert :: Int -> Run a -> IMap a -> IMap a
insert :: forall a. Int -> Run a -> IMap a -> IMap a
insert Int
k Run a
r IMap a
m
    | Run a -> Int
forall a. Run a -> Int
len Run a
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = IMap a
m
    | Bool
otherwise = IMap a
m { _runs = IM.insert k r (_runs (delete k r m)) }

{-# INLINE delete #-}
delete :: Int -> Run ignored -> IMap a -> IMap a
delete :: forall ignored a. Int -> Run ignored -> IMap a -> IMap a
delete Int
k Run ignored
r IMap a
m
    | Run ignored -> Int
forall a. Run a -> Int
len Run ignored
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = IMap a
m
    | Bool
otherwise = IMap a
m { _runs = IM.union (_runs lt) (_runs gt) }
    where
    (IMap a
lt, IMap a
ge) = Int -> IMap a -> (IMap a, IMap a)
forall a. Int -> IMap a -> (IMap a, IMap a)
splitLE (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) IMap a
m
    (IMap a
_ , IMap a
gt) = Int -> IMap a -> (IMap a, IMap a)
forall a. Int -> IMap a -> (IMap a, IMap a)
splitLE (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Run ignored -> Int
forall a. Run a -> Int
len Run ignored
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) IMap a
ge

-- | Given a range of keys (as specified by a starting key and a length for
-- consistency with other functions in this module), restrict the map to keys
-- in that range. @restrict k r m@ is equivalent to @intersectionWith const m
-- (insert k r empty)@ but potentially more efficient.
restrict :: Int -> Run ignored -> IMap a -> IMap a
restrict :: forall ignored a. Int -> Run ignored -> IMap a -> IMap a
restrict Int
k Run ignored
r = IMap a -> IMap a
forall a. a -> a
id
    (IMap a -> IMap a) -> (IMap a -> IMap a) -> IMap a -> IMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IMap a, IMap a) -> IMap a
forall a b. (a, b) -> b
snd
    ((IMap a, IMap a) -> IMap a)
-> (IMap a -> (IMap a, IMap a)) -> IMap a -> IMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IMap a -> (IMap a, IMap a)
forall a. Int -> IMap a -> (IMap a, IMap a)
splitLE (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
    (IMap a -> (IMap a, IMap a))
-> (IMap a -> IMap a) -> IMap a -> (IMap a, IMap a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IMap a, IMap a) -> IMap a
forall a b. (a, b) -> a
fst
    ((IMap a, IMap a) -> IMap a)
-> (IMap a -> (IMap a, IMap a)) -> IMap a -> IMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IMap a -> (IMap a, IMap a)
forall a. Int -> IMap a -> (IMap a, IMap a)
splitLE (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Run ignored -> Int
forall a. Run a -> Int
len Run ignored
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

lookup :: Int -> IMap a -> Maybe a
lookup :: forall a. Int -> IMap a -> Maybe a
lookup Int
k IMap a
m = case Int -> IntMap (Run a) -> Maybe (Int, Run a)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
k (IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs IMap a
m) of
    Just (Int
k', Run Int
n a
a) | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k'Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n -> a -> Maybe a
forall a. a -> Maybe a
Just a
a
    Maybe (Int, Run a)
_ -> Maybe a
forall a. Maybe a
Nothing

-- | @splitLE n m@ produces a tuple @(le, gt)@ where @le@ has all the
-- associations of @m@ where the keys are @<= n@ and @gt@ has all the
-- associations of @m@ where the keys are @> n@.
splitLE :: Int -> IMap a -> (IMap a, IMap a)
splitLE :: forall a. Int -> IMap a -> (IMap a, IMap a)
splitLE Int
k IMap a
m = case Int -> IntMap (Run a) -> Maybe (Int, Run a)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
k (IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs IMap a
m) of
    Maybe (Int, Run a)
Nothing -> (IMap a
forall a. IMap a
empty, IMap a
m)
    Just (Int
k', r :: Run a
r@(Run Int
n a
_)) -> case (Int
k' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
k, Int
k' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k) of
        (Bool
True , Bool
False) -> (IMap a
m { _runs = lt }, IMap a
m { _runs = gt })
        (Bool
True , Bool
True ) -> (IMap a
m { _runs = IM.insert k r lt }, IMap a
m { _runs = gt })
        (Bool
False, Bool
_    ) -> ( IMap a
m { _runs = IM.insert k'    r { len =     1 + k - k' } lt' }
                          , IMap a
m { _runs = IM.insert (k+1) r { len = n - 1 - k + k' } gt' }
                          )
        where
        (IntMap (Run a)
lt', IntMap (Run a)
gt') = Int -> IntMap (Run a) -> (IntMap (Run a), IntMap (Run a))
forall a. Int -> IntMap a -> (IntMap a, IntMap a)
IM.split Int
k' (IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs IMap a
m)
    where
    (IntMap (Run a)
lt, IntMap (Run a)
gt) = Int -> IntMap (Run a) -> (IntMap (Run a), IntMap (Run a))
forall a. Int -> IntMap a -> (IntMap a, IntMap a)
IM.split Int
k (IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs IMap a
m)

-- | Increment all keys by the given amount. This is like
-- 'IM.mapKeysMonotonic', but restricted to partially-applied addition.
addToKeys :: Int -> IMap a -> IMap a
addToKeys :: forall a. Int -> IMap a -> IMap a
addToKeys Int
n IMap a
m = IMap a
m { _runs = IM.mapKeysMonotonic (n+) (_runs m) }

-- TODO: This is pretty inefficient. IntMap offers some splitting functions
-- that should make it possible to be more efficient here (though the
-- implementation would be significantly messier).
intersectionWith :: (a -> b -> c) -> IMap a -> IMap b -> IMap c
intersectionWith :: forall a b c. (a -> b -> c) -> IMap a -> IMap b -> IMap c
intersectionWith a -> b -> c
f (IMap IntMap (Run a)
runsa) (IMap IntMap (Run b)
runsb)
    = IntMap (Run c) -> IMap c
forall a. IntMap (Run a) -> IMap a
IMap (IntMap (Run c) -> IMap c)
-> ([(Int, Run c)] -> IntMap (Run c)) -> [(Int, Run c)] -> IMap c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Run c)] -> IntMap (Run c)
forall a. [(Int, a)] -> IntMap a
IM.fromDistinctAscList ([(Int, Run c)] -> IMap c) -> [(Int, Run c)] -> IMap c
forall a b. (a -> b) -> a -> b
$ [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge (IntMap (Run a) -> [(Int, Run a)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap (Run a)
runsa) (IntMap (Run b) -> [(Int, Run b)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap (Run b)
runsb)
    where
    merge :: [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge as :: [(Int, Run a)]
as@((Int
ka, Run a
ra):[(Int, Run a)]
at) bs :: [(Int, Run b)]
bs@((Int
kb, Run b
rb):[(Int, Run b)]
bt)
        | Int
ka' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
kb = [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge [(Int, Run a)]
at [(Int, Run b)]
bs
        | Int
kb' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ka = [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge [(Int, Run a)]
as [(Int, Run b)]
bt
        | Bool
otherwise = (Int
kc, Int -> c -> Run c
forall a. Int -> a -> Run a
Run (Int
kc' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
kc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) c
vc) (Int, Run c) -> [(Int, Run c)] -> [(Int, Run c)]
forall a. a -> [a] -> [a]
: case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
ka' Int
kb' of
            Ordering
LT -> [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge [(Int, Run a)]
at [(Int, Run b)]
bs
            Ordering
EQ -> [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge [(Int, Run a)]
at [(Int, Run b)]
bt
            Ordering
GT -> [(Int, Run a)] -> [(Int, Run b)] -> [(Int, Run c)]
merge [(Int, Run a)]
as [(Int, Run b)]
bt
        where
        ka' :: Int
ka' = Int
ka Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Run a -> Int
forall a. Run a -> Int
len Run a
ra Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        kb' :: Int
kb' = Int
kb Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Run b -> Int
forall a. Run a -> Int
len Run b
rb Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        kc :: Int
kc  = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
ka  Int
kb
        kc' :: Int
kc' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
ka' Int
kb'
        vc :: c
vc  = a -> b -> c
f (Run a -> a
forall a. Run a -> a
val Run a
ra) (Run b -> b
forall a. Run a -> a
val Run b
rb)
    merge [(Int, Run a)]
_ [(Int, Run b)]
_ = []

mapMaybe :: (a -> Maybe b) -> IMap a -> IMap b
mapMaybe :: forall a b. (a -> Maybe b) -> IMap a -> IMap b
mapMaybe a -> Maybe b
f (IMap IntMap (Run a)
runs) = IntMap (Run b) -> IMap b
forall a. IntMap (Run a) -> IMap a
IMap ((Run a -> Maybe (Run b)) -> IntMap (Run a) -> IntMap (Run b)
forall a b. (a -> Maybe b) -> IntMap a -> IntMap b
IM.mapMaybe ((a -> Maybe b) -> Run a -> Maybe (Run b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Run a -> f (Run b)
traverse a -> Maybe b
f) IntMap (Run a)
runs)

fromList :: [(Int, Run a)] -> IMap a
fromList :: forall a. [(Int, Run a)] -> IMap a
fromList = (IMap a -> (Int, Run a) -> IMap a)
-> IMap a -> [(Int, Run a)] -> IMap a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IMap a
m (Int
k, Run a
r) -> Int -> Run a -> IMap a -> IMap a
forall a. Int -> Run a -> IMap a -> IMap a
insert Int
k Run a
r IMap a
m) IMap a
forall a. IMap a
empty

-- | This function is unsafe because 'IMap's that compare equal may split their
-- runs into different chunks; consumers must promise that they do not treat
-- run boundaries specially.
unsafeToAscList :: IMap a -> [(Int, Run a)]
unsafeToAscList :: forall a. IMap a -> [(Int, Run a)]
unsafeToAscList = IntMap (Run a) -> [(Int, Run a)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList (IntMap (Run a) -> [(Int, Run a)])
-> (IMap a -> IntMap (Run a)) -> IMap a -> [(Int, Run a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs

-- | This function is unsafe because it assumes there is no overlap between its
-- arguments. That is, in the call @unsafeUnion a b@, the caller must guarantee
-- that if @lookup k a = Just v@ then @lookup k b = Nothing@ and vice versa.
unsafeUnion :: IMap a -> IMap a -> IMap a
unsafeUnion :: forall a. IMap a -> IMap a -> IMap a
unsafeUnion IMap a
a IMap a
b = IMap { _runs :: IntMap (Run a)
_runs = IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs IMap a
a IntMap (Run a) -> IntMap (Run a) -> IntMap (Run a)
forall a. IntMap a -> IntMap a -> IntMap a
`IM.union` IMap a -> IntMap (Run a)
forall a. IMap a -> IntMap (Run a)
_runs IMap a
b }