git.iko.soy/elk
on 63243ca help sign in
elk/libraries/Dict.elm
354 lines 11 kB history raw
1module Dict (empty,singleton,insert,update
2 ,lookup,findWithDefault
3 ,remove,member
4 ,foldl,foldr,map
5 ,union,intersect,diff
6 ,keys,values
7 ,toList,fromList
8 ) where
10{-| A dictionary mapping unique keys to values. The keys can be any comparable
11type. This includes `Int`, `Float`, `Time`, `Char`, `String`, and tuples or
12lists of comparable types.
14Insert, remove, and query operations all take *O(log n)* time.
16# Build
17@docs empty, singleton, insert, update, remove
19# Query
20@docs member, lookup, findWithDefault
22# Combine
23@docs union, intersect, diff
25# Lists
26@docs keys, values, toList, fromList
28# Transform
29@docs map, foldl, foldr
31-}
34import open Basics
35import open Maybe
36import Native.Error
37import List
38import String
39import Native.Utils
41-- BBlack and NBlack should only be used during the deletion
42-- algorithm. Any other occurrence is a bug and should fail an assert.
43data NColor = Red
44 | Black
45 -- ^ Double Black, counts as 2 blacks for the invariant
46 | BBlack
47 -- ^ Negative Black, counts as -1 blacks for the invariant
48 | NBlack
51showNColor : NColor -> String
52showNColor c = case c of
53 Red -> "Red"
54 Black -> "Black"
55 BBlack -> "BBlack"
56 NBlack -> "NBlack"
58data LeafColor = LBlack
59 -- ^ Double Black, counts as 2
60 | LBBlack
62showLColor : LeafColor -> String
63showLColor c = case c of
64 LBlack -> "LBlack"
65 LBBlack -> "LBBlack"
67data Dict k v = RBNode NColor k v (Dict k v) (Dict k v)
68 | RBEmpty LeafColor
70{-| Create an empty dictionary. -}
71empty : Dict comparable v
72empty = RBEmpty LBlack
74min : Dict k v -> (k,v)
75min t =
76 case t of
77 RBNode _ k v (RBEmpty LBlack) _ -> (k,v)
78 RBNode _ _ _ l _ -> min l
79 RBEmpty LBlack -> Native.Error.raise "(min Empty) is not defined"
81max : Dict k v -> (k, v)
82max t =
83 case t of
84 RBNode _ k v _ (RBEmpty _) -> (k,v)
85 RBNode _ _ _ _ r -> max r
86 RBEmpty _ -> Native.Error.raise "(max Empty) is not defined"
88{-| Lookup the value associated with a key. -}
89lookup : comparable -> Dict comparable v -> Maybe v
90lookup k t =
91 case t of
92 RBEmpty LBlack -> Nothing
93 RBNode _ k' v l r ->
94 case Native.Utils.compare k k' of
95 LT -> lookup k l
96 EQ -> Just v
97 GT -> lookup k r
99{-| Find the value associated with a key. If the key is not found,
100return the default value. -}
101findWithDefault : v -> comparable -> Dict comparable v -> v
102findWithDefault base k t =
103 case t of
104 RBEmpty LBlack -> base
105 RBNode _ k' v l r ->
106 case Native.Utils.compare k k' of
107 LT -> findWithDefault base k l
108 EQ -> v
109 GT -> findWithDefault base k r
111{-| Determine if a key is in a dictionary. -}
112member : comparable -> Dict comparable v -> Bool
113-- Does t contain k?
114member k t = isJust <| lookup k t
116ensureBlackRoot : Dict k v -> Dict k v
117ensureBlackRoot t =
118 case t of
119 RBNode Red k v l r -> RBNode Black k v l r
120 RBNode Black _ _ _ _ -> t
121 RBEmpty LBlack -> t
123{-| Insert a key-value pair into a dictionary. Replaces value when there is
124a collision. -}
125insert : comparable -> v -> Dict comparable v -> Dict comparable v
126insert k v t = let u _ = Just v in
127 update k u t
129{-| Remove a key-value pair from a dictionary. If the key is not found,
130no changes are made. -}
131remove : comparable -> Dict comparable v -> Dict comparable v
132remove k t = let u _ = Nothing in
133 update k u t
135data Flag = Insert | Remove | Same
136showFlag : Flag -> String
137showFlag f = case f of
138 Insert -> "Insert"
139 Remove -> "Remove"
140 Same -> "Same"
142{-| Update the value of a dictionary for a specific key with a given function. -}
143update : comparable -> (Maybe v -> Maybe v) -> Dict comparable v -> Dict comparable v
144update k u t =
145 let up t = case t of
146 RBEmpty LBlack -> case u Nothing of
147 Nothing -> (Same, empty)
148 Just v -> (Insert, RBNode Red k v empty empty)
149 RBNode c k' v l r -> case Native.Utils.compare k k' of
150 EQ -> case u (Just v) of
151 Nothing -> (Remove, rem c l r)
152 Just v' -> (Same, RBNode c k' v' l r)
153 LT -> let (fl, l') = up l in
154 case fl of
155 Same -> (Same, RBNode c k' v l' r)
156 Insert -> (Insert, balance c k' v l' r)
157 Remove -> (Remove, bubble c k' v l' r)
158 GT -> let (fl, r') = up r in
159 case fl of
160 Same -> (Same, RBNode c k' v l r')
161 Insert -> (Insert, balance c k' v l r')
162 Remove -> (Remove, bubble c k' v l r')
163 (fl, t') = up t
164 in case fl of
165 Same -> t'
166 Insert -> ensureBlackRoot t'
167 Remove -> blacken t'
169{-| Create a dictionary with one key-value pair. -}
170singleton : comparable -> v -> Dict comparable v
171singleton k v = insert k v (RBEmpty LBlack)
173isBBlack : Dict k v -> Bool
174isBBlack t = case t of
175 RBNode BBlack _ _ _ _ -> True
176 RBEmpty LBBlack -> True
177 _ -> False
179moreBlack : NColor -> NColor
180moreBlack c = case c of
181 Black -> BBlack
182 Red -> Black
183 NBlack -> Red
184 BBlack -> Native.Error.raise "Can't make a double black node more black!"
186lessBlack : NColor -> NColor
187lessBlack c = case c of
188 BBlack -> Black
189 Black -> Red
190 Red -> NBlack
191 NBlack -> Native.Error.raise "Can't make a negative black node less black!"
193lessBlackTree : Dict k v -> Dict k v
194lessBlackTree t = case t of
195 RBNode c k v l r -> RBNode (lessBlack c) k v l r
196 RBEmpty LBBlack -> RBEmpty LBlack
198reportRemBug : String -> NColor -> String -> String -> a
199reportRemBug msg c lgot rgot =
200 Native.Error.raise . String.concat <| [
201 "Internal red-black tree invariant violated, expected ",
202 msg,
203 "and got",
204 showNColor c,
205 " ",
206 lgot,
207 " ",
208 rgot,
209 "\nPlease report this bug to https://github.com/evancz/Elm/issues"
210 ]
212-- Remove the top node from the tree, may leave behind BBlacks
213rem : NColor -> Dict k v -> Dict k v -> Dict k v
214rem c l r = case (l, r) of
215 ((RBEmpty _), (RBEmpty _)) -> case c of
216 Red -> RBEmpty LBlack
217 Black -> RBEmpty LBBlack
218 ((RBEmpty cl), (RBNode cr k' v' l' r')) ->
219 case (c, cl, cr) of
220 (Black, LBlack, Red) -> RBNode Black k' v' l' r'
221 _ -> reportRemBug "Black, LBlack, Red" c (showLColor cl) (showNColor cr)
222 ((RBNode cl k' v' l' r'), (RBEmpty cr)) ->
223 case (c, cl, cr) of
224 (Black, Red, LBlack) -> RBNode Black k' v' l' r'
225 _ -> reportRemBug "Black, Red, LBlack" c (showNColor cl) (showLColor cr)
226 -- l and r are both RBNodes
227 ((RBNode cl kl vl ll rl), (RBNode cr kr vr lr rr)) ->
228 let l = RBNode cl kl vl ll rl
229 r = RBNode cr kr vr lr rr
230 (k, v) = max l
231 l' = remove_max cl kl vl ll rl
232 in bubble c k v l' r
234-- Kills a BBlack or moves it upward, may leave behind NBlack
235bubble : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
236bubble c k v l r = if isBBlack l || isBBlack r
237 then balance (moreBlack c) k v (lessBlackTree l) (lessBlackTree r)
238 else RBNode c k v l r
240-- Removes rightmost node, may leave root as BBlack
241remove_max : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
242remove_max c k v l r = case r of
243 RBEmpty _ -> rem c l r
244 RBNode cr kr vr lr rr
245 -> bubble c k v l (remove_max cr kr vr lr rr)
247-- generalized tree balancing act
248balance : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
249balance c k v l r =
251 balance_node (RBNode c k v l r)
253blackish : Dict k v -> Bool
254blackish t = case t of
255 RBNode c _ _ _ _ -> c == Black || c == BBlack
256 RBEmpty _ -> True
258balance_node : Dict k v -> Dict k v
259balance_node t =
260 let assemble col xk xv yk yv zk zv a b c d =
261 RBNode (lessBlack col) yk yv (RBNode Black xk xv a b) (RBNode Black zk zv c d)
262 in
263 if blackish t
264 then case t of
265 RBNode col zk zv (RBNode Red yk yv (RBNode Red xk xv a b) c) d ->
266 assemble col xk xv yk yv zk zv a b c d
267 RBNode col zk zv (RBNode Red xk xv a (RBNode Red yk yv b c)) d ->
268 assemble col xk xv yk yv zk zv a b c d
269 RBNode col xk xv a (RBNode Red zk zv (RBNode Red yk yv b c) d) ->
270 assemble col xk xv yk yv zk zv a b c d
271 RBNode col xk xv a (RBNode Red yk yv b (RBNode Red zk zv c d)) ->
272 assemble col xk xv yk yv zk zv a b c d
274 RBNode BBlack xk xv a (RBNode NBlack zk zv (RBNode Black yk yv b c) d) ->
275 case d of
276 (RBNode Black _ _ _ _) ->
277 RBNode Black yk yv (RBNode Black xk xv a b) (balance Black zk zv c (redden d))
278 _ -> t
280 RBNode BBlack zk zv (RBNode NBlack xk xv a (RBNode Black yk yv b c)) d ->
281 case a of
282 (RBNode Black _ _ _ _) ->
283 RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
284 _ -> t
285 _ -> t
286 else t
288-- make the top node black
289blacken : Dict k v -> Dict k v
290blacken t = case t of
291 RBEmpty _ -> RBEmpty LBlack
292 RBNode _ k v l r -> RBNode Black k v l r
294-- make the top node red
295redden : Dict k v -> Dict k v
296redden t = case t of
297 RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
298 RBNode _ k v l r -> RBNode Red k v l r
300{-| Apply a function to all values in a dictionary. -}
301map : (a -> b) -> Dict comparable a -> Dict comparable b
302map f t =
303 case t of
304 RBEmpty LBlack -> RBEmpty LBlack
305 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
307{-| Fold over the key-value pairs in a dictionary, in order from lowest
308key to highest key. -}
309foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
310foldl f acc t =
311 case t of
312 RBEmpty LBlack -> acc
313 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
315{-| Fold over the key-value pairs in a dictionary, in order from highest
316key to lowest key. -}
317foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
318foldr f acc t =
319 case t of
320 RBEmpty LBlack -> acc
321 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
323{-| Combine two dictionaries. If there is a collision, preference is given
324to the first dictionary. -}
325union : Dict comparable v -> Dict comparable v -> Dict comparable v
326union t1 t2 = foldl insert t2 t1
328{-| Keep a key-value pair when its key appears in the second dictionary.
329Preference is given to values in the first dictionary. -}
330intersect : Dict comparable v -> Dict comparable v -> Dict comparable v
331intersect t1 t2 =
332 let combine k v t = if k `member` t2 then insert k v t else t
333 in foldl combine empty t1
335{-| Keep a key-value pair when its key does not appear in the second dictionary.
336Preference is given to the first dictionary. -}
337diff : Dict comparable v -> Dict comparable v -> Dict comparable v
338diff t1 t2 = foldl (\k v t -> remove k t) t1 t2
340{-| Get all of the keys in a dictionary. -}
341keys : Dict comparable v -> [comparable]
342keys t = foldr (\k v acc -> k :: acc) [] t
344{-| Get all of the values in a dictionary. -}
345values : Dict comparable v -> [v]
346values t = foldr (\k v acc -> v :: acc) [] t
348{-| Convert a dictionary into an association list of key-value pairs. -}
349toList : Dict comparable v -> [(comparable,v)]
350toList t = foldr (\k v acc -> (k,v) :: acc) [] t
352{-| Convert an association list into a dictionary. -}
353fromList : [(comparable,v)] -> Dict comparable v
354fromList assocs = List.foldl (\(k,v) d -> insert k v d) empty assocs