modify CCPersistentHashtbl.merge and add CCMap.merge_safe

src/core/CCMap.ml

@@ -24,6 +24,12 @@ module type S = sig
       [k] is removed from [m], and if the result is [Some v'] then
       [add k v' m] is returned. *)
 
+  val merge_safe :
+    f:(key -> [`Left of 'a | `Right of 'b | `Both of 'a * 'b] -> 'c option) ->
+    'a t -> 'b t -> 'c t
+  (** [merge_safe ~f a b] merges the maps [a] and [b] together.
+      @since NEXT_RELEASE *)
+
   val of_seq : (key * 'a) sequence -> 'a t
 
   val add_seq : 'a t -> (key * 'a) sequence -> 'a t
@@ -75,6 +81,15 @@ module Make(O : Map.OrderedType) = struct
     | None -> remove k m
     | Some v' -> add k v' m
 
+  let merge_safe ~f a b =
+    merge
+      (fun k v1 v2 -> match v1, v2 with
+         | None, None -> assert false
+         | Some v1, None -> f k (`Left v1)
+         | None, Some v2 -> f k (`Right v2)
+         | Some v1, Some v2 -> f k (`Both (v1,v2)))
+      a b
+
   let add_seq m s =
     let m = ref m in
     s (fun (k,v) -> m := add k v !m);

src/core/CCMap.mli

@@ -27,6 +27,12 @@ module type S = sig
       [k] is removed from [m], and if the result is [Some v'] then
       [add k v' m] is returned. *)
 
+  val merge_safe :
+    f:(key -> [`Left of 'a | `Right of 'b | `Both of 'a * 'b] -> 'c option) ->
+    'a t -> 'b t -> 'c t
+  (** [merge_safe ~f a b] merges the maps [a] and [b] together.
+      @since NEXT_RELEASE *)
+
   val of_seq : (key * 'a) sequence -> 'a t
 
   val add_seq : 'a t -> (key * 'a) sequence -> 'a t

src/data/CCPersistentHashtbl.ml

@@ -67,8 +67,9 @@ module type S = sig
   (** Fresh copy of the table; the underlying structure is not shared
       anymore, so using both tables alternatively will be efficient *)
 
-  val merge : (key -> 'a option -> 'b option -> 'c option) ->
-              'a t -> 'b t -> 'c t
+  val merge :
+    (key -> [`Left of 'a | `Right of 'b | `Both of 'a * 'b] -> 'c option) ->
+    'a t -> 'b t -> 'c t
   (** Merge two tables together into a new table. The function's argument
       correspond to values associated with the key (if present); if the
       function returns [None] the key will not appear in the result. *)
@@ -543,8 +544,11 @@ module Make(H : HashedType) : S with type key = H.t = struct
     let tbl = create (max (length t1) (length t2)) in
     let tbl = fold
       (fun tbl k v1 ->
-        let v2 = try Some (find t2 k) with Not_found -> None in
-        match f k (Some v1) v2 with
+        let comb =
+          try `Both (v1, find t2 k)
+          with Not_found -> `Left v1
+        in
+        match f k comb with
         | None -> tbl
         | Some v' -> replace tbl k v')
       tbl t1
@@ -552,7 +556,7 @@ module Make(H : HashedType) : S with type key = H.t = struct
     fold
       (fun tbl k v2 ->
         if mem t1 k then tbl
-        else match f k None (Some v2) with
+        else match f k (`Right v2) with
           | None -> tbl
           | Some v' -> replace tbl k v'
       ) tbl t2

src/data/CCPersistentHashtbl.mli

@@ -74,8 +74,9 @@ module type S = sig
   (** Fresh copy of the table; the underlying structure is not shared
       anymore, so using both tables alternatively will be efficient *)
 
-  val merge : (key -> 'a option -> 'b option -> 'c option) ->
-              'a t -> 'b t -> 'c t
+  val merge :
+    f:(key -> [`Left of 'a | `Right of 'b | `Both of 'a * 'b] -> 'c option) ->
+    'a t -> 'b t -> 'c t
   (** Merge two tables together into a new table. The function's argument
       correspond to values associated with the key (if present); if the
       function returns [None] the key will not appear in the result. *)