[lib] Add `Map.update` from OCaml 4.06

It will take more than a year to bump the OCaml version, this is in
response of a request by @Skyskimmer.

We also update our internal repr to make it closer to the one in
modern OCaml.

3 files changed, 130 insertions, 30 deletions

diff --git a/clib/cMap.ml b/clib/cMap.ml
index e4ce6c7c02..016d8bdeca 100644
--- a/clib/cMap.ml
+++ b/clib/cMap.ml
@@ -36,6 +36,7 @@ sig
   val fold_right : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
   val height : 'a t -> int
   val filter_range : (key -> int) -> 'a t -> 'a t
+  val update: key -> ('a option -> 'a option) -> 'a t -> 'a t
   module Smart :
   sig
     val map : ('a -> 'a) -> 'a t -> 'a t
@@ -64,6 +65,7 @@ sig
   val fold_right : (M.t -> 'a -> 'b -> 'b) -> 'a map -> 'b -> 'b
   val height : 'a map -> int
   val filter_range : (M.t -> int) -> 'a map -> 'a map
+  val update: M.t -> ('a option -> 'a option) -> 'a map -> 'a map
   module Smart :
   sig
     val map : ('a -> 'a) -> 'a map -> 'a map
@@ -94,8 +96,8 @@ struct
   type set = S.t
 
   type 'a _map =
-  | MEmpty
-  | MNode of 'a map * M.t * 'a * 'a map * int
+    | MEmpty
+    | MNode of {l:'a map; v:F.key; d:'a; r:'a map; h:int}
 
   type _set =
   | SEmpty
@@ -108,41 +110,41 @@ struct
 
   let rec set k v (s : 'a map) : 'a map = match map_prj s with
   | MEmpty -> raise Not_found
-  | MNode (l, k', v', r, h) ->
+  | MNode {l; v=k'; d=v'; r; h} ->
     let c = M.compare k k' in
     if c < 0 then
       let l' = set k v l in
       if l == l' then s
-      else map_inj (MNode (l', k', v', r, h))
+      else map_inj (MNode {l=l'; v=k'; d=v'; r; h})
     else if c = 0 then
       if v' == v then s
-      else map_inj (MNode (l, k', v, r, h))
+      else map_inj (MNode {l; v=k'; d=v; r; h})
     else
       let r' = set k v r in
       if r == r' then s
-      else map_inj (MNode (l, k', v', r', h))
+      else map_inj (MNode {l; v=k'; d=v'; r=r'; h})
 
   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
+  | MNode {l; v; d; r; h} ->
+    let c = M.compare k v in
     if c < 0 then
       let l' = modify k f l in
       if l == l' then s
-      else map_inj (MNode (l', k', v, r, h))
+      else map_inj (MNode {l=l'; v; d; r; h})
     else if c = 0 then
-      let v' = f k' v in
-      if v' == v then s
-      else map_inj (MNode (l, k', v', r, h))
+      let d' = f v d in
+      if d' == d then s
+      else map_inj (MNode {l; v; d=d'; r; h})
     else
       let r' = modify k f r in
       if r == r' then s
-      else map_inj (MNode (l, k', v, r', h))
+      else map_inj (MNode {l; v; d; r=r'; h})
 
   let rec domain (s : 'a map) : set = match map_prj s with
   | MEmpty -> set_inj SEmpty
-  | MNode (l, k, _, r, h) ->
-    set_inj (SNode (domain l, k, domain r, h))
+  | MNode {l; v; r; h; _} ->
+    set_inj (SNode (domain l, v, 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. *)
@@ -150,31 +152,31 @@ struct
   let rec bind f (s : set) : 'a map = match set_prj s with
   | SEmpty -> map_inj MEmpty
   | SNode (l, k, r, h) ->
-    map_inj (MNode (bind f l, k, f k, bind f r, h))
+    map_inj (MNode { l=bind f l; v=k; d=f k; r=bind f r; h})
   (** Dual operation of [domain]. *)
 
   let rec fold_left f (s : 'a map) accu = match map_prj s with
   | MEmpty -> accu
-  | MNode (l, k, v, r, h) ->
+  | MNode {l; v=k; d=v; r; h} ->
     let accu = f k v (fold_left f l accu) in
     fold_left f r accu
 
   let rec fold_right f (s : 'a map) accu = match map_prj s with
   | MEmpty -> accu
-  | MNode (l, k, v, r, h) ->
+  | MNode {l; v=k; d=v; r; h} ->
     let accu = f k v (fold_right f r accu) in
     fold_right f l accu
 
   let height s = match map_prj s with
   | MEmpty -> 0
-  | MNode (_, _, _, _, h) -> h
+  | MNode {h;_} -> h
 
   (* Filter based on a range *)
   let filter_range in_range m =
     let rec aux m = function
       | MEmpty -> m
-      | MNode (l, k, v, r, _) ->
-        let vr = in_range k in
+      | MNode {l; v; d; r; _} ->
+        let vr = in_range v in
         (* the range is below the current value *)
         if vr < 0 then aux m (map_prj l)
         (* the range is above the current value *)
@@ -183,29 +185,102 @@ struct
         else
           let m = aux m (map_prj l) in
           let m = aux m (map_prj r) in
-          F.add k v m
+          F.add v d m
     in aux F.empty (map_prj m)
 
+  (* Imported from OCaml upstream until we can bump the version *)
+  let create l x d r =
+    let hl = height l and hr = height r in
+    map_inj @@ MNode{l; v=x; d; r; h=(if hl >= hr then hl + 1 else hr + 1)}
+
+  let bal l x d r =
+    let hl = match map_prj l with MEmpty -> 0 | MNode {h} -> h in
+    let hr = match map_prj r with MEmpty -> 0 | MNode {h} -> h in
+    if hl > hr + 2 then begin
+      match map_prj l with
+      | MEmpty -> invalid_arg "Map.bal"
+      | MNode{l=ll; v=lv; d=ld; r=lr} ->
+        if height ll >= height lr then
+          create ll lv ld (create lr x d r)
+        else begin
+          match map_prj lr with
+          | MEmpty -> invalid_arg "Map.bal"
+          | MNode{l=lrl; v=lrv; d=lrd; r=lrr}->
+            create (create ll lv ld lrl) lrv lrd (create lrr x d r)
+        end
+    end else if hr > hl + 2 then begin
+      match map_prj r with
+      | MEmpty -> invalid_arg "Map.bal"
+      | MNode{l=rl; v=rv; d=rd; r=rr} ->
+        if height rr >= height rl then
+          create (create l x d rl) rv rd rr
+        else begin
+          match map_prj rl with
+          | MEmpty -> invalid_arg "Map.bal"
+          | MNode{l=rll; v=rlv; d=rld; r=rlr} ->
+            create (create l x d rll) rlv rld (create rlr rv rd rr)
+        end
+    end else
+      map_inj @@ MNode{l; v=x; d; r; h=(if hl >= hr then hl + 1 else hr + 1)}
+
+  let rec remove_min_binding m = match map_prj m with
+    | MEmpty -> invalid_arg "Map.remove_min_elt"
+    | MNode {l;v;d;r;_} ->
+      match map_prj l with
+      | MEmpty -> r
+      | _ -> bal (remove_min_binding l) v d r
+
+  let merge t1 t2 =
+    match (map_prj t1, map_prj t2) with
+      (MEmpty, t) -> map_inj t
+    | (t, MEmpty) -> map_inj t
+    | (_, _) ->
+      let (x, d) = F.min_binding t2 in
+      bal t1 x d (remove_min_binding t2)
+
+  let rec update x f m = match map_prj m with
+    | MEmpty ->
+      begin match f None with
+        | None -> map_inj MEmpty
+        | Some data -> map_inj @@ MNode{l=map_inj MEmpty; v=x; d=data; r=map_inj MEmpty; h=1}
+      end
+    | MNode {l; v; d; r; h} as m ->
+      let c = M.compare x v in
+      if c = 0 then begin
+        match f (Some d) with
+        | None -> merge l r
+        | Some data ->
+          if d == data then map_inj m else
+            map_inj @@ MNode{l; v=x; d=data; r; h}
+      end else if c < 0 then
+        let ll = update x f l in
+        if l == ll then map_inj m else bal ll v d r
+      else
+        let rr = update x f r in
+        if r == rr then map_inj m else bal l v d rr
+
+  (* End of Imported OCaml *)
+
   module Smart =
   struct
 
     let rec map f (s : 'a map) = match map_prj s with
     | MEmpty -> map_inj MEmpty
-    | MNode (l, k, v, r, h) ->
+    | MNode {l; v=k; d=v; r; h} ->
       let l' = map f l in
       let r' = map f r in
       let v' = f v in
       if l == l' && r == r' && v == v' then s
-      else map_inj (MNode (l', k, v', r', h))
+      else map_inj (MNode {l=l'; v=k; d=v'; r=r'; h})
 
     let rec mapi f (s : 'a map) = match map_prj s with
     | MEmpty -> map_inj MEmpty
-    | MNode (l, k, v, r, h) ->
+    | MNode {l; v=k; d=v; r; h} ->
       let l' = mapi f l in
       let r' = mapi f r in
       let v' = f k v in
       if l == l' && r == r' && v == v' then s
-      else map_inj (MNode (l', k, v', r', h))
+      else map_inj (MNode {l=l'; v=k; d=v'; r=r'; h})
 
   end
 
@@ -214,9 +289,9 @@ struct
 
     let rec map f (s : 'a map) : 'b map = match map_prj s with
     | MEmpty -> map_inj MEmpty
-    | MNode (l, k, v, r, h) ->
+    | MNode {l; v=k; d=v; r; h} ->
       let (k, v) = f k v in
-      map_inj (MNode (map f l, k, v, map f r, h))
+      map_inj (MNode {l=map f l; v=k; d=v; r=map f r; h})
 
   end
 
@@ -227,14 +302,14 @@ struct
 
     let rec fold_left f s accu = match map_prj s with
     | MEmpty -> return accu
-    | MNode (l, k, v, r, h) ->
+    | MNode {l; v=k; d=v; r; h} ->
       fold_left f l accu >>= fun accu ->
       f k v accu >>= fun accu ->
       fold_left f r accu
 
     let rec fold_right f s accu = match map_prj s with
     | MEmpty -> return accu
-    | MNode (l, k, v, r, h) ->
+    | MNode {l; v=k; d=v; r; h} ->
       fold_right f r accu >>= fun accu ->
       f k v accu >>= fun accu ->
       fold_right f l accu
diff --git a/clib/cMap.mli b/clib/cMap.mli
index ca6ddb2f4e..9bbb8d50dd 100644
--- a/clib/cMap.mli
+++ b/clib/cMap.mli
@@ -66,6 +66,18 @@ sig
       [filter_range] returns the submap of [m] whose keys are in
       range. Note that [in_range] has to define a continouous range. *)
 
+  val update: key -> ('a option -> 'a option) -> 'a t -> 'a t
+  (** [update x f m] returns a map containing the same bindings as
+      [m], except for the binding of [x]. Depending on the value of
+      [y] where [y] is [f (find_opt x m)], the binding of [x] is
+      added, removed or updated. If [y] is [None], the binding is
+      removed if it exists; otherwise, if [y] is [Some z] then [x]
+      is associated to [z] in the resulting map.  If [x] was already
+      bound in [m] to a value that is physically equal to [z], [m]
+      is returned unchanged (the result of the function is then
+      physically equal to [m]).
+    *)
+
   module Smart :
   sig
     val map : ('a -> 'a) -> 'a t -> 'a t
diff --git a/clib/hMap.ml b/clib/hMap.ml
index 5d634b7af0..92c6cbd434 100644
--- a/clib/hMap.ml
+++ b/clib/hMap.ml
@@ -408,6 +408,19 @@ struct
   let filter_range f s =
     filter (fun x _ -> f x = 0) s
 
+  (* Not as efficient as the original version *)
+  let update k f m =
+    try
+      match f (Some (find k m)) with
+      | None -> remove k m
+      | Some v ->
+        let m = remove k m in
+        add k v m
+    with Not_found ->
+      match f None with
+      | None -> m
+      | Some v -> add k v m
+
   module Unsafe =
   struct
     let map f s =