@@ -1,4 +1,4 @@
1−module Dict (empty,singleton,insert
1+module Dict (empty,singleton,insert,update
22 ,lookup,findWithDefault
33 ,remove,member
44 ,foldl,foldr,map
@@ -14,7 +14,7 @@ lists of comparable types.
1414Insert, remove, and query operations all take *O(log n)* time.
1515
1616# Build
17−@docs empty, singleton, insert, remove
17+@docs empty, singleton, insert, update, remove
1818
1919# Query
2020@docs member, lookup, findWithDefault
@@ -148,17 +148,23 @@ 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 = -- Invariant: t is a valid left-leaning rb tree
152− let ins t =
151+insert k v t = let updater _ = v
152+ in update k updater t
153+
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 =
153158 case t of
154− RBEmpty LBlack -> RBNode Red k v (RBEmpty LBlack) (RBEmpty LBlack)
159+ -- Insert case
160+ RBEmpty LBlack -> RBNode Red k (u Nothing) (RBEmpty LBlack) (RBEmpty LBlack)
155161 RBNode c k' v' l r ->
156162 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)
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)
160166 in fixUp h
161− in ensureBlackRoot (ins t)
167+ in ensureBlackRoot (up t)
162168
163169{-| Create a dictionary with one key-value pair. -}
164170singleton : comparable -> v -> Dict comparable v