@@ -1,4 +1,3 @@
1−
21module Dict (empty,singleton,insert
32 ,lookup,findWithDefault
43 ,remove,member
@@ -38,113 +37,39 @@ import Native.Error
3837import List as List
3938import Native.Utils
4039
40+-- BBlack and NBlack should only be used during the deletion
41+-- algorithm. Any other occurrence is a bug and should fail an assert.
4142data NColor = Red | Black
43+ | BBlack | NBlack
44+
45+data LeafColor = LBlack | LBBlack
4246
43−data Dict k v = RBNode NColor k v (Dict k v) (Dict k v) | RBEmpty
47+data Dict k v = RBNode NColor k v (Dict k v) (Dict k v)
48+ | RBEmpty LeafColor
4449
4550{-| Create an empty dictionary. -}
4651empty : Dict comparable v
47−empty = RBEmpty
48−
49−{-- Helpers for checking invariants
50−
51−-- Check that the tree has an equal number of black nodes on each path
52−equal_pathLen t =
53− let path_numBlacks t =
54− case t of
55− RBEmpty -> 1
56− RBNode col _ _ l r ->
57− let { bl = path_numBlacks l ; br = path_numBlacks r } in
58− if bl /= br || bl == 0-1 || br == 0-1
59− then 0-1
60− else bl + (if col == Red then 0 else 1)
61− in 0-1 /= path_numBlacks t
62−
63−rootBlack t =
64− case t of
65− RBEmpty -> True
66− RBNode Black _ _ _ _ -> True
67− _ -> False
68−
69−redBlack_children t =
70− case t of
71− { RBNode Red _ _ (RBNode Red _ _ _ _) _ -> False
72− ; RBNode Red _ _ _ (RBNode Red _ _ _ _) -> False
73− ; RBEmpty -> True
74− ; RBNode _ _ _ l r -> redBlack_children l && redBlack_children r
75− }
76−
77−findExtreme f t =
78− case t of
79− { RBEmpty -> Nothing
80− ; RBNode c k _ l r ->
81− case findExtreme f (f (l,r)) of
82− { Nothing -> Just k
83− ; Just k' -> Just k' }
84− }
85−
86−findminRbt t = findExtreme fst t
87−findmaxRbt t = findExtreme snd t
88−
89−-- "Option LT than"
90−-- Returns True if either xo or yo is Nothing
91−-- Otherwise returns the result of comparing the values using f
92−
93−optionRelation f u xo yo =
94− case (xo,yo) of
95− { (Nothing,_) -> u
96− ; (_,Nothing) -> u
97− ; (Just x, Just y) -> f x y }
98−
99−olt xo yo = optionRelation (< ) True xo yo
100−olte xo yo = optionRelation (<=) True xo yo
101−
102−ordered t =
103− case t of
104− { RBEmpty -> True
105− ; RBNode c k v l r ->
106− let (lmax,rmin) = (findmaxRbt l, findminRbt r) in
107− olte lmax (Just k) && olte (Just k) rmin && ordered l && ordered r
108− }
109−
110−-- Check that there aren't any right red nodes in the tree *)
111−leftLeaning t =
112− case t of
113− { RBEmpty -> True
114− ; RBNode _ _ _ (RBNode Black _ _ _ _) (RBNode Red _ _ _ _) -> False
115− ; RBNode _ _ _ RBEmpty (RBNode Red _ _ _ _) -> False
116− ; RBNode _ _ _ l r -> (leftLeaning l) && (leftLeaning r)
117− }
118−
119−invariants_hold t =
120− ordered t && rootBlack t && redBlack_children t &&
121− equal_pathLen t && leftLeaning t
122−
123−--** End invariant helpers *****
124−--}
125−
52+empty = RBEmpty LBlack
12653
12754min : Dict k v -> (k,v)
12855min t =
12956 case t of
130− RBNode _ k v RBEmpty _ -> (k,v)
57+ RBNode _ k v (RBEmpty LBlack) _ -> (k,v)
13158 RBNode _ _ _ l _ -> min l
132− RBEmpty -> Native.Error.raise "(min Empty) is not defined"
59+ RBEmpty LBlack -> Native.Error.raise "(min Empty) is not defined"
13360
134−{--
61+max : Dict k v -> (k, v)
13562max t =
13663 case t of
137− { RBNode _ k v _ RBEmpty -> (k,v)
138− ; RBNode _ _ _ _ r -> max r
139− ; RBEmpty -> Native.Error.raise "(max Empty) is not defined"
140− }
141−--}
64+ RBNode _ k v _ (RBEmpty _) -> (k,v)
65+ RBNode _ _ _ _ r -> max r
66+ RBEmpty _ -> Native.Error.raise "(max Empty) is not defined"
14267
14368{-| Lookup the value associated with a key. -}
14469lookup : comparable -> Dict comparable v -> Maybe v
14570lookup k t =
14671 case t of
147− RBEmpty -> Nothing
72+ RBEmpty LBlack -> Nothing
14873 RBNode _ k' v l r ->
14974 case Native.Utils.compare k k' of
15075 LT -> lookup k l
@@ -156,26 +81,13 @@ return the default value. -}
15681findWithDefault : v -> comparable -> Dict comparable v -> v
15782findWithDefault base k t =
15883 case t of
159− RBEmpty -> base
84+ RBEmpty LBlack -> base
16085 RBNode _ k' v l r ->
16186 case Native.Utils.compare k k' of
16287 LT -> findWithDefault base k l
16388 EQ -> v
16489 GT -> findWithDefault base k r
16590
166−{--
167−-- Find the value associated with a key. If the key is not found, there will be a runtime error.
168−find k t =
169− case t of
170− { RBEmpty -> Native.Error.raise "Key was not found in dictionary!"
171− ; RBNode _ k' v l r ->
172− case Native.Utils.compare k k' of
173− { LT -> find k l
174− ; EQ -> v
175− ; GT -> find k r }
176− }
177−--}
178−
17991{-| Determine if a key is in a dictionary. -}
18092member : comparable -> Dict comparable v -> Bool
18193-- Does t contain k?
@@ -230,7 +142,8 @@ ensureBlackRoot : Dict k v -> Dict k v
230142ensureBlackRoot t =
231143 case t of
232144 RBNode Red k v l r -> RBNode Black k v l r
233− _ -> t
145+ RBNode Black _ _ _ _ -> t
146+ RBEmpty LBlack -> t
234147
235148{-| Insert a key-value pair into a dictionary. Replaces value when there is
236149a collision. -}
@@ -238,7 +151,7 @@ insert : comparable -> v -> Dict comparable v -> Dict comparable v
238151insert k v t = -- Invariant: t is a valid left-leaning rb tree
239152 let ins t =
240153 case t of
241− RBEmpty -> RBNode Red k v RBEmpty RBEmpty
154+ RBEmpty LBlack -> RBNode Red k v (RBEmpty LBlack) (RBEmpty LBlack)
242155 RBNode c k' v' l r ->
243156 let h = case Native.Utils.compare k k' of
244157 LT -> RBNode c k' v' (ins l) r
@@ -246,135 +159,143 @@ insert k v t = -- Invariant: t is a valid left-leaning rb tree
246159 GT -> RBNode c k' v' l (ins r)
247160 in fixUp h
248161 in ensureBlackRoot (ins t)
249−{--
250− if not (invariants_hold t) then
251− Native.Error.raise "invariants broken before insert"
252− else (let new_t = ensureBlackRoot (ins t) in
253− if not (invariants_hold new_t) then
254− Native.Error.raise "invariants broken after insert"
255− else new_t)
256−--}
257162
258163{-| Create a dictionary with one key-value pair. -}
259164singleton : comparable -> v -> Dict comparable v
260−singleton k v = insert k v RBEmpty
261−
262−isRed : Dict k v -> Bool
263−isRed t =
264− case t of
265− RBNode Red _ _ _ _ -> True
266− _ -> False
267−
268−isRedLeft : Dict k v -> Bool
269−isRedLeft t =
270− case t of
271− RBNode _ _ _ (RBNode Red _ _ _ _) _ -> True
272− _ -> False
273−
274−isRedLeftLeft : Dict k v -> Bool
275−isRedLeftLeft t =
276− case t of
277− RBNode _ _ _ (RBNode _ _ _ (RBNode Red _ _ _ _) _) _ -> True
278− _ -> False
279−
280−isRedRight : Dict k v -> Bool
281−isRedRight t =
282− case t of
283− RBNode _ _ _ _ (RBNode Red _ _ _ _) -> True
284− _ -> False
285−
286−isRedRightLeft : Dict k v -> Bool
287−isRedRightLeft t =
288− case t of
289− RBNode _ _ _ _ (RBNode _ _ _ (RBNode Red _ _ _ _) _) -> True
290− _ -> False
291−
292−moveRedLeft : Dict k v -> Dict k v
293−moveRedLeft t =
294− let t' = color_flip t in
295− case t' of
296− RBNode c k v l r ->
297− case r of
298− RBNode _ _ _ (RBNode Red _ _ _ _) _ ->
299− color_flip (rotateLeft (RBNode c k v l (rotateRight r)))
300− _ -> t'
301− _ -> t'
302−
303−moveRedRight : Dict k v -> Dict k v
304−moveRedRight t =
305− let t' = color_flip t in
306− if isRedLeftLeft t' then color_flip (rotateRight t') else t'
307−
308−moveRedLeftIfNeeded : Dict k v -> Dict k v
309−moveRedLeftIfNeeded t =
310− if isRedLeft t || isRedLeftLeft t then t else moveRedLeft t
311−
312−moveRedRightIfNeeded : Dict k v -> Dict k v
313−moveRedRightIfNeeded t =
314− if isRedRight t || isRedRightLeft t then t else moveRedRight t
315−
316−deleteMin : Dict k v -> Dict k v
317−deleteMin t =
318− let del t =
319− case t of
320− RBNode _ _ _ RBEmpty _ -> RBEmpty
321− _ -> case moveRedLeftIfNeeded t of
322− RBNode c k v l r -> fixUp (RBNode c k v (del l) r)
323− RBEmpty -> RBEmpty
324− in ensureBlackRoot (del t)
325−
326−{--
327−deleteMax t =
328− let del t =
329− let t' = if isRedLeft t then rotateRight t else t in
330− case t' of
331− { RBNode _ _ _ _ RBEmpty -> RBEmpty
332− ; _ -> let t'' = moveRedRightIfNeeded t' in
333− case t'' of
334− { RBNode c k v l r -> fixUp (RBNode c k v l (del r))
335− ; RBEmpty -> RBEmpty } }
336− in ensureBlackRoot (del t)
337−--}
165+singleton k v = insert k v (RBEmpty LBlack)
166+
167+{- Remove helpers: everything from here to remove should only be used
168+ internally by remove as they would otherwise break rb-invariants -}
169+
170+isBBlack : Dict k v -> Bool
171+isBBlack t = case t of
172+ RBNode c _ _ _ _ -> case c of
173+ BBlack -> True
174+ _ -> False
175+ RBEmpty LBBlack -> True
176+ _ -> False
177+
178+moreBlack : NColor -> NColor
179+moreBlack c = case c of
180+ BBlack -> BBlack
181+ Black -> BBlack
182+ Red -> Black
183+ NBlack -> Red
184+lessBlack : NColor -> NColor
185+lessBlack c = case c of
186+ BBlack -> Black
187+ Black -> Red
188+ Red -> NBlack
189+ NBlack -> NBlack
190+
191+moreBlackTree : Dict k v -> Dict k v
192+moreBlackTree t = case t of
193+ RBNode c k v l r -> RBNode (moreBlack c) k v l r
194+ RBEmpty _ -> RBEmpty LBBlack
195+
196+lessBlackTree : Dict k v -> Dict k v
197+lessBlackTree t = case t of
198+ RBNode c k v l r -> RBNode (lessBlack c) k v l r
199+ RBEmpty _ -> RBEmpty LBlack
200+
201+-- Finds and deletes k in t
202+del : comparable -> Dict comparable v -> Dict comparable v
203+del k t = case t of
204+ RBEmpty _ -> t
205+ RBNode c k' v l r -> case Native.Utils.compare k k' of
206+ LT -> bubble c k' v (del k l) r
207+ EQ -> rem t
208+ GT -> bubble c k' v l (del k r)
209+
210+-- Remove the top node from the tree, may leave behind BBlacks
211+rem : Dict k v -> Dict k v
212+rem t = case t of
213+ RBNode c k v (RBEmpty _) (RBEmpty _) -> case c of
214+ Red -> RBEmpty LBlack
215+ Black -> RBEmpty LBBlack
216+ RBNode Black _ _ (RBEmpty _) (RBNode _ k v l r) ->
217+ RBNode Black k v l r
218+ RBNode Black _ _ (RBNode _ k v l r) (RBEmpty _) ->
219+ RBNode Black k v l r
220+ -- l and r are both RBNodes
221+ RBNode c _ _ l r ->
222+ let (k, v) = max l
223+ l' = remove_max l
224+ in bubble c k v l' r
225+
226+-- Kills a BBlack or moves it upward, may leave behind NBlack
227+bubble : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
228+bubble c k v l r = if isBBlack l || isBBlack r
229+ then balance (moreBlack c) k v (lessBlackTree l) (lessBlackTree r)
230+ else RBNode c k v l r
231+
232+-- Removes rightmost node, may leave root as BBlack
233+remove_max : Dict k v -> Dict k v
234+remove_max t = case t of
235+ RBNode c k v l (RBEmpty _) -> rem t
236+ RBNode c k v l r -> bubble c k v l (remove_max r)
237+
238+-- generalized tree balancing act
239+balance : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
240+balance c k v l r = balance_node (RBNode c k v l r)
241+
242+blackish : Dict k v -> Bool
243+blackish (RBNode c _ _ _ _) = c == Black || c == BBlack
244+
245+balance_node : Dict k v -> Dict k v
246+balance_node t =
247+ let assemble col xk xv yk yv zk zv a b c d =
248+ RBNode (lessBlack col) yk yv (RBNode Black xk xv a b) (RBNode Black zk zv c d)
249+ in
250+ if blackish t
251+ then case t of
252+ RBNode col zk zv (RBNode Red yk yv (RBNode Red xk xv a b) c) d ->
253+ assemble col xk xv yk yv zk zv a b c d
254+ RBNode col zk zv (RBNode Red xk xv a (RBNode Red yk yv b c)) d ->
255+ assemble col xk xv yk yv zk zv a b c d
256+ RBNode col xk xv a (RBNode Red zk zv (RBNode Red yk yv b c) d) ->
257+ assemble col xk xv yk yv zk zv a b c d
258+ RBNode col xk xv a (RBNode Red yk yv b (RBNode Red zk zv c d)) ->
259+ assemble col xk xv yk yv zk zv a b c d
260+
261+ RBNode BBlack xk xv a (RBNode NBlack zk zv (RBNode Black yk yv b c) d) ->
262+ case d of
263+ (RBNode Black _ _ _ _) ->
264+ RBNode Black yk yv (RBNode Black xk xv a b) (balance Black zk zv c (redden d))
265+ _ -> t
266+
267+ RBNode BBlack zk zv (RBNode NBlack xk xv a (RBNode Black yk yv b c)) d ->
268+ case a of
269+ (RBNode Black _ _ _ _) ->
270+ RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
271+ _ -> t
272+
273+ _ -> t
274+
275+ else t
276+
277+-- make the top node black
278+blacken : Dict k v -> Dict k v
279+blacken t = case t of
280+ RBEmpty _ -> RBEmpty LBlack
281+ RBNode _ k v l r -> RBNode Black k v l r
282+
283+-- make the top node red
284+redden : Dict k v -> Dict k v
285+redden t = case t of
286+ RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
287+ RBNode _ k v l r -> RBNode Red k v l r
338288
339289{-| Remove a key-value pair from a dictionary. If the key is not found,
340290no changes are made. -}
341291remove : comparable -> Dict comparable v -> Dict comparable v
342−remove k t =
343− let eq_and_noRightNode t =
344− case t of { RBNode _ k' _ _ RBEmpty -> k == k' ; _ -> False }
345− eq t = case t of { RBNode _ k' _ _ _ -> k == k' ; _ -> False }
346− delLT t = case moveRedLeftIfNeeded t of
347− RBNode c k' v l r -> fixUp (RBNode c k' v (del l) r)
348− RBEmpty -> Native.Error.raise "delLT on Empty"
349− delEQ t = case t of -- Replace with successor
350− RBNode c _ _ l r -> let (k',v') = min r in
351− fixUp (RBNode c k' v' l (deleteMin r))
352− RBEmpty -> Native.Error.raise "delEQ called on a Empty"
353− delGT t = case t of
354− RBNode c k' v l r -> fixUp (RBNode c k' v l (del r))
355− RBEmpty -> Native.Error.raise "delGT called on a Empty"
356− del t = case t of
357− RBEmpty -> RBEmpty
358− RBNode _ k' _ _ _ ->
359− if k < k' then delLT t else
360− let u = if isRedLeft t then rotateRight t else t in
361− if eq_and_noRightNode u then RBEmpty else
362− let t' = moveRedRightIfNeeded t in
363− if eq t' then delEQ t' else delGT t'
364− in if member k t then ensureBlackRoot (del t) else t
365−{--
366− if not (invariants_hold t) then
367− Native.Error.raise "invariants broken before remove"
368− else (let t' = ensureBlackRoot (del t) in
369− if invariants_hold t' then t' else
370− Native.Error.raise "invariants broken after remove")
371−--}
292+remove k t = blacken <| del k t
372293
373294{-| Apply a function to all values in a dictionary. -}
374295map : (a -> b) -> Dict comparable a -> Dict comparable b
375296map f t =
376297 case t of
377− RBEmpty -> RBEmpty
298+ RBEmpty LBlack -> RBEmpty LBlack
378299 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
379300
380301{-| Fold over the key-value pairs in a dictionary, in order from lowest
@@ -382,7 +303,7 @@ key to highest key. -}
382303foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
383304foldl f acc t =
384305 case t of
385− RBEmpty -> acc
306+ RBEmpty LBlack -> acc
386307 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
387308
388309{-| Fold over the key-value pairs in a dictionary, in order from highest
@@ -390,7 +311,7 @@ key to lowest key. -}
390311foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
391312foldr f acc t =
392313 case t of
393− RBEmpty -> acc
314+ RBEmpty LBlack -> acc
394315 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
395316
396317{-| Combine two dictionaries. If there is a collision, preference is given