@@ -1,8 +1,9 @@
11{-# LANGUAGE DeriveDataTypeable #-}
22{-# LANGUAGE DeriveFunctor #-}
33
4−-- | An 'OMap' behaves much like a 'Map', with all the same asymptotics, but
5−-- also remembers the order that keys were inserted.
4+-- | An 'OMap' behaves much like a 'Map', with mostly the same asymptotics, but
5+-- also remembers the order that keys were inserted. All operations whose
6+-- asymptotics are worse than 'Map' have documentation saying so.
67module Data.Map.Ordered
78 ( OMap
89 -- * Trivial maps
@@ -18,8 +19,10 @@ module Data.Map.Ordered
1819 --
1920 -- * If both sides contain the same key, the tuple's value wins
2021 , (<|), (|<), (>|), (|>)
22+ , (<>|), (|<>), unionWithL, unionWithR
2123 -- * Deletion
2224 , delete, filter, (\\)
25+ , (|/\), (/\|), intersectionWith
2326 -- * Query
2427 , null, size, member, notMember, lookup
2528 -- * Indexing
@@ -67,10 +70,21 @@ oMapDataType = mkDataType "Data.Map.Ordered.Map" [fromListConstr]
6770
6871infixr 5 <|, |< -- copy :
6972infixl 5 >|, |>
73+infixr 6 <>|, |<> -- copy <>
7074
7175(<|) , (|<) :: Ord k => (,) k v -> OMap k v -> OMap k v
7276(>|) , (|>) :: Ord k => OMap k v -> (,) k v -> OMap k v
7377
78+-- | When a key occurs in both maps, prefer the value from the first map.
79+--
80+-- See asymptotics of 'unionWithR'.
81+(<>|) :: Ord k => OMap k v -> OMap k v -> OMap k v
82+
83+-- | When a key occurs in both maps, prefer the value from the first map.
84+--
85+-- See asymptotics of 'unionWithL'.
86+(|<>) :: Ord k => OMap k v -> OMap k v -> OMap k v
87+
7488(k, v) <| OMap tvs kvs = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
7589 t = maybe (nextLowerTag kvs) fst (M.lookup k tvs)
7690
@@ -85,7 +99,41 @@ o >| (k, v) = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
8599OMap tvs kvs |> (k, v) = OMap (M.insert k (t, v) tvs) (M.insert t (k, v) kvs) where
86100 t = maybe (nextHigherTag kvs) fst (M.lookup k tvs)
87101
102+(<>|) = unionWithR (const const)
103+(|<>) = unionWithL (const const)
104+
105+-- | Take the union. The first 'OMap' \'s argument's indices are lower than the
106+-- second. If a key appears in both maps, the first argument's index takes
107+-- precedence, and the supplied function is used to combine the values.
108+--
109+-- /O(r*log(r))/ where /r/ is the size of the result
110+unionWithL :: Ord k => (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
111+unionWithL = unionWithInternal (\t t' -> t )
112+
113+-- | Take the union. The first 'OMap' \'s argument's indices are lower than the
114+-- second. If a key appears in both maps, the second argument's index takes
115+-- precedence, and the supplied function is used to combine the values.
116+--
117+-- /O(r*log(r))/ where /r/ is the size of the result
118+unionWithR :: Ord k => (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
119+unionWithR = unionWithInternal (\t t' -> t')
120+
121+unionWithInternal :: Ord k => (Tag -> Tag -> Tag) -> (k -> v -> v -> v) -> OMap k v -> OMap k v -> OMap k v
122+unionWithInternal fT fKV (OMap tvs kvs) (OMap tvs' kvs') = fromTV tvs'' where
123+ bump = case maxTag kvs of
124+ Nothing -> 0
125+ Just k -> -k-1
126+ bump' = case minTag kvs' of
127+ Nothing -> 0
128+ Just k -> -k
129+ tvs'' = M.unionWithKey (\k (t,v) (t',v') -> (fT t t', fKV k v v'))
130+ (fmap (\(t,v) -> (bump +t,v)) tvs )
131+ (fmap (\(t,v) -> (bump'+t,v)) tvs')
132+
88133-- | @m \\\\ n@ deletes all the keys that exist in @n@ from @m@
134+--
135+-- /O(m*log(n))/ where /m/ is the size of the smaller map and /n/ is the size
136+-- of the larger map.
89137(\\) :: Ord k => OMap k v -> OMap k v' -> OMap k v
90138o@(OMap tvs kvs) \\ o'@(OMap tvs' kvs') = if size o < size o'
91139 then filter (const . (`notMember` o')) o
@@ -129,6 +177,36 @@ delete k o@(OMap tvs kvs) = case M.lookup k tvs of
129177 Nothing -> o
130178 Just (t, _) -> OMap (M.delete k tvs) (M.delete t kvs)
131179
180+-- | Intersection. (The @/\\@ is intended to look a bit like the standard
181+-- mathematical notation for set intersection.)
182+--
183+-- See asymptotics of 'intersectionWith'.
184+(/\|) :: Ord k => OMap k v -> OMap k v' -> OMap k v
185+o /\| o' = intersectionWith (\k v' v -> v) o' o
186+
187+-- | Intersection. (The @/\\@ is intended to look a bit like the standard
188+-- mathematical notation for set intersection.)
189+--
190+-- See asymptotics of 'intersectionWith'.
191+(|/\) :: Ord k => OMap k v -> OMap k v' -> OMap k v
192+o |/\ o' = intersectionWith (\k v v' -> v) o o'
193+
194+-- | Take the intersection. The first 'OMap' \'s argument's indices are used for
195+-- the result.
196+--
197+-- /O(m*log(n\/(m+1)) + r*log(r))/ where /m/ is the size of the smaller map, /n/
198+-- is the size of the larger map, and /r/ is the size of the result.
199+intersectionWith ::
200+ Ord k =>
201+ (k -> v -> v' -> v'') ->
202+ OMap k v -> OMap k v' -> OMap k v''
203+intersectionWith f (OMap tvs kvs) (OMap tvs' kvs') = fromTV
204+ $ M.intersectionWithKey (\k (t,v) (t',v') -> (t, f k v v')) tvs tvs'
205+
206+fromTV :: Ord k => Map k (Tag, v) -> OMap k v
207+fromTV tvs = OMap tvs kvs where
208+ kvs = M.fromList [(t,(k,v)) | (k,(t,v)) <- M.toList tvs]
209+
132210findIndex :: Ord k => k -> OMap k v -> Maybe Index
133211findIndex k o@(OMap tvs kvs) = do
134212 (t, _) <- M.lookup k tvs