git.iko.soy/elk
on 577fc17 help sign in
elk/libraries/Dict.elm
412 lines 13 kB history raw
1module Dict (empty,singleton,insert,update
2 ,get,getOrElse,getOrFail
3 ,remove,member
4 ,filter
5 ,partition
6 ,foldl,foldr,map
7 ,union,intersect,diff
8 ,keys,values
9 ,toList,fromList
10 ) where
12{-| A dictionary mapping unique keys to values. The keys can be any comparable
13type. This includes `Int`, `Float`, `Time`, `Char`, `String`, and tuples or
14lists of comparable types.
16Insert, remove, and query operations all take *O(log n)* time.
18# Build
19@docs empty, singleton, insert, update, remove
21# Query
22@docs member, get, getOrElse, getOrFail
24# Combine
25@docs union, intersect, diff
27# Lists
28@docs keys, values, toList, fromList
30# Transform
31@docs map, foldl, foldr, filter, partition
33-}
36import Basics (..)
37import Maybe (..)
38import Native.Error
39import List
40import List ((::))
41import Native.Utils
43-- BBlack and NBlack should only be used during the deletion
44-- algorithm. Any other occurrence is a bug and should fail an assert.
45data NColor = Red
46 | Black
47 -- ^ Double Black, counts as 2 blacks for the invariant
48 | BBlack
49 -- ^ Negative Black, counts as -1 blacks for the invariant
50 | NBlack
53showNColor : NColor -> String
54showNColor c = case c of
55 Red -> "Red"
56 Black -> "Black"
57 BBlack -> "BBlack"
58 NBlack -> "NBlack"
60data LeafColor = LBlack
61 -- ^ Double Black, counts as 2
62 | LBBlack
64showLColor : LeafColor -> String
65showLColor c = case c of
66 LBlack -> "LBlack"
67 LBBlack -> "LBBlack"
69data Dict k v = RBNode NColor k v (Dict k v) (Dict k v)
70 | RBEmpty LeafColor
72{-| Create an empty dictionary. -}
73empty : Dict comparable v
74empty = RBEmpty LBlack
76min : Dict k v -> (k,v)
77min t =
78 case t of
79 RBNode _ k v (RBEmpty LBlack) _ -> (k,v)
80 RBNode _ _ _ l _ -> min l
81 RBEmpty LBlack -> Native.Error.raise "(min Empty) is not defined"
83max : Dict k v -> (k, v)
84max t =
85 case t of
86 RBNode _ k v _ (RBEmpty _) -> (k,v)
87 RBNode _ _ _ _ r -> max r
88 RBEmpty _ -> Native.Error.raise "(max Empty) is not defined"
90{-| Get the value associated with a key. If the key is not found, return
91`Nothing`. This is useful when you are not sure if a key will be in the
92dictionary.
94 animals = fromList [ ("Tom", Cat), ("Jerry", Mouse) ]
96 get "Tom" animals == Just Cat
97 get "Mouse" animals == Just Mouse
98 get "Spike" animals == Nothing
100The `getOrElse` and `getOrFail` are built-in ways to handle common ways of
101using the resulting `Maybe`.
102-}
103get : comparable -> Dict comparable v -> Maybe v
104get k t =
105 case t of
106 RBEmpty LBlack -> Nothing
107 RBNode _ k' v l r ->
108 case Native.Utils.compare k k' of
109 LT -> get k l
110 EQ -> Just v
111 GT -> get k r
113{-| Get the value associated with a key. If the key is not found,
114return a default value.
116 animals = fromList [ ("Tom", Cat), ("Jerry", Mouse) ]
118 getOrElse Dog "Tom" animals == Cat
119 getOrElse Dog "Mouse" animals == Mouse
120 getOrElse Dog "Spike" animals == Dog
121-}
122getOrElse : v -> comparable -> Dict comparable v -> v
123getOrElse base k t =
124 case t of
125 RBEmpty LBlack -> base
126 RBNode _ k' v l r ->
127 case Native.Utils.compare k k' of
128 LT -> getOrElse base k l
129 EQ -> v
130 GT -> getOrElse base k r
132{-| Get the value associated with a key.
134 animals = fromList [ ("Tom", Cat), ("Jerry", Mouse) ]
136 getOrFail "Tom" animals == Cat
137 getOrFail "Mouse" animals == Mouse
138 getOrFail "Spike" animals == -- Runtime Error!
140Warning: this function will result in a runtime error if the key is not found,
141so it is best to use `get` or `getOrElse` unless you are sure the key will be
142found.
143-}
144getOrFail : comparable -> Dict comparable v -> v
145getOrFail k t =
146 case t of
147 RBEmpty LBlack -> Native.Error.raise "key not found when using 'getOrFail'"
148 RBNode _ k' v l r ->
149 case Native.Utils.compare k k' of
150 LT -> getOrFail k l
151 EQ -> v
152 GT -> getOrFail k r
154{-| Determine if a key is in a dictionary. -}
155member : comparable -> Dict comparable v -> Bool
156-- Does t contain k?
157member k t = isJust <| get k t
159ensureBlackRoot : Dict k v -> Dict k v
160ensureBlackRoot t =
161 case t of
162 RBNode Red k v l r -> RBNode Black k v l r
163 RBNode Black _ _ _ _ -> t
164 RBEmpty LBlack -> t
166{-| Insert a key-value pair into a dictionary. Replaces value when there is
167a collision. -}
168insert : comparable -> v -> Dict comparable v -> Dict comparable v
169insert k v t = let u _ = Just v in
170 update k u t
172{-| Remove a key-value pair from a dictionary. If the key is not found,
173no changes are made. -}
174remove : comparable -> Dict comparable v -> Dict comparable v
175remove k t = let u _ = Nothing in
176 update k u t
178data Flag = Insert | Remove | Same
179showFlag : Flag -> String
180showFlag f = case f of
181 Insert -> "Insert"
182 Remove -> "Remove"
183 Same -> "Same"
185{-| Update the value of a dictionary for a specific key with a given function. -}
186update : comparable -> (Maybe v -> Maybe v) -> Dict comparable v -> Dict comparable v
187update k u t =
188 let up t = case t of
189 RBEmpty LBlack -> case u Nothing of
190 Nothing -> (Same, empty)
191 Just v -> (Insert, RBNode Red k v empty empty)
192 RBNode c k' v l r -> case Native.Utils.compare k k' of
193 EQ -> case u (Just v) of
194 Nothing -> (Remove, rem c l r)
195 Just v' -> (Same, RBNode c k' v' l r)
196 LT -> let (fl, l') = up l in
197 case fl of
198 Same -> (Same, RBNode c k' v l' r)
199 Insert -> (Insert, balance c k' v l' r)
200 Remove -> (Remove, bubble c k' v l' r)
201 GT -> let (fl, r') = up r in
202 case fl of
203 Same -> (Same, RBNode c k' v l r')
204 Insert -> (Insert, balance c k' v l r')
205 Remove -> (Remove, bubble c k' v l r')
206 (fl, t') = up t
207 in case fl of
208 Same -> t'
209 Insert -> ensureBlackRoot t'
210 Remove -> blacken t'
212{-| Create a dictionary with one key-value pair. -}
213singleton : comparable -> v -> Dict comparable v
214singleton k v = insert k v (RBEmpty LBlack)
216isBBlack : Dict k v -> Bool
217isBBlack t = case t of
218 RBNode BBlack _ _ _ _ -> True
219 RBEmpty LBBlack -> True
220 _ -> False
222moreBlack : NColor -> NColor
223moreBlack c = case c of
224 Black -> BBlack
225 Red -> Black
226 NBlack -> Red
227 BBlack -> Native.Error.raise "Can't make a double black node more black!"
229lessBlack : NColor -> NColor
230lessBlack c = case c of
231 BBlack -> Black
232 Black -> Red
233 Red -> NBlack
234 NBlack -> Native.Error.raise "Can't make a negative black node less black!"
236lessBlackTree : Dict k v -> Dict k v
237lessBlackTree t = case t of
238 RBNode c k v l r -> RBNode (lessBlack c) k v l r
239 RBEmpty LBBlack -> RBEmpty LBlack
241reportRemBug : String -> NColor -> String -> String -> a
242reportRemBug msg c lgot rgot =
243 Native.Error.raise <| List.concat [
244 "Internal red-black tree invariant violated, expected ",
245 msg,
246 "and got",
247 showNColor c,
248 " ",
249 lgot,
250 " ",
251 rgot,
252 "\nPlease report this bug to https://github.com/elm-lang/Elm/issues"
253 ]
255-- Remove the top node from the tree, may leave behind BBlacks
256rem : NColor -> Dict k v -> Dict k v -> Dict k v
257rem c l r = case (l, r) of
258 ((RBEmpty _), (RBEmpty _)) -> case c of
259 Red -> RBEmpty LBlack
260 Black -> RBEmpty LBBlack
261 ((RBEmpty cl), (RBNode cr k' v' l' r')) ->
262 case (c, cl, cr) of
263 (Black, LBlack, Red) -> RBNode Black k' v' l' r'
264 _ -> reportRemBug "Black, LBlack, Red" c (showLColor cl) (showNColor cr)
265 ((RBNode cl k' v' l' r'), (RBEmpty cr)) ->
266 case (c, cl, cr) of
267 (Black, Red, LBlack) -> RBNode Black k' v' l' r'
268 _ -> reportRemBug "Black, Red, LBlack" c (showNColor cl) (showLColor cr)
269 -- l and r are both RBNodes
270 ((RBNode cl kl vl ll rl), (RBNode cr kr vr lr rr)) ->
271 let l = RBNode cl kl vl ll rl
272 r = RBNode cr kr vr lr rr
273 (k, v) = max l
274 l' = remove_max cl kl vl ll rl
275 in bubble c k v l' r
277-- Kills a BBlack or moves it upward, may leave behind NBlack
278bubble : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
279bubble c k v l r = if isBBlack l || isBBlack r
280 then balance (moreBlack c) k v (lessBlackTree l) (lessBlackTree r)
281 else RBNode c k v l r
283-- Removes rightmost node, may leave root as BBlack
284remove_max : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
285remove_max c k v l r = case r of
286 RBEmpty _ -> rem c l r
287 RBNode cr kr vr lr rr
288 -> bubble c k v l (remove_max cr kr vr lr rr)
290-- generalized tree balancing act
291balance : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
292balance c k v l r =
294 balance_node (RBNode c k v l r)
296blackish : Dict k v -> Bool
297blackish t = case t of
298 RBNode c _ _ _ _ -> c == Black || c == BBlack
299 RBEmpty _ -> True
301balance_node : Dict k v -> Dict k v
302balance_node t =
303 let assemble col xk xv yk yv zk zv a b c d =
304 RBNode (lessBlack col) yk yv (RBNode Black xk xv a b) (RBNode Black zk zv c d)
305 in
306 if blackish t
307 then case t of
308 RBNode col zk zv (RBNode Red yk yv (RBNode Red xk xv a b) c) d ->
309 assemble col xk xv yk yv zk zv a b c d
310 RBNode col zk zv (RBNode Red xk xv a (RBNode Red yk yv b c)) d ->
311 assemble col xk xv yk yv zk zv a b c d
312 RBNode col xk xv a (RBNode Red zk zv (RBNode Red yk yv b c) d) ->
313 assemble col xk xv yk yv zk zv a b c d
314 RBNode col xk xv a (RBNode Red yk yv b (RBNode Red zk zv c d)) ->
315 assemble col xk xv yk yv zk zv a b c d
317 RBNode BBlack xk xv a (RBNode NBlack zk zv (RBNode Black yk yv b c) d) ->
318 case d of
319 (RBNode Black _ _ _ _) ->
320 RBNode Black yk yv (RBNode Black xk xv a b) (balance Black zk zv c (redden d))
321 _ -> t
323 RBNode BBlack zk zv (RBNode NBlack xk xv a (RBNode Black yk yv b c)) d ->
324 case a of
325 (RBNode Black _ _ _ _) ->
326 RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
327 _ -> t
328 _ -> t
329 else t
331-- make the top node black
332blacken : Dict k v -> Dict k v
333blacken t = case t of
334 RBEmpty _ -> RBEmpty LBlack
335 RBNode _ k v l r -> RBNode Black k v l r
337-- make the top node red
338redden : Dict k v -> Dict k v
339redden t = case t of
340 RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
341 RBNode _ k v l r -> RBNode Red k v l r
343{-| Apply a function to all values in a dictionary. -}
344map : (a -> b) -> Dict comparable a -> Dict comparable b
345map f t =
346 case t of
347 RBEmpty LBlack -> RBEmpty LBlack
348 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
350{-| Fold over the key-value pairs in a dictionary, in order from lowest
351key to highest key. -}
352foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
353foldl f acc t =
354 case t of
355 RBEmpty LBlack -> acc
356 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
358{-| Fold over the key-value pairs in a dictionary, in order from highest
359key to lowest key. -}
360foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
361foldr f acc t =
362 case t of
363 RBEmpty LBlack -> acc
364 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
366{-| Combine two dictionaries. If there is a collision, preference is given
367to the first dictionary. -}
368union : Dict comparable v -> Dict comparable v -> Dict comparable v
369union t1 t2 = foldl insert t2 t1
371{-| Keep a key-value pair when its key appears in the second dictionary.
372Preference is given to values in the first dictionary. -}
373intersect : Dict comparable v -> Dict comparable v -> Dict comparable v
374intersect t1 t2 = filter (\k _ -> k `member` t2) t1
376{-| Keep a key-value pair when its key does not appear in the second dictionary.
377Preference is given to the first dictionary. -}
378diff : Dict comparable v -> Dict comparable v -> Dict comparable v
379diff t1 t2 = foldl (\k v t -> remove k t) t1 t2
381{-| Get all of the keys in a dictionary. -}
382keys : Dict comparable v -> [comparable]
383keys t = foldr (\k v acc -> k :: acc) [] t
385{-| Get all of the values in a dictionary. -}
386values : Dict comparable v -> [v]
387values t = foldr (\k v acc -> v :: acc) [] t
389{-| Convert a dictionary into an association list of key-value pairs. -}
390toList : Dict comparable v -> [(comparable,v)]
391toList t = foldr (\k v acc -> (k,v) :: acc) [] t
393{-| Convert an association list into a dictionary. -}
394fromList : [(comparable,v)] -> Dict comparable v
395fromList assocs = List.foldl (\(k,v) d -> insert k v d) empty assocs
397{-| Keep a key-value pair when it satisfies a predicate. -}
398filter : (comparable -> v -> Bool) -> Dict comparable v -> Dict comparable v
399filter p dict =
400 let add k v t = if p k v then insert k v t else t
401 in foldl add empty dict
403{-| Partition a dictionary according to a predicate. The first dictionary
404contains all key-value pairs which satisfy the predicate, and the second
405contains the rest.
406-}
407partition : (comparable -> v -> Bool) -> Dict comparable v -> (Dict comparable v, Dict comparable v)
408partition p dict =
409 let add k v (t1, t2) = if p k v
410 then (insert k v t1, t2)
411 else (t1, insert k v t2)
412 in foldl add (empty, empty) dict