@@ -1,4 +1,3 @@
1−
21module Dict (empty,singleton,insert
32 ,lookup,findWithDefault
43 ,remove,member
@@ -8,6 +7,10 @@ module Dict (empty,singleton,insert
87 ,toList,fromList
98 ) where
109
10+-- Cleanup: Change RBEmpty to 2 constructors: Leaf and BBLeaf
11+-- Consider making the types safer or moving all the remove methods to
12+-- a let
13+
1114{-| A dictionary mapping unique keys to values. The keys can be any comparable
1215type. This includes `Int`, `Float`, `Time`, `Char`, `String`, and tuples or
1316lists of comparable types.
@@ -38,13 +41,19 @@ import Native.Error
3841import List as List
3942import Native.Utils
4043
44+-- BBlack and NBlack should only be used during the deletion
45+-- algorithm. Any other occurrence is a bug and should fail an assert.
4146data NColor = Red | Black
47+ | BBlack | NBlack
4248
43−data Dict k v = RBNode NColor k v (Dict k v) (Dict k v) | RBEmpty
49+data LeafColor = LBlack | LBBlack
50+
51+data Dict k v = RBNode NColor k v (Dict k v) (Dict k v)
52+ | RBEmpty LeafColor
4453
4554{-| Create an empty dictionary. -}
4655empty : Dict comparable v
47−empty = RBEmpty
56+empty = RBEmpty LBlack
4857
4958{-- Helpers for checking invariants
5059
@@ -127,24 +136,22 @@ invariants_hold t =
127136min : Dict k v -> (k,v)
128137min t =
129138 case t of
130− RBNode _ k v RBEmpty _ -> (k,v)
139+ RBNode _ k v (RBEmpty LBlack) _ -> (k,v)
131140 RBNode _ _ _ l _ -> min l
132− RBEmpty -> Native.Error.raise "(min Empty) is not defined"
141+ RBEmpty LBlack -> Native.Error.raise "(min Empty) is not defined"
133142
134−{--
143+max : Dict k v -> (k, v)
135144max t =
136145 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−--}
146+ RBNode _ k v _ (RBEmpty _) -> (k,v)
147+ RBNode _ _ _ _ r -> max r
148+ RBEmpty _ -> Native.Error.raise "(max Empty) is not defined"
142149
143150{-| Lookup the value associated with a key. -}
144151lookup : comparable -> Dict comparable v -> Maybe v
145152lookup k t =
146153 case t of
147− RBEmpty -> Nothing
154+ RBEmpty LBlack -> Nothing
148155 RBNode _ k' v l r ->
149156 case Native.Utils.compare k k' of
150157 LT -> lookup k l
@@ -156,7 +163,7 @@ return the default value. -}
156163findWithDefault : v -> comparable -> Dict comparable v -> v
157164findWithDefault base k t =
158165 case t of
159− RBEmpty -> base
166+ RBEmpty LBlack -> base
160167 RBNode _ k' v l r ->
161168 case Native.Utils.compare k k' of
162169 LT -> findWithDefault base k l
@@ -230,7 +237,8 @@ ensureBlackRoot : Dict k v -> Dict k v
230237ensureBlackRoot t =
231238 case t of
232239 RBNode Red k v l r -> RBNode Black k v l r
233− _ -> t
240+ RBNode Black _ _ _ _ -> t
241+ RBEmpty LBlack -> t
234242
235243{-| Insert a key-value pair into a dictionary. Replaces value when there is
236244a collision. -}
@@ -238,7 +246,7 @@ insert : comparable -> v -> Dict comparable v -> Dict comparable v
238246insert k v t = -- Invariant: t is a valid left-leaning rb tree
239247 let ins t =
240248 case t of
241− RBEmpty -> RBNode Red k v RBEmpty RBEmpty
249+ RBEmpty LBlack -> RBNode Red k v (RBEmpty LBlack) (RBEmpty LBlack)
242250 RBNode c k' v' l r ->
243251 let h = case Native.Utils.compare k k' of
244252 LT -> RBNode c k' v' (ins l) r
@@ -257,71 +265,17 @@ insert k v t = -- Invariant: t is a valid left-leaning rb tree
257265
258266{-| Create a dictionary with one key-value pair. -}
259267singleton : 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
268+singleton k v = insert k v (RBEmpty LBlack)
315269
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)
270+-- deleteMin : Dict k v -> Dict k v
271+-- deleteMin t =
272+-- let del t =
273+-- case t of
274+-- RBNode _ _ _ RBEmpty _ -> RBEmpty
275+-- _ -> case moveRedLeftIfNeeded t of
276+-- RBNode c k v l r -> fixUp (RBNode c k v (del l) r)
277+-- RBEmpty -> RBEmpty
278+-- in ensureBlackRoot (del t)
325279
326280{--
327281deleteMax t =
@@ -336,32 +290,133 @@ deleteMax t =
336290 in ensureBlackRoot (del t)
337291--}
338292
293+-- Finds and deletes k in t
294+del : comparable -> Dict comparable v -> Dict comparable v
295+del k t = case t of
296+ RBEmpty _ -> t
297+ RBNode c k' v l r -> case Native.Utils.compare k k' of
298+ LT -> bubble c k' v (del k l) r
299+ EQ -> rem t
300+ GT -> bubble c k' v l (del k r)
301+
302+-- Remove the top node from the tree, may leave behind BBlacks
303+rem : Dict k v -> Dict k v
304+rem t = case t of
305+ RBNode c k v (RBEmpty _) (RBEmpty _) -> case c of
306+ Red -> RBEmpty LBlack
307+ Black -> RBEmpty LBBlack
308+ RBNode Black _ _ (RBEmpty _) (RBNode _ k v l r) ->
309+ RBNode Black k v l r
310+ RBNode Black _ _ (RBNode _ k v l r) (RBEmpty _) ->
311+ RBNode Black k v l r
312+ -- l and r are both RBNodes
313+ RBNode c _ _ l r ->
314+ let (k, v) = max l
315+ l' = remove_max l
316+ in bubble c k v l' r
317+
318+isBBlack : Dict k v -> Bool
319+isBBlack t = case t of
320+ RBNode c _ _ _ _ -> case c of
321+ BBlack -> True
322+ _ -> False
323+ RBEmpty LBBlack -> True
324+ _ -> False
325+
326+moreBlack : NColor -> NColor
327+moreBlack c = case c of
328+ BBlack -> BBlack
329+ Black -> BBlack
330+ Red -> Black
331+ NBlack -> Red
332+lessBlack : NColor -> NColor
333+lessBlack c = case c of
334+ BBlack -> Black
335+ Black -> Red
336+ Red -> NBlack
337+ NBlack -> NBlack
338+
339+moreBlackTree : Dict k v -> Dict k v
340+moreBlackTree t = case t of
341+ RBNode c k v l r -> RBNode (moreBlack c) k v l r
342+ RBEmpty _ -> RBEmpty LBBlack
343+
344+lessBlackTree : Dict k v -> Dict k v
345+lessBlackTree t = case t of
346+ RBNode c k v l r -> RBNode (lessBlack c) k v l r
347+ RBEmpty _ -> RBEmpty LBlack
348+
349+-- Kills a BBlack or moves it upward, may leave behind NBlack
350+bubble : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
351+bubble c k v l r = if isBBlack l || isBBlack r
352+ then balance (moreBlack c) k v (lessBlackTree l) (lessBlackTree r)
353+ else RBNode c k v l r
354+
355+-- Removes rightmost node, may leave root as BBlack
356+remove_max : Dict k v -> Dict k v
357+remove_max t = case t of
358+ RBNode c k v l (RBEmpty _) -> rem t
359+ RBNode c k v l r -> bubble c k v l (remove_max r)
360+
361+-- generalized tree balancing act
362+balance : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
363+balance c k v l r = balance_node (RBNode c k v l r)
364+
365+blackish : Dict k v -> Bool
366+blackish (RBNode c _ _ _ _) = c == Black || c == BBlack
367+
368+balance_node : Dict k v -> Dict k v
369+balance_node t =
370+ let assemble col xk xv yk yv zk zv a b c d =
371+ RBNode (lessBlack col) yk yv (RBNode Black xk xv a b) (RBNode Black zk zv c d)
372+ in
373+ if blackish t
374+ then case t of
375+ RBNode col zk zv (RBNode Red yk yv (RBNode Red xk xv a b) c) d ->
376+ assemble col xk xv yk yv zk zv a b c d
377+ RBNode col zk zv (RBNode Red xk xv a (RBNode Red yk yv b c)) d ->
378+ assemble col xk xv yk yv zk zv a b c d
379+ RBNode col xk xv a (RBNode Red zk zv (RBNode Red yk yv b c) d) ->
380+ assemble col xk xv yk yv zk zv a b c d
381+ RBNode col xk xv a (RBNode Red yk yv b (RBNode Red zk zv c d)) ->
382+ assemble col xk xv yk yv zk zv a b c d
383+
384+ RBNode BBlack xk xv a (RBNode NBlack zk zv (RBNode Black yk yv b c) d) ->
385+ case d of
386+ (RBNode Black _ _ _ _) ->
387+ RBNode Black yk yv (RBNode Black xk xv a b) (balance Black zk zv c (redden d))
388+ _ -> t
389+
390+ RBNode BBlack zk zv (RBNode NBlack xk xv a (RBNode Black yk yv b c)) d ->
391+ case a of
392+ (RBNode Black _ _ _ _) ->
393+ RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
394+ _ -> t
395+
396+ _ -> t
397+
398+ else t
399+
400+
401+
402+
403+-- make the top node black
404+blacken : Dict k v -> Dict k v
405+blacken t = case t of
406+ RBEmpty _ -> RBEmpty LBlack
407+ RBNode _ k v l r -> RBNode Black k v l r
408+
409+-- make the top node red
410+redden : Dict k v -> Dict k v
411+redden t = case t of
412+ RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
413+ RBNode _ k v l r -> RBNode Red k v l r
414+
339415{-| Remove a key-value pair from a dictionary. If the key is not found,
340416no changes are made. -}
341417remove : 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
418+remove k t = blacken <| del k t
419+
365420{--
366421 if not (invariants_hold t) then
367422 Native.Error.raise "invariants broken before remove"
@@ -374,7 +429,7 @@ remove k t =
374429map : (a -> b) -> Dict comparable a -> Dict comparable b
375430map f t =
376431 case t of
377− RBEmpty -> RBEmpty
432+ RBEmpty LBlack -> RBEmpty LBlack
378433 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
379434
380435{-| Fold over the key-value pairs in a dictionary, in order from lowest
@@ -382,7 +437,7 @@ key to highest key. -}
382437foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
383438foldl f acc t =
384439 case t of
385− RBEmpty -> acc
440+ RBEmpty LBlack -> acc
386441 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
387442
388443{-| Fold over the key-value pairs in a dictionary, in order from highest
@@ -390,7 +445,7 @@ key to lowest key. -}
390445foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
391446foldr f acc t =
392447 case t of
393− RBEmpty -> acc
448+ RBEmpty LBlack -> acc
394449 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
395450
396451{-| Combine two dictionaries. If there is a collision, preference is given