1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
{-# 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, v) <| OMap tvs kvs = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
t = maybe (nextLowerTag kvs) fst (M.lookup k tvs)
(k, v) |< o = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
t = nextLowerTag kvs
OMap tvs kvs = delete k o
o >| (k, v) = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
t = nextHigherTag kvs
OMap tvs kvs = delete k o
OMap tvs kvs |> (k, v) = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
t = maybe (nextHigherTag kvs) fst (M.lookup k 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 = unionWithInternal (\t t' -> 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 = unionWithInternal (\t t' -> t')
unionWithInternal :: Ord k => (Tag -> Tag -> Tag) -> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
unionWithInternal fT fKV (OMap tvs kvs) (OMap tvs' kvs') = fromTV tvs'' where
bump = case maxTag kvs of
Nothing -> 0
Just k -> -k-1
bump' = case minTag kvs' of
Nothing -> 0
Just k -> -k
tvs'' = M.unionWithKey (\k (t,v) (t',v') -> (fT t t', fKV k v v'))
(fmap (\(t,v) -> (bump +t,v)) tvs )
(fmap (\(t,v) -> (bump'+t,v)) tvs')
singleton :: (k, v) -> OMap k v
singleton kv@(k, v) = OMap (M.singleton k (0, v)) (M.singleton 0 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 = foldl' (|>) 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 f (OMap tvs kvs) (OMap tvs' kvs') = fromTV
$ M.intersectionWithKey (\k (t,v) (t',v') -> (t, f k v v')) tvs 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 f k om@(OMap tvs kvs) = case M.lookup k tvs of
Just (t, _) -> OMap
(M.alter (fmap (t,) . f . fmap snd) k tvs)
(M.alter (fmap (k,) . f . fmap snd) t kvs)
Nothing -> maybe om ((om |>) . (k, )) $ f Nothing
|
