@@ -148,24 +148,44 @@ ensureBlackRoot t =
148148{-| Insert a key-value pair into a dictionary. Replaces value when there is
149149a collision. -}
150150insert : comparable -> v -> Dict comparable v -> Dict comparable v
151−insert k v t = let updater _ = v
152− in update k updater t
151+insert k v t = let u _ = Just v in
152+ update k u t
153153
154−{-| Update the value of a dictionary for a specific key with a given function. -}
155−update : comparable -> (Maybe v -> v) -> Dict comparable v -> Dict comparable v
156−update k u t = -- Invariant: t is a valid left-leaning rb tree
157− let up t =
158− case t of
159− -- Insert case
160− RBEmpty LBlack -> RBNode Red k (u Nothing) (RBEmpty LBlack) (RBEmpty LBlack)
161− RBNode c k' v' l r ->
162− let h = case Native.Utils.compare k k' of
163− LT -> RBNode c k' v' (up l) r
164− EQ -> RBNode c k' (u (Just v')) l r -- update case
165− GT -> RBNode c k' v' l (up r)
166− in fixUp h
167− in ensureBlackRoot (up t)
154+{-| Remove a key-value pair from a dictionary. If the key is not found,
155+no changes are made. -}
156+remove : comparable -> Dict comparable v -> Dict comparable v
157+remove k t = let u _ = Nothing in
158+ update k u t
159+
160+data Flag = Insert | Remove | Same
168161
162+{-| Update the value of a dictionary for a specific key with a given function. -}
163+update : comparable -> (Maybe v -> Maybe v) -> Dict comparable v -> Dict comparable v
164+update k u t =
165+ let up t = case t of
166+ RBEmpty LBlack -> case u Nothing of
167+ Nothing -> (Same, empty)
168+ Just v -> (Insert, RBNode Red k v empty empty)
169+ RBNode c k' v l r -> case Native.Utils.compare k k' of
170+ EQ -> case u (Just v) of
171+ Nothing -> (Remove, rem t)
172+ Just v' -> (Same, RBNode c k' v' l r)
173+ LT -> let (fl, l') = up l in
174+ case fl of
175+ Same -> (Same, RBNode c k' v l' r)
176+ Insert -> (Insert, fixUp <| RBNode c k' v l' r)
177+ Remove -> (Remove, bubble c k' v l' r)
178+ GT -> let (fl, r') = up r in
179+ case fl of
180+ Same -> (Same, RBNode c k' v l r')
181+ Insert -> (Insert, fixUp <| RBNode c k' v l r')
182+ Remove -> (Remove, bubble c k' v l r')
183+ (fl, t') = up t
184+ in case fl of
185+ Same -> t'
186+ Insert -> ensureBlackRoot t'
187+ Remove -> blacken t'
188+
169189{-| Create a dictionary with one key-value pair. -}
170190singleton : comparable -> v -> Dict comparable v
171191singleton k v = insert k v (RBEmpty LBlack)
@@ -204,15 +224,6 @@ lessBlackTree t = case t of
204224 RBNode c k v l r -> RBNode (lessBlack c) k v l r
205225 RBEmpty _ -> RBEmpty LBlack
206226
207−-- Finds and deletes k in t
208−del : comparable -> Dict comparable v -> Dict comparable v
209−del k t = case t of
210− RBEmpty _ -> t
211− RBNode c k' v l r -> case Native.Utils.compare k k' of
212− LT -> bubble c k' v (del k l) r
213− EQ -> rem t
214− GT -> bubble c k' v l (del k r)
215−
216227-- Remove the top node from the tree, may leave behind BBlacks
217228rem : Dict k v -> Dict k v
218229rem t = case t of
@@ -292,11 +303,6 @@ redden t = case t of
292303 RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
293304 RBNode _ k v l r -> RBNode Red k v l r
294305
295−{-| Remove a key-value pair from a dictionary. If the key is not found,
296−no changes are made. -}
297−remove : comparable -> Dict comparable v -> Dict comparable v
298−remove k t = blacken <| del k t
299−
300306{-| Apply a function to all values in a dictionary. -}
301307map : (a -> b) -> Dict comparable a -> Dict comparable b
302308map f t =