{-# LANGUAGE TupleSections #-}

-- | An 'OMap' behaves much like a 'M.Map', with mostly the same asymptotics, but
-- also remembers the order that keys were inserted. All operations whose
-- asymptotics are worse than 'M.Map' have documentation saying so.
module Data.Map.Ordered.Strict
	( OMap
	-- * Trivial maps
	, empty, singleton
	-- * Insertion
	-- | Conventions:
	--
	-- * The open side of an angle bracket points to an 'OMap'
	--
	-- * The pipe appears on the side whose indices take precedence if both sides contain the same key
	--
	-- * The left argument's indices are lower than the right argument's indices
	--
	-- * If both sides contain the same key, the tuple's value wins
	, (<|), (|<), (>|), (|>)
	, (<>|), (|<>), unionWithL, unionWithR
	, Bias(Bias, unbiased), L, R
	-- * Deletion
	, delete, filter, (\\)
	, (|/\), (/\|), intersectionWith
	-- * Query
	, null, size, member, notMember, lookup
	-- * Indexing
	, Index, findIndex, elemAt
	-- * List conversions
	, fromList, assocs, toAscList
	-- * 'M.Map' conversion
	, toMap
	) where

import Data.Foldable (foldl')
import qualified Data.Map.Strict as M
import Data.Map.Ordered.Internal
	( OMap(..), empty, (<>|), (|<>), delete, filter, (\\), (|/\), (/\|), null, size
	, member, notMember, lookup, findIndex, elemAt, assocs, toAscList, fromTV, toMap )
import Data.Map.Util
import Prelude hiding (filter, lookup, null)

infixr 5 <|, |< -- copy :
infixl 5 >|, |>

(<|) , (|<) :: Ord k => (,)  k v -> OMap k v -> OMap k v
(>|) , (|>) :: Ord k => OMap k v -> (,)  k v -> OMap k v

(k
k, v
v) <| :: forall k v. Ord k => (k, v) -> OMap k v -> OMap k v
<| OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs = Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
forall k v. Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
OMap (k -> (Tag, v) -> Map k (Tag, v) -> Map k (Tag, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (Tag
t, v
v) Map k (Tag, v)
tvs) (Tag -> (k, v) -> Map Tag (k, v) -> Map Tag (k, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Tag
t (k
k, v
v) Map Tag (k, v)
kvs) where
	t :: Tag
t = Tag -> ((Tag, v) -> Tag) -> Maybe (Tag, v) -> Tag
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map Tag (k, v) -> Tag
forall a. Map Tag a -> Tag
nextLowerTag Map Tag (k, v)
kvs) (Tag, v) -> Tag
forall a b. (a, b) -> a
fst (k -> Map k (Tag, v) -> Maybe (Tag, v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k (Tag, v)
tvs)

(k
k, v
v) |< :: forall k v. Ord k => (k, v) -> OMap k v -> OMap k v
|< OMap k v
o = Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
forall k v. Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
OMap (k -> (Tag, v) -> Map k (Tag, v) -> Map k (Tag, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (Tag
t, v
v) Map k (Tag, v)
tvs) (Tag -> (k, v) -> Map Tag (k, v) -> Map Tag (k, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Tag
t (k
k, v
v) Map Tag (k, v)
kvs) where
	t :: Tag
t = Map Tag (k, v) -> Tag
forall a. Map Tag a -> Tag
nextLowerTag Map Tag (k, v)
kvs
	OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs = k -> OMap k v -> OMap k v
forall k v. Ord k => k -> OMap k v -> OMap k v
delete k
k OMap k v
o

OMap k v
o >| :: forall k v. Ord k => OMap k v -> (k, v) -> OMap k v
>| (k
k, v
v) = Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
forall k v. Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
OMap (k -> (Tag, v) -> Map k (Tag, v) -> Map k (Tag, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (Tag
t, v
v) Map k (Tag, v)
tvs) (Tag -> (k, v) -> Map Tag (k, v) -> Map Tag (k, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Tag
t (k
k, v
v) Map Tag (k, v)
kvs) where
	t :: Tag
t = Map Tag (k, v) -> Tag
forall a. Map Tag a -> Tag
nextHigherTag Map Tag (k, v)
kvs
	OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs = k -> OMap k v -> OMap k v
forall k v. Ord k => k -> OMap k v -> OMap k v
delete k
k OMap k v
o

OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs |> :: forall k v. Ord k => OMap k v -> (k, v) -> OMap k v
|> (k
k, v
v) = Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
forall k v. Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
OMap (k -> (Tag, v) -> Map k (Tag, v) -> Map k (Tag, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (Tag
t, v
v) Map k (Tag, v)
tvs) (Tag -> (k, v) -> Map Tag (k, v) -> Map Tag (k, v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Tag
t (k
k, v
v) Map Tag (k, v)
kvs) where
	t :: Tag
t = Tag -> ((Tag, v) -> Tag) -> Maybe (Tag, v) -> Tag
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map Tag (k, v) -> Tag
forall a. Map Tag a -> Tag
nextHigherTag Map Tag (k, v)
kvs) (Tag, v) -> Tag
forall a b. (a, b) -> a
fst (k -> Map k (Tag, v) -> Maybe (Tag, v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k (Tag, v)
tvs)

-- | Take the union. The first 'OMap' \'s argument's indices are lower than the
-- second. If a key appears in both maps, the first argument's index takes
-- precedence, and the supplied function is used to combine the values.
--
-- /O(r*log(r))/ where /r/ is the size of the result
--
-- @since 0.2
unionWithL :: Ord k => (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithL :: forall k v.
Ord k =>
(k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithL = (Tag -> Tag -> Tag)
-> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
forall k v.
Ord k =>
(Tag -> Tag -> Tag)
-> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithInternal (\Tag
t Tag
t' -> Tag
t )

-- | Take the union. The first 'OMap' \'s argument's indices are lower than the
-- second. If a key appears in both maps, the second argument's index takes
-- precedence, and the supplied function is used to combine the values.
--
-- /O(r*log(r))/ where /r/ is the size of the result
--
-- @since 0.2
unionWithR :: Ord k => (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithR :: forall k v.
Ord k =>
(k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithR = (Tag -> Tag -> Tag)
-> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
forall k v.
Ord k =>
(Tag -> Tag -> Tag)
-> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithInternal (\Tag
t Tag
t' -> Tag
t')

unionWithInternal :: Ord k => (Tag -> Tag -> Tag) -> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithInternal :: forall k v.
Ord k =>
(Tag -> Tag -> Tag)
-> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithInternal Tag -> Tag -> Tag
fT k -> v -> v -> v
fKV (OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs) (OMap Map k (Tag, v)
tvs' Map Tag (k, v)
kvs') = Map k (Tag, v) -> OMap k v
forall k v. Ord k => Map k (Tag, v) -> OMap k v
fromTV Map k (Tag, v)
tvs'' where
	bump :: Tag
bump  = case Map Tag (k, v) -> Maybe Tag
forall a. Map Tag a -> Maybe Tag
maxTag Map Tag (k, v)
kvs  of
		Maybe Tag
Nothing -> Tag
0
		Just Tag
k  -> -Tag
kTag -> Tag -> Tag
forall a. Num a => a -> a -> a
-Tag
1
	bump' :: Tag
bump' = case Map Tag (k, v) -> Maybe Tag
forall a. Map Tag a -> Maybe Tag
minTag Map Tag (k, v)
kvs' of
		Maybe Tag
Nothing -> Tag
0
		Just Tag
k  -> -Tag
k
	tvs'' :: Map k (Tag, v)
tvs'' = (k -> (Tag, v) -> (Tag, v) -> (Tag, v))
-> Map k (Tag, v) -> Map k (Tag, v) -> Map k (Tag, v)
forall k a.
Ord k =>
(k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWithKey (\k
k (Tag
t,v
v) (Tag
t',v
v') -> (Tag -> Tag -> Tag
fT Tag
t Tag
t', k -> v -> v -> v
fKV k
k v
v v
v'))
		(((Tag, v) -> (Tag, v)) -> Map k (Tag, v) -> Map k (Tag, v)
forall a b. (a -> b) -> Map k a -> Map k b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Tag
t,v
v) -> (Tag
bump Tag -> Tag -> Tag
forall a. Num a => a -> a -> a
+Tag
t,v
v)) Map k (Tag, v)
tvs )
		(((Tag, v) -> (Tag, v)) -> Map k (Tag, v) -> Map k (Tag, v)
forall a b. (a -> b) -> Map k a -> Map k b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Tag
t,v
v) -> (Tag
bump'Tag -> Tag -> Tag
forall a. Num a => a -> a -> a
+Tag
t,v
v)) Map k (Tag, v)
tvs')

singleton :: (k, v) -> OMap k v
singleton :: forall k v. (k, v) -> OMap k v
singleton kv :: (k, v)
kv@(k
k, v
v) = Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
forall k v. Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
OMap (k -> (Tag, v) -> Map k (Tag, v)
forall k a. k -> a -> Map k a
M.singleton k
k (Tag
0, v
v)) (Tag -> (k, v) -> Map Tag (k, v)
forall k a. k -> a -> Map k a
M.singleton Tag
0 (k, v)
kv)

-- | If a key appears multiple times, the first occurrence is used for ordering
-- and the last occurrence is used for its value. The library author welcomes
-- comments on whether this default is sane.
fromList :: Ord k => [(k, v)] -> OMap k v
fromList :: forall k v. Ord k => [(k, v)] -> OMap k v
fromList = (OMap k v -> (k, v) -> OMap k v)
-> OMap k v -> [(k, v)] -> OMap k v
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' OMap k v -> (k, v) -> OMap k v
forall k v. Ord k => OMap k v -> (k, v) -> OMap k v
(|>) OMap k v
forall k v. OMap k v
empty

-- | Take the intersection. The first 'OMap' \'s argument's indices are used for
-- the result.
--
-- /O(m*log(n\/(m+1)) + r*log(r))/ where /m/ is the size of the smaller map, /n/
-- is the size of the larger map, and /r/ is the size of the result.
--
-- @since 0.2
intersectionWith ::
	Ord k =>
	(k -> v -> v' -> v'') ->
	OMap k v -> OMap k v' -> OMap k v''
intersectionWith :: forall k v v' v''.
Ord k =>
(k -> v -> v' -> v'') -> OMap k v -> OMap k v' -> OMap k v''
intersectionWith k -> v -> v' -> v''
f (OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs) (OMap Map k (Tag, v')
tvs' Map Tag (k, v')
kvs') = Map k (Tag, v'') -> OMap k v''
forall k v. Ord k => Map k (Tag, v) -> OMap k v
fromTV
	(Map k (Tag, v'') -> OMap k v'') -> Map k (Tag, v'') -> OMap k v''
forall a b. (a -> b) -> a -> b
$ (k -> (Tag, v) -> (Tag, v') -> (Tag, v''))
-> Map k (Tag, v) -> Map k (Tag, v') -> Map k (Tag, v'')
forall k a b c.
Ord k =>
(k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWithKey (\k
k (Tag
t,v
v) (Tag
t',v'
v') -> (Tag
t, k -> v -> v' -> v''
f k
k v
v v'
v')) Map k (Tag, v)
tvs Map k (Tag, v')
tvs'

-- | Alter the value (or its absence) associated with a key.
--
-- @since 0.2.4
alter :: Ord k => (Maybe v -> Maybe v) -> k -> OMap k v -> OMap k v
alter :: forall k v.
Ord k =>
(Maybe v -> Maybe v) -> k -> OMap k v -> OMap k v
alter Maybe v -> Maybe v
f k
k om :: OMap k v
om@(OMap Map k (Tag, v)
tvs Map Tag (k, v)
kvs) = case k -> Map k (Tag, v) -> Maybe (Tag, v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k (Tag, v)
tvs of
	Just (Tag
t, v
_) -> Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
forall k v. Map k (Tag, v) -> Map Tag (k, v) -> OMap k v
OMap
		((Maybe (Tag, v) -> Maybe (Tag, v))
-> k -> Map k (Tag, v) -> Map k (Tag, v)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter ((v -> (Tag, v)) -> Maybe v -> Maybe (Tag, v)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Tag
t,) (Maybe v -> Maybe (Tag, v))
-> (Maybe (Tag, v) -> Maybe v) -> Maybe (Tag, v) -> Maybe (Tag, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe v -> Maybe v
f (Maybe v -> Maybe v)
-> (Maybe (Tag, v) -> Maybe v) -> Maybe (Tag, v) -> Maybe v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Tag, v) -> v) -> Maybe (Tag, v) -> Maybe v
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Tag, v) -> v
forall a b. (a, b) -> b
snd) k
k Map k (Tag, v)
tvs)
		((Maybe (k, v) -> Maybe (k, v))
-> Tag -> Map Tag (k, v) -> Map Tag (k, v)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter ((v -> (k, v)) -> Maybe v -> Maybe (k, v)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (k
k,) (Maybe v -> Maybe (k, v))
-> (Maybe (k, v) -> Maybe v) -> Maybe (k, v) -> Maybe (k, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe v -> Maybe v
f (Maybe v -> Maybe v)
-> (Maybe (k, v) -> Maybe v) -> Maybe (k, v) -> Maybe v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> v) -> Maybe (k, v) -> Maybe v
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (k, v) -> v
forall a b. (a, b) -> b
snd) Tag
t Map Tag (k, v)
kvs)
	Maybe (Tag, v)
Nothing -> OMap k v -> (v -> OMap k v) -> Maybe v -> OMap k v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe OMap k v
om ((OMap k v
om OMap k v -> (k, v) -> OMap k v
forall k v. Ord k => OMap k v -> (k, v) -> OMap k v
|>) ((k, v) -> OMap k v) -> (v -> (k, v)) -> v -> OMap k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k
k, )) (Maybe v -> OMap k v) -> Maybe v -> OMap k v
forall a b. (a -> b) -> a -> b
$ Maybe v -> Maybe v
f Maybe v
forall a. Maybe a
Nothing