git.iko.soy/elk
on 727d708 help sign in
elk/libraries/Dict.elm
402 lines 12 kB history raw
2module Dict (empty,singleton,insert
3 ,lookup,findWithDefault
4 ,remove,member
5 ,foldl,foldr,map
6 ,union,intersect,diff
7 ,keys,values
8 ,toList,fromList
9 ) where
11import Maybe as Maybe
12import Native.Error as Error
13import List as List
14import Native.Utils (compare)
16data NColor = Red | Black
18data Dict k v = RBNode NColor k v (Dict k v) (Dict k v) | RBEmpty
20-- Create an empty dictionary.
21empty : Dict (Comparable k) v
22empty = RBEmpty
24{-- Helpers for checking invariants
26-- Check that the tree has an equal number of black nodes on each path
27equal_pathLen t =
28 let path_numBlacks t =
29 case t of
30 RBEmpty -> 1
31 RBNode col _ _ l r ->
32 let { bl = path_numBlacks l ; br = path_numBlacks r } in
33 if bl /= br || bl == 0-1 || br == 0-1
34 then 0-1
35 else bl + (if col == Red then 0 else 1)
36 in 0-1 /= path_numBlacks t
38rootBlack t =
39 case t of
40 RBEmpty -> True
41 RBNode Black _ _ _ _ -> True
42 _ -> False
44redBlack_children t =
45 case t of
46 { RBNode Red _ _ (RBNode Red _ _ _ _) _ -> False
47 ; RBNode Red _ _ _ (RBNode Red _ _ _ _) -> False
48 ; RBEmpty -> True
49 ; RBNode _ _ _ l r -> redBlack_children l && redBlack_children r
50 }
52findExtreme f t =
53 case t of
54 { RBEmpty -> Nothing
55 ; RBNode c k _ l r ->
56 case findExtreme f (f (l,r)) of
57 { Nothing -> Just k
58 ; Just k' -> Just k' }
59 }
61findminRbt t = findExtreme fst t
62findmaxRbt t = findExtreme snd t
64-- "Option LT than"
65-- Returns True if either xo or yo is Nothing
66-- Otherwise returns the result of comparing the values using f
68optionRelation f u xo yo =
69 case (xo,yo) of
70 { (Nothing,_) -> u
71 ; (_,Nothing) -> u
72 ; (Just x, Just y) -> f x y }
74olt xo yo = optionRelation (< ) True xo yo
75olte xo yo = optionRelation (<=) True xo yo
77ordered t =
78 case t of
79 { RBEmpty -> True
80 ; RBNode c k v l r ->
81 let (lmax,rmin) = (findmaxRbt l, findminRbt r) in
82 olte lmax (Just k) && olte (Just k) rmin && ordered l && ordered r
83 }
85-- Check that there aren't any right red nodes in the tree *)
86leftLeaning t =
87 case t of
88 { RBEmpty -> True
89 ; RBNode _ _ _ (RBNode Black _ _ _ _) (RBNode Red _ _ _ _) -> False
90 ; RBNode _ _ _ RBEmpty (RBNode Red _ _ _ _) -> False
91 ; RBNode _ _ _ l r -> (leftLeaning l) && (leftLeaning r)
92 }
94invariants_hold t =
95 ordered t && rootBlack t && redBlack_children t &&
96 equal_pathLen t && leftLeaning t
98--** End invariant helpers *****
99--}
102min : Dict k v -> (k,v)
103min t =
104 case t of
105 RBNode _ k v RBEmpty _ -> (k,v)
106 RBNode _ _ _ l _ -> min l
107 RBEmpty -> Error.raise "(min Empty) is not defined"
109{--
110max t =
111 case t of
112 { RBNode _ k v _ RBEmpty -> (k,v)
113 ; RBNode _ _ _ _ r -> max r
114 ; RBEmpty -> Error.raise "(max Empty) is not defined"
115 }
116--}
118-- Lookup the value associated with a key.
119lookup : Comparable k -> Dict (Comparable k) v -> Maybe v
120lookup k t =
121 case t of
122 RBEmpty -> Maybe.Nothing
123 RBNode _ k' v l r ->
124 case compare k k' of
125 LT -> lookup k l
126 EQ -> Maybe.Just v
127 GT -> lookup k r
129-- Find the value associated with a key. If the key is not found,
130-- return the default value.
131findWithDefault : v -> Comparable k -> Dict (Comparable k) v -> v
132findWithDefault base k t =
133 case t of
134 RBEmpty -> base
135 RBNode _ k' v l r ->
136 case compare k k' of
137 LT -> findWithDefault base k l
138 EQ -> v
139 GT -> findWithDefault base k r
141{--
142-- Find the value associated with a key. If the key is not found, there will be a runtime error.
143find k t =
144 case t of
145 { RBEmpty -> Error.raise "Key was not found in dictionary!"
146 ; RBNode _ k' v l r ->
147 case compare k k' of
148 { LT -> find k l
149 ; EQ -> v
150 ; GT -> find k r }
151 }
152--}
154-- Determine if a key is in a dictionary.
155member : Comparable k -> Dict (Comparable k) v -> Bool
156-- Does t contain k?
157member k t = Maybe.isJust $ lookup k t
159rotateLeft : Dict k v -> Dict k v
160rotateLeft t =
161 case t of
162 RBNode cy ky vy a (RBNode cz kz vz b c) -> RBNode cy kz vz (RBNode Red ky vy a b) c
163 _ -> Error.raise "rotateLeft of a node without enough children"
165-- rotateRight -- the reverse, and
166-- makes Y have Z's color, and makes Z Red.
167rotateRight : Dict k v -> Dict k v
168rotateRight t =
169 case t of
170 RBNode cz kz vz (RBNode cy ky vy a b) c -> RBNode cz ky vy a (RBNode Red kz vz b c)
171 _ -> Error.raise "rotateRight of a node without enough children"
173rotateLeftIfNeeded : Dict k v -> Dict k v
174rotateLeftIfNeeded t =
175 case t of
176 RBNode _ _ _ _ (RBNode Red _ _ _ _) -> rotateLeft t
177 _ -> t
179rotateRightIfNeeded : Dict k v -> Dict k v
180rotateRightIfNeeded t =
181 case t of
182 RBNode _ _ _ (RBNode Red _ _ (RBNode Red _ _ _ _) _) _ -> rotateRight t
183 _ -> t
185otherColor c = case c of { Red -> Black ; Black -> Red }
187color_flip : Dict k v -> Dict k v
188color_flip t =
189 case t of
190 RBNode c1 bk bv (RBNode c2 ak av la ra) (RBNode c3 ck cv lc rc) ->
191 RBNode (otherColor c1) bk bv
192 (RBNode (otherColor c2) ak av la ra)
193 (RBNode (otherColor c3) ck cv lc rc)
194 _ -> Error.raise "color_flip called on a Empty or Node with a Empty child"
196color_flipIfNeeded : Dict k v -> Dict k v
197color_flipIfNeeded t =
198 case t of
199 RBNode _ _ _ (RBNode Red _ _ _ _) (RBNode Red _ _ _ _) -> color_flip t
200 _ -> t
202fixUp t = color_flipIfNeeded (rotateRightIfNeeded (rotateLeftIfNeeded t))
204ensureBlackRoot : Dict k v -> Dict k v
205ensureBlackRoot t =
206 case t of
207 RBNode Red k v l r -> RBNode Black k v l r
208 _ -> t
210-- Insert a key-value pair into a dictionary. Replaces value when there is
211-- a collision.
212insert : Comparable k -> v -> Dict (Comparable k) v -> Dict (Comparable k) v
213insert k v t = -- Invariant: t is a valid left-leaning rb tree
214 let ins t =
215 case t of
216 RBEmpty -> RBNode Red k v RBEmpty RBEmpty
217 RBNode c k' v' l r ->
218 let h = case compare k k' of
219 LT -> RBNode c k' v' (ins l) r
220 EQ -> RBNode c k' v l r -- replace
221 GT -> RBNode c k' v' l (ins r)
222 in fixUp h
223 in ensureBlackRoot (ins t)
224{--
225 if not (invariants_hold t) then
226 Error.raise "invariants broken before insert"
227 else (let new_t = ensureBlackRoot (ins t) in
228 if not (invariants_hold new_t) then
229 Error.raise "invariants broken after insert"
230 else new_t)
231--}
233-- Create a dictionary with one key-value pair.
234singleton : Comparable k -> v -> Dict (Comparable k) v
235singleton k v = insert k v RBEmpty
237isRed : Dict k v -> Bool
238isRed t =
239 case t of
240 RBNode Red _ _ _ _ -> True
241 _ -> False
243isRedLeft : Dict k v -> Bool
244isRedLeft t =
245 case t of
246 RBNode _ _ _ (RBNode Red _ _ _ _) _ -> True
247 _ -> False
249isRedLeftLeft : Dict k v -> Bool
250isRedLeftLeft t =
251 case t of
252 RBNode _ _ _ (RBNode _ _ _ (RBNode Red _ _ _ _) _) _ -> True
253 _ -> False
255isRedRight : Dict k v -> Bool
256isRedRight t =
257 case t of
258 RBNode _ _ _ _ (RBNode Red _ _ _ _) -> True
259 _ -> False
261isRedRightLeft : Dict k v -> Bool
262isRedRightLeft t =
263 case t of
264 RBNode _ _ _ _ (RBNode _ _ _ (RBNode Red _ _ _ _) _) -> True
265 _ -> False
267moveRedLeft : Dict k v -> Dict k v
268moveRedLeft t =
269 let t' = color_flip t in
270 case t' of
271 RBNode c k v l r ->
272 case r of
273 RBNode _ _ _ (RBNode Red _ _ _ _) _ ->
274 color_flip (rotateLeft (RBNode c k v l (rotateRight r)))
275 _ -> t'
276 _ -> t'
278moveRedRight : Dict k v -> Dict k v
279moveRedRight t =
280 let t' = color_flip t in
281 if isRedLeftLeft t' then color_flip (rotateRight t') else t'
283moveRedLeftIfNeeded : Dict k v -> Dict k v
284moveRedLeftIfNeeded t =
285 if isRedLeft t || isRedLeftLeft t then t else moveRedLeft t
287moveRedRightIfNeeded : Dict k v -> Dict k v
288moveRedRightIfNeeded t =
289 if isRedRight t || isRedRightLeft t then t else moveRedRight t
291deleteMin : Dict k v -> Dict k v
292deleteMin t =
293 let del t =
294 case t of
295 RBNode _ _ _ RBEmpty _ -> RBEmpty
296 _ -> case moveRedLeftIfNeeded t of
297 RBNode c k v l r -> fixUp (RBNode c k v (del l) r)
298 RBEmpty -> RBEmpty
299 in ensureBlackRoot (del t)
301{--
302deleteMax t =
303 let del t =
304 let t' = if isRedLeft t then rotateRight t else t in
305 case t' of
306 { RBNode _ _ _ _ RBEmpty -> RBEmpty
307 ; _ -> let t'' = moveRedRightIfNeeded t' in
308 case t'' of
309 { RBNode c k v l r -> fixUp (RBNode c k v l (del r))
310 ; RBEmpty -> RBEmpty } }
311 in ensureBlackRoot (del t)
312--}
314-- Remove a key-value pair from a dictionary. If the key is not found,
315-- no changes are made.
316remove : Comparable k -> Dict (Comparable k) v -> Dict (Comparable k) v
317remove k t =
318 let eq_and_noRightNode t =
319 case t of { RBNode _ k' _ _ RBEmpty -> k == k' ; _ -> False }
320 eq t = case t of { RBNode _ k' _ _ _ -> k == k' ; _ -> False }
321 delLT t = case moveRedLeftIfNeeded t of
322 RBNode c k' v l r -> fixUp (RBNode c k' v (del l) r)
323 RBEmpty -> Error.raise "delLT on Empty"
324 delEQ t = case t of -- Replace with successor
325 RBNode c _ _ l r -> let (k',v') = min r in
326 fixUp (RBNode c k' v' l (deleteMin r))
327 RBEmpty -> Error.raise "delEQ called on a Empty"
328 delGT t = case t of
329 RBNode c k' v l r -> fixUp (RBNode c k' v l (del r))
330 RBEmpty -> Error.raise "delGT called on a Empty"
331 del t = case t of
332 RBEmpty -> RBEmpty
333 RBNode _ k' _ _ _ ->
334 if k < k' then delLT t else
335 let u = if isRedLeft t then rotateRight t else t in
336 if eq_and_noRightNode u then RBEmpty else
337 let t' = moveRedRightIfNeeded t in
338 if eq t' then delEQ t' else delGT t'
339 in if member k t then ensureBlackRoot (del t) else t
340{--
341 if not (invariants_hold t) then
342 Error.raise "invariants broken before remove"
343 else (let t' = ensureBlackRoot (del t) in
344 if invariants_hold t' then t' else
345 Error.raise "invariants broken after remove")
346--}
348-- Apply a function to all values in a dictionary.
349map : (a -> b) -> Dict (Comparable k) a -> Dict (Comparable k) b
350map f t =
351 case t of
352 RBEmpty -> RBEmpty
353 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
355-- Fold over the key-value pairs in a dictionary, in order from lowest
356-- key to highest key.
357foldl : (Comparable k -> v -> b -> b) -> b -> Dict (Comparable k) v -> b
358foldl f acc t =
359 case t of
360 RBEmpty -> acc
361 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
363-- Fold over the key-value pairs in a dictionary, in order from highest
364-- key to lowest key.
365foldr : (Comparable k -> v -> b -> b) -> b -> Dict (Comparable k) v -> b
366foldr f acc t =
367 case t of
368 RBEmpty -> acc
369 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
371-- Combine two dictionaries. If there is a collision, preference is given
372-- to the first dictionary.
373union : Dict (Comparable k) v -> Dict (Comparable k) v -> Dict (Comparable k) v
374union t1 t2 = foldl insert t2 t1
376-- Keep a key-value pair when its key appears in the second dictionary.
377-- Preference is given to values in the first dictionary.
378intersect : Dict (Comparable k) v -> Dict (Comparable k) v -> Dict (Comparable k) v
379intersect t1 t2 =
380 let combine k v t = if k `member` t2 then insert k v t else t
381 in foldl combine empty t1
383-- Keep a key-value pair when its key does not appear in the second dictionary.
384-- Preference is given to the first dictionary.
385diff : Dict (Comparable k) v -> Dict (Comparable k) v -> Dict (Comparable k) v
386diff t1 t2 = foldl (\k v t -> remove k t) t1 t2
388-- Get all of the keys in a dictionary.
389keys : Dict (Comparable k) v -> [Comparable k]
390keys t = foldr (\k v acc -> k :: acc) [] t
392-- Get all of the values in a dictionary.
393values : Dict (Comparable k) v -> [v]
394values t = foldr (\k v acc -> v :: acc) [] t
396-- Convert a dictionary into an association list of key-value pairs.
397toList : Dict (Comparable k) v -> [(Comparable k,v)]
398toList t = foldr (\k v acc -> (k,v) :: acc) [] t
400-- Convert an association list into a dictionary.
401fromList : [(Comparable k,v)] -> Dict (Comparable k) v
402fromList assocs = List.foldl (\(k,v) d -> insert k v d) empty assocs