@@ -34,15 +34,35 @@ Insert, remove, and query operations all take *O(log n)* time.
3434import open Basics
3535import open Maybe
3636import Native.Error
37−import List as List
37+import List
38+import String
3839import Native.Utils
3940
4041-- BBlack and NBlack should only be used during the deletion
4142-- algorithm. Any other occurrence is a bug and should fail an assert.
42−data NColor = Red | Black
43− | BBlack | NBlack
44−
45−data LeafColor = LBlack | LBBlack
43+data 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
49+
50+
51+showNColor : NColor -> String
52+showNColor c = case c of
53+ Red -> "Red"
54+ Black -> "Black"
55+ BBlack -> "BBlack"
56+ NBlack -> "NBlack"
57+
58+data LeafColor = LBlack
59+ -- ^ Double Black, counts as 2
60+ | LBBlack
61+
62+showLColor : LeafColor -> String
63+showLColor c = case c of
64+ LBlack -> "LBlack"
65+ LBBlack -> "LBBlack"
4666
4767data Dict k v = RBNode NColor k v (Dict k v) (Dict k v)
4868 | RBEmpty LeafColor
@@ -93,51 +113,6 @@ member : comparable -> Dict comparable v -> Bool
93113-- Does t contain k?
94114member k t = isJust <| lookup k t
95115
96−rotateLeft : Dict k v -> Dict k v
97−rotateLeft t =
98− case t of
99− RBNode cy ky vy a (RBNode cz kz vz b c) -> RBNode cy kz vz (RBNode Red ky vy a b) c
100− _ -> Native.Error.raise "rotateLeft of a node without enough children"
101−
102−-- rotateRight -- the reverse, and
103−-- makes Y have Z's color, and makes Z Red.
104−rotateRight : Dict k v -> Dict k v
105−rotateRight t =
106− case t of
107− RBNode cz kz vz (RBNode cy ky vy a b) c -> RBNode cz ky vy a (RBNode Red kz vz b c)
108− _ -> Native.Error.raise "rotateRight of a node without enough children"
109−
110−rotateLeftIfNeeded : Dict k v -> Dict k v
111−rotateLeftIfNeeded t =
112− case t of
113− RBNode _ _ _ _ (RBNode Red _ _ _ _) -> rotateLeft t
114− _ -> t
115−
116−rotateRightIfNeeded : Dict k v -> Dict k v
117−rotateRightIfNeeded t =
118− case t of
119− RBNode _ _ _ (RBNode Red _ _ (RBNode Red _ _ _ _) _) _ -> rotateRight t
120− _ -> t
121−
122−otherColor c = case c of { Red -> Black ; Black -> Red }
123−
124−color_flip : Dict k v -> Dict k v
125−color_flip t =
126− case t of
127− RBNode c1 bk bv (RBNode c2 ak av la ra) (RBNode c3 ck cv lc rc) ->
128− RBNode (otherColor c1) bk bv
129− (RBNode (otherColor c2) ak av la ra)
130− (RBNode (otherColor c3) ck cv lc rc)
131− _ -> Native.Error.raise "color_flip called on a Empty or Node with a Empty child"
132−
133−color_flipIfNeeded : Dict k v -> Dict k v
134−color_flipIfNeeded t =
135− case t of
136− RBNode _ _ _ (RBNode Red _ _ _ _) (RBNode Red _ _ _ _) -> color_flip t
137− _ -> t
138−
139−fixUp t = color_flipIfNeeded (rotateRightIfNeeded (rotateLeftIfNeeded t))
140−
141116ensureBlackRoot : Dict k v -> Dict k v
142117ensureBlackRoot t =
143118 case t of
@@ -158,6 +133,11 @@ remove k t = let u _ = Nothing in
158133 update k u t
159134
160135data Flag = Insert | Remove | Same
136+showFlag : Flag -> String
137+showFlag f = case f of
138+ Insert -> "Insert"
139+ Remove -> "Remove"
140+ Same -> "Same"
161141
162142{-| Update the value of a dictionary for a specific key with a given function. -}
163143update : comparable -> (Maybe v -> Maybe v) -> Dict comparable v -> Dict comparable v
@@ -168,76 +148,87 @@ update k u t =
168148 Just v -> (Insert, RBNode Red k v empty empty)
169149 RBNode c k' v l r -> case Native.Utils.compare k k' of
170150 EQ -> case u (Just v) of
171− Nothing -> (Remove, rem t)
151+ Nothing -> (Remove, rem c l r)
172152 Just v' -> (Same, RBNode c k' v' l r)
173153 LT -> let (fl, l') = up l in
174154 case fl of
175155 Same -> (Same, RBNode c k' v l' r)
176− Insert -> (Insert, fixUp <| RBNode c k' v l' r)
156+ Insert -> (Insert, balance c k' v l' r)
177157 Remove -> (Remove, bubble c k' v l' r)
178158 GT -> let (fl, r') = up r in
179159 case fl of
180160 Same -> (Same, RBNode c k' v l r')
181− Insert -> (Insert, fixUp <| RBNode c k' v l r')
161+ Insert -> (Insert, balance c k' v l r')
182162 Remove -> (Remove, bubble c k' v l r')
183163 (fl, t') = up t
184164 in case fl of
185165 Same -> t'
186166 Insert -> ensureBlackRoot t'
187167 Remove -> blacken t'
188−
168+
189169{-| Create a dictionary with one key-value pair. -}
190170singleton : comparable -> v -> Dict comparable v
191171singleton k v = insert k v (RBEmpty LBlack)
192172
193−{- Remove helpers: everything from here to remove should only be used
194− internally by remove as they would otherwise break rb-invariants -}
195−
196173isBBlack : Dict k v -> Bool
197174isBBlack t = case t of
198− RBNode c _ _ _ _ -> case c of
199− BBlack -> True
200− _ -> False
175+ RBNode BBlack _ _ _ _ -> True
201176 RBEmpty LBBlack -> True
202177 _ -> False
203178
204179moreBlack : NColor -> NColor
205180moreBlack c = case c of
206− BBlack -> BBlack
207181 Black -> BBlack
208182 Red -> Black
209183 NBlack -> Red
184+ BBlack -> Native.Error.raise "Can't make a double black node more black!"
185+
210186lessBlack : NColor -> NColor
211187lessBlack c = case c of
212188 BBlack -> Black
213189 Black -> Red
214190 Red -> NBlack
215− NBlack -> NBlack
216−
217−moreBlackTree : Dict k v -> Dict k v
218−moreBlackTree t = case t of
219− RBNode c k v l r -> RBNode (moreBlack c) k v l r
220− RBEmpty _ -> RBEmpty LBBlack
191+ NBlack -> Native.Error.raise "Can't make a negative black node less black!"
221192
222193lessBlackTree : Dict k v -> Dict k v
223194lessBlackTree t = case t of
224195 RBNode c k v l r -> RBNode (lessBlack c) k v l r
225− RBEmpty _ -> RBEmpty LBlack
196+ RBEmpty LBBlack -> RBEmpty LBlack
197+
198+reportRemBug : String -> NColor -> String -> String -> a
199+reportRemBug 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+ ]
226211
227212-- Remove the top node from the tree, may leave behind BBlacks
228−rem : Dict k v -> Dict k v
229−rem t = case t of
230− RBNode c k v (RBEmpty _) (RBEmpty _) -> case c of
231− Red -> RBEmpty LBlack
213+rem : NColor -> Dict k v -> Dict k v -> Dict k v
214+rem c l r = case (l, r) of
215+ ((RBEmpty _), (RBEmpty _)) -> case c of
216+ Red -> RBEmpty LBlack
232217 Black -> RBEmpty LBBlack
233− RBNode Black _ _ (RBEmpty _) (RBNode _ k v l r) ->
234− RBNode Black k v l r
235− RBNode Black _ _ (RBNode _ k v l r) (RBEmpty _) ->
236− RBNode Black k v l r
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)
237226 -- l and r are both RBNodes
238− RBNode c _ _ l r ->
239− let (k, v) = max l
240− l' = remove_max l
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
241232 in bubble c k v l' r
242233
243234-- Kills a BBlack or moves it upward, may leave behind NBlack
@@ -245,19 +236,24 @@ bubble : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
245236bubble c k v l r = if isBBlack l || isBBlack r
246237 then balance (moreBlack c) k v (lessBlackTree l) (lessBlackTree r)
247238 else RBNode c k v l r
248−
239+
249240-- Removes rightmost node, may leave root as BBlack
250−remove_max : Dict k v -> Dict k v
251−remove_max t = case t of
252− RBNode c k v l (RBEmpty _) -> rem t
253− RBNode c k v l r -> bubble c k v l (remove_max r)
241+remove_max : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
242+remove_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)
254246
255247-- generalized tree balancing act
256248balance : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
257−balance c k v l r = balance_node (RBNode c k v l r)
249+balance c k v l r =
250+
251+ balance_node (RBNode c k v l r)
258252
259253blackish : Dict k v -> Bool
260−blackish (RBNode c _ _ _ _) = c == Black || c == BBlack
254+blackish t = case t of
255+ RBNode c _ _ _ _ -> c == Black || c == BBlack
256+ RBEmpty _ -> True
261257
262258balance_node : Dict k v -> Dict k v
263259balance_node t =
@@ -286,9 +282,7 @@ balance_node t =
286282 (RBNode Black _ _ _ _) ->
287283 RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
288284 _ -> t
289−
290285 _ -> t
291−
292286 else t
293287
294288-- make the top node black