@@ -1,4 +1,4 @@
1−module Dict (empty,singleton,insert
1+module Dict (empty,singleton,insert,update
2 ,lookup,findWithDefault
2 ,lookup,findWithDefault
3 ,remove,member
3 ,remove,member
4 ,foldl,foldr,map
4 ,foldl,foldr,map
@@ -14,7 +14,7 @@ lists of comparable types.
14Insert, remove, and query operations all take *O(log n)* time.
14Insert, remove, and query operations all take *O(log n)* time.
17−@docs empty, singleton, insert, remove
17+@docs empty, singleton, insert, update, remove
20@docs member, lookup, findWithDefault
20@docs member, lookup, findWithDefault
@@ -148,18 +148,44 @@ ensureBlackRoot t =
148{-| 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
149a collision. -}
149a collision. -}
150insert : comparable -> v -> Dict comparable v -> Dict comparable v
150insert : comparable -> v -> Dict comparable v -> Dict comparable v
151−insert k v t = -- Invariant: t is a valid left-leaning rb tree
151+insert k v t = let u _ = Just v in
152− let ins t =
152+ update k u t
154− RBEmpty LBlack -> RBNode Red k v (RBEmpty LBlack) (RBEmpty LBlack)
155− RBNode c k' v' l r ->
156− let h = case Native.Utils.compare k k' of
157− LT -> RBNode c k' v' (ins l) r
158− EQ -> RBNode c k' v l r -- replace
159− GT -> RBNode c k' v' l (ins r)
161− in ensureBlackRoot (ins 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
160+data Flag = Insert | Remove | Same
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
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
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
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')
186+ Insert -> ensureBlackRoot t'
187+ Remove -> blacken t'
163{-| Create a dictionary with one key-value pair. -}
189{-| Create a dictionary with one key-value pair. -}
164singleton : comparable -> v -> Dict comparable v
190singleton : comparable -> v -> Dict comparable v
165singleton k v = insert k v (RBEmpty LBlack)
191singleton k v = insert k v (RBEmpty LBlack)
@@ -198,15 +224,6 @@ lessBlackTree t = case t of
198 RBNode c k v l r -> RBNode (lessBlack c) k v l r
224 RBNode c k v l r -> RBNode (lessBlack c) k v l r
199 RBEmpty _ -> RBEmpty LBlack
225 RBEmpty _ -> RBEmpty LBlack
201−-- Finds and deletes k in t
202−del : comparable -> Dict comparable v -> Dict comparable v
205− RBNode c k' v l r -> case Native.Utils.compare k k' of
206− LT -> bubble c k' v (del k l) r
208− GT -> bubble c k' v l (del k r)
210-- Remove the top node from the tree, may leave behind BBlacks
227-- Remove the top node from the tree, may leave behind BBlacks
211rem : Dict k v -> Dict k v
228rem : Dict k v -> Dict k v
212rem t = case t of
229rem t = case t of
@@ -286,11 +303,6 @@ redden t = case t of
286 RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
303 RBEmpty _ -> Native.Error.raise "can't make a Leaf red"
287 RBNode _ k v l r -> RBNode Red k v l r
304 RBNode _ k v l r -> RBNode Red k v l r
289−{-| Remove a key-value pair from a dictionary. If the key is not found,
290−no changes are made. -}
291−remove : comparable -> Dict comparable v -> Dict comparable v
292−remove k t = blacken <| del k t
294{-| Apply a function to all values in a dictionary. -}
306{-| Apply a function to all values in a dictionary. -}
295map : (a -> b) -> Dict comparable a -> Dict comparable b
307map : (a -> b) -> Dict comparable a -> Dict comparable b