1 files changed, 24 insertions, 6 deletions
diff --git a/lib/cMap.ml b/lib/cMap.ml
index 4c94f122ab..d5aa8c46be 100644
--- a/lib/cMap.ml
+++ b/lib/cMap.ml
@@ -18,12 +18,14 @@ module type ExtS =
 sig
   include Map.S
   module Set : Set.S with type elt = key
+  val modify : key -> (key -> 'a -> 'a) -> 'a t -> 'a t
   val domain : 'a t -> Set.t
   val bind : (key -> 'a) -> Set.t -> 'a t
 end
 
 module MapExt (M : Map.OrderedType) :
 sig
+  val modify : M.t -> (M.t -> 'a -> 'a) -> 'a Map.Make(M).t -> 'a Map.Make(M).t
   val domain : 'a Map.Make(M).t -> Set.Make(M).t
   val bind : (M.t -> 'a) -> Set.Make(M).t -> 'a Map.Make(M).t
 end =
@@ -45,18 +47,34 @@ struct
   | SEmpty
   | SNode of set * M.t * set * int
 
-  let rec domain (s : 'a map) : set = match Obj.magic s with
-  | MEmpty -> Obj.magic SEmpty
+  let map_prj : 'a map -> 'a _map = Obj.magic
+  let map_inj : 'a _map -> 'a map = Obj.magic
+  let set_prj : set -> _set = Obj.magic
+  let set_inj : _set -> set = Obj.magic
+
+  let rec modify k f (s : 'a map) : 'a map = match map_prj s with
+  | MEmpty -> raise Not_found
+  | MNode (l, k', v, r, h) ->
+    let c = M.compare k k' in
+    if c < 0 then
+      map_inj (MNode (modify k f l, k', v, r, h))
+    else if c = 0 then
+      map_inj (MNode (l, k', f k' v, r, h))
+    else
+      map_inj (MNode (l, k', v, modify k f r, h))
+
+  let rec domain (s : 'a map) : set = match map_prj s with
+  | MEmpty -> set_inj SEmpty
   | MNode (l, k, _, r, h) ->
-    Obj.magic (SNode (domain l, k, domain r, h))
+    set_inj (SNode (domain l, k, domain r, h))
   (** This function is essentially identity, but OCaml current stdlib does not
       take advantage of the similarity of the two structures, so we introduce
       this unsafe loophole. *)
 
-  let rec bind f (s : set) : 'a map = match Obj.magic s with
-  | SEmpty -> Obj.magic MEmpty
+  let rec bind f (s : set) : 'a map = match set_prj s with
+  | SEmpty -> map_inj MEmpty
   | SNode (l, k, r, h) ->
-    Obj.magic (MNode (bind f l, k, f k, bind f r, h))
+    map_inj (MNode (bind f l, k, f k, bind f r, h))
   (** Dual operation of [domain]. *)
 
 end