@@ -1,4 +1,3 @@
2module Dict (empty,singleton,insert
1module Dict (empty,singleton,insert
3 ,lookup,findWithDefault
2 ,lookup,findWithDefault
4 ,remove,member
3 ,remove,member
@@ -38,113 +37,39 @@ import Native.Error
38import List as List
37import List as List
39import Native.Utils
38import Native.Utils
40+-- BBlack and NBlack should only be used during the deletion
41+-- algorithm. Any other occurrence is a bug and should fail an assert.
41data NColor = Red | Black
42data NColor = Red | Black
45+data LeafColor = LBlack | LBBlack
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)
45{-| Create an empty dictionary. -}
50{-| Create an empty dictionary. -}
46empty : Dict comparable v
51empty : Dict comparable v
47−empty = RBEmpty
52+empty = RBEmpty LBlack
49−{-- Helpers for checking invariants
51−-- Check that the tree has an equal number of black nodes on each path
53− let path_numBlacks t =
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
60− else bl + (if col == Red then 0 else 1)
61− in 0-1 /= path_numBlacks t
66− RBNode Black _ _ _ _ -> True
71− { RBNode Red _ _ (RBNode Red _ _ _ _) _ -> False
72− ; RBNode Red _ _ _ (RBNode Red _ _ _ _) -> False
74− ; RBNode _ _ _ l r -> redBlack_children l && redBlack_children r
80− ; RBNode c k _ l r ->
81− case findExtreme f (f (l,r)) of
83− ; Just k' -> Just k' }
86−findminRbt t = findExtreme fst t
87−findmaxRbt t = findExtreme snd t
90−-- Returns True if either xo or yo is Nothing
91−-- Otherwise returns the result of comparing the values using f
93−optionRelation f u xo yo =
97− ; (Just x, Just y) -> f x y }
99−olt xo yo = optionRelation (< ) True xo yo
100−olte xo yo = optionRelation (<=) True xo yo
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
110−-- Check that there aren't any right red nodes in the tree *)
114− ; RBNode _ _ _ (RBNode Black _ _ _ _) (RBNode Red _ _ _ _) -> False
115− ; RBNode _ _ _ RBEmpty (RBNode Red _ _ _ _) -> False
116− ; RBNode _ _ _ l r -> (leftLeaning l) && (leftLeaning r)
120− ordered t && rootBlack t && redBlack_children t &&
121− equal_pathLen t && leftLeaning t
123−--** End invariant helpers *****
127min : Dict k v -> (k,v)
54min : Dict k v -> (k,v)
129 case t of
56 case t of
130− RBNode _ k v RBEmpty _ -> (k,v)
57+ RBNode _ k v (RBEmpty LBlack) _ -> (k,v)
131 RBNode _ _ _ l _ -> min l
58 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"
134−{--
61+max : Dict k v -> (k, v)
136 case t of
63 case t of
137− { RBNode _ k v _ RBEmpty -> (k,v)
64+ RBNode _ k v _ (RBEmpty _) -> (k,v)
138− ; RBNode _ _ _ _ r -> max r
65+ RBNode _ _ _ _ r -> max r
139− ; RBEmpty -> Native.Error.raise "(max Empty) is not defined"
66+ RBEmpty _ -> Native.Error.raise "(max Empty) is not defined"
143{-| Lookup the value associated with a key. -}
68{-| Lookup the value associated with a key. -}
144lookup : comparable -> Dict comparable v -> Maybe v
69lookup : comparable -> Dict comparable v -> Maybe v
145lookup k t =
70lookup k t =
146 case t of
71 case t of
147− RBEmpty -> Nothing
72+ RBEmpty LBlack -> Nothing
148 RBNode _ k' v l r ->
73 RBNode _ k' v l r ->
149 case Native.Utils.compare k k' of
74 case Native.Utils.compare k k' of
150 LT -> lookup k l
75 LT -> lookup k l
@@ -156,26 +81,13 @@ return the default value. -}
156findWithDefault : v -> comparable -> Dict comparable v -> v
81findWithDefault : v -> comparable -> Dict comparable v -> v
157findWithDefault base k t =
82findWithDefault base k t =
158 case t of
83 case t of
159− RBEmpty -> base
84+ RBEmpty LBlack -> base
160 RBNode _ k' v l r ->
85 RBNode _ k' v l r ->
161 case Native.Utils.compare k k' of
86 case Native.Utils.compare k k' of
162 LT -> findWithDefault base k l
87 LT -> findWithDefault base k l
164 GT -> findWithDefault base k r
89 GT -> findWithDefault base k r
167−-- Find the value associated with a key. If the key is not found, there will be a runtime error.
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
179{-| Determine if a key is in a dictionary. -}
91{-| Determine if a key is in a dictionary. -}
180member : comparable -> Dict comparable v -> Bool
92member : comparable -> Dict comparable v -> Bool
181-- Does t contain k?
93-- Does t contain k?
@@ -230,7 +142,8 @@ ensureBlackRoot : Dict k v -> Dict k v
230ensureBlackRoot t =
142ensureBlackRoot t =
231 case t of
143 case t of
232 RBNode Red k v l r -> RBNode Black k v l r
144 RBNode Red k v l r -> RBNode Black k v l r
233− _ -> t
145+ RBNode Black _ _ _ _ -> t
235{-| Insert a key-value pair into a dictionary. Replaces value when there is
148{-| Insert a key-value pair into a dictionary. Replaces value when there is
236a collision. -}
149a collision. -}
@@ -238,7 +151,7 @@ insert : comparable -> v -> Dict comparable v -> Dict comparable v
238insert k v t = -- Invariant: t is a valid left-leaning rb tree
151insert k v t = -- Invariant: t is a valid left-leaning rb tree
239 let ins t =
152 let ins t =
240 case t of
153 case t of
241− RBEmpty -> RBNode Red k v RBEmpty RBEmpty
154+ RBEmpty LBlack -> RBNode Red k v (RBEmpty LBlack) (RBEmpty LBlack)
242 RBNode c k' v' l r ->
155 RBNode c k' v' l r ->
243 let h = case Native.Utils.compare k k' of
156 let h = case Native.Utils.compare k k' of
244 LT -> RBNode c k' v' (ins l) r
157 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
246 GT -> RBNode c k' v' l (ins r)
159 GT -> RBNode c k' v' l (ins r)
247 in fixUp h
160 in fixUp h
248 in ensureBlackRoot (ins t)
161 in ensureBlackRoot (ins t)
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"
258{-| Create a dictionary with one key-value pair. -}
163{-| Create a dictionary with one key-value pair. -}
259singleton : comparable -> v -> Dict comparable v
164singleton : comparable -> v -> Dict comparable v
260−singleton k v = insert k v RBEmpty
165+singleton k v = insert k v (RBEmpty LBlack)
262−isRed : Dict k v -> Bool
167+{- Remove helpers: everything from here to remove should only be used
263−isRed t =
168+ internally by remove as they would otherwise break rb-invariants -}
265− RBNode Red _ _ _ _ -> True
268−isRedLeft : Dict k v -> Bool
170+isBBlack : Dict k v -> Bool
270− case t of
171+isBBlack t = case t of
271− RBNode _ _ _ (RBNode Red _ _ _ _) _ -> True
172+ RBNode c _ _ _ _ -> case c of
175+ RBEmpty LBBlack -> True
272− _ -> False
176+ _ -> False
274−isRedLeftLeft : Dict k v -> Bool
178+moreBlack : NColor -> NColor
276− case t of
179+moreBlack c = case c of
277− RBNode _ _ _ (RBNode _ _ _ (RBNode Red _ _ _ _) _) _ -> True
180+ BBlack -> BBlack
278− _ -> False
181+ Black -> BBlack
280−isRedRight : Dict k v -> Bool
183+ NBlack -> Red
281−isRedRight t =
184+lessBlack : NColor -> NColor
185+lessBlack c = case c of
191+moreBlackTree : Dict k v -> Dict k v
282− case t of
192+moreBlackTree t = case t of
283− RBNode _ _ _ _ (RBNode Red _ _ _ _) -> True
193+ RBNode c k v l r -> RBNode (moreBlack c) k v l r
284− _ -> False
194+ RBEmpty _ -> RBEmpty LBBlack
286−isRedRightLeft : Dict k v -> Bool
196+lessBlackTree : Dict k v -> Dict k v
288− case t of
197+lessBlackTree t = case t of
289− RBNode _ _ _ _ (RBNode _ _ _ (RBNode Red _ _ _ _) _) -> True
198+ RBNode c k v l r -> RBNode (lessBlack c) k v l r
290− _ -> False
199+ RBEmpty _ -> RBEmpty LBlack
292−moveRedLeft : Dict k v -> Dict k v
201+-- Finds and deletes k in t
293−moveRedLeft t =
202+del : comparable -> Dict comparable v -> Dict comparable v
294− let t' = color_flip t in
203+del k t = case t of
295− case t' of
204+ RBEmpty _ -> t
296− RBNode c k v l r ->
205+ RBNode c k' v l r -> case Native.Utils.compare k k' of
297− case r of
206+ LT -> bubble c k' v (del k l) r
298− RBNode _ _ _ (RBNode Red _ _ _ _) _ ->
207+ EQ -> rem t
299− color_flip (rotateLeft (RBNode c k v l (rotateRight r)))
208+ GT -> bubble c k' v l (del k r)
301− _ -> t'
210+-- Remove the top node from the tree, may leave behind BBlacks
303−moveRedRight : Dict k v -> Dict k v
211+rem : Dict k v -> Dict k v
304−moveRedRight t =
212+rem t = case t of
305− let t' = color_flip t in
213+ RBNode c k v (RBEmpty _) (RBEmpty _) -> case c of
306− if isRedLeftLeft t' then color_flip (rotateRight t') else t'
214+ Red -> RBEmpty LBlack
307−
215+ Black -> RBEmpty LBBlack
308−moveRedLeftIfNeeded : Dict k v -> Dict k v
216+ RBNode Black _ _ (RBEmpty _) (RBNode _ k v l r) ->
309−moveRedLeftIfNeeded t =
217+ RBNode Black k v l r
310− if isRedLeft t || isRedLeftLeft t then t else moveRedLeft t
218+ RBNode Black _ _ (RBNode _ k v l r) (RBEmpty _) ->
311−
219+ RBNode Black k v l r
312−moveRedRightIfNeeded : Dict k v -> Dict k v
220+ -- l and r are both RBNodes
313−moveRedRightIfNeeded t =
221+ RBNode c _ _ l r ->
314− if isRedRight t || isRedRightLeft t then t else moveRedRight t
222+ let (k, v) = max l
315−
223+ l' = remove_max l
224+ in bubble c k v l' r
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
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)
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)
242+blackish : Dict k v -> Bool
243+blackish (RBNode c _ _ _ _) = c == Black || c == BBlack
245+balance_node : Dict k v -> Dict k v
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)
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
261+ RBNode BBlack xk xv a (RBNode NBlack zk zv (RBNode Black yk yv b c) d) ->
263+ (RBNode Black _ _ _ _) ->
264+ RBNode Black yk yv (RBNode Black xk xv a b) (balance Black zk zv c (redden d))
267+ RBNode BBlack zk zv (RBNode NBlack xk xv a (RBNode Black yk yv b c)) d ->
269+ (RBNode Black _ _ _ _) ->
270+ RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
277+-- make the top node black
316−deleteMin : Dict k v -> Dict k v
278+blacken : Dict k v -> Dict k v
319− case t of
279+blacken t = case t of
320− RBNode _ _ _ RBEmpty _ -> RBEmpty
280+ RBEmpty _ -> RBEmpty LBlack
321− _ -> case moveRedLeftIfNeeded t of
322− RBNode c k v l r -> fixUp (RBNode c k v (del l) r)
281+ RBNode _ k v l r -> RBNode Black k v l r
323− RBEmpty -> RBEmpty
282+
324− in ensureBlackRoot (del t)
283+-- make the top node red
325−
284+redden : Dict k v -> Dict k v
329− let t' = if isRedLeft t then rotateRight t else t in
330− case t' of
285+redden t = case t of
331− { RBNode _ _ _ _ RBEmpty -> RBEmpty
286+ RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
332− ; _ -> let t'' = moveRedRightIfNeeded t' in
334− { RBNode c k v l r -> fixUp (RBNode c k v l (del r))
287+ RBNode _ k v l r -> RBNode Red k v l r
335− ; RBEmpty -> RBEmpty } }
336− in ensureBlackRoot (del t)
339{-| Remove a key-value pair from a dictionary. If the key is not found,
289{-| Remove a key-value pair from a dictionary. If the key is not found,
340no changes are made. -}
290no changes are made. -}
341remove : comparable -> Dict comparable v -> Dict comparable v
291remove : comparable -> Dict comparable v -> Dict comparable v
342−remove k t =
292+remove k t = blacken <| del 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"
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"
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
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")
373{-| Apply a function to all values in a dictionary. -}
294{-| Apply a function to all values in a dictionary. -}
374map : (a -> b) -> Dict comparable a -> Dict comparable b
295map : (a -> b) -> Dict comparable a -> Dict comparable b
376 case t of
297 case t of
377− RBEmpty -> RBEmpty
298+ RBEmpty LBlack -> RBEmpty LBlack
378 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
299 RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)
380{-| Fold over the key-value pairs in a dictionary, in order from lowest
301{-| Fold over the key-value pairs in a dictionary, in order from lowest
@@ -382,7 +303,7 @@ key to highest key. -}
382foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
303foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
383foldl f acc t =
304foldl f acc t =
384 case t of
305 case t of
385− RBEmpty -> acc
306+ RBEmpty LBlack -> acc
386 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
307 RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r
388{-| Fold over the key-value pairs in a dictionary, in order from highest
309{-| Fold over the key-value pairs in a dictionary, in order from highest
@@ -390,7 +311,7 @@ key to lowest key. -}
390foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
311foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
391foldr f acc t =
312foldr f acc t =
392 case t of
313 case t of
393− RBEmpty -> acc
314+ RBEmpty LBlack -> acc
394 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
315 RBNode _ k v l r -> foldr f (f k v (foldr f acc r)) l
396{-| Combine two dictionaries. If there is a collision, preference is given
317{-| Combine two dictionaries. If there is a collision, preference is given