From c10b240474b8f81898b2a0d975c6da217d0146ab Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 16 Oct 2015 15:22:23 +0200
Subject: [PATCH] add `CCHashtbl.update`

---
 src/core/CCHashtbl.ml  | 41 ++++++++++++++++++++++++++++++++++++++++-
 src/core/CCHashtbl.mli | 16 ++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/core/CCHashtbl.ml b/src/core/CCHashtbl.ml
index 0321c464..65412453 100644
--- a/src/core/CCHashtbl.ml
+++ b/src/core/CCHashtbl.ml
@@ -71,6 +71,25 @@ let of_list l =
   List.iter (fun (k,v) -> Hashtbl.add tbl k v) l;
   tbl
 
+let update tbl ~f ~k =
+  let v = get tbl k in
+  match v, f k v with
+  | None, None -> ()
+  | None, Some v' -> Hashtbl.add tbl k v'
+  | Some _, Some v' -> Hashtbl.replace tbl k v'
+  | Some _, None -> Hashtbl.remove tbl k
+
+(*$R
+  let tbl = Hashtbl.create 32 in
+  update tbl ~k:1 ~f:(fun _ _ -> Some "1");
+  assert_equal (Some "1") (get tbl 1);
+  update tbl ~k:2 ~f:(fun _ v->match v with Some _ -> assert false | None -> Some "2");
+  assert_equal (Some "2") (get tbl 2);
+  assert_equal 2 (Hashtbl.length tbl);
+  update tbl ~k:1 ~f:(fun _ _ -> None);
+  assert_equal None (get tbl 1);
+*)
+
 let print pp_k pp_v fmt m =
   Format.fprintf fmt "@[<hov2>tbl {@,";
   let first = ref true in
@@ -121,10 +140,22 @@ module type S = sig
   val of_list : (key * 'a) list -> 'a t
   (** From the given list of bindings, added in order *)
 
+  val update : 'a t -> f:(key -> 'a option -> 'a option) -> k:key -> unit
+  (** [update tbl ~f ~k] updates key [k] by calling [f k (Some v)] if
+      [k] was mapped to [v], or [f k None] otherwise; if the call
+      returns [None] then [k] is removed/stays removed, if the call
+      returns [Some v'] then the binding [k -> v'] is inserted
+      using {!Hashtbl.replace}
+      @since NEXT_RELEASE *)
+
   val print : key printer -> 'a printer -> 'a t printer
+  (** Printer for tables
+      @since 0.13 *)
 end
 
-module Make(X : Hashtbl.HashedType) = struct
+module Make(X : Hashtbl.HashedType)
+  : S with type key = X.t and type 'a t = 'a Hashtbl.Make(X).t
+= struct
   include Hashtbl.Make(X)
 
   let get tbl x =
@@ -143,6 +174,14 @@ module Make(X : Hashtbl.HashedType) = struct
       (fun x y acc -> f x y :: acc)
       h []
 
+  let update tbl ~f ~k =
+    let v = get tbl k in
+    match v, f k v with
+    | None, None -> ()
+    | None, Some v' -> add tbl k v'
+    | Some _, Some v' -> replace tbl k v'
+    | Some _, None -> remove tbl k
+
   let to_seq tbl k = iter (fun key v -> k (key,v)) tbl
 
   let of_seq seq =
diff --git a/src/core/CCHashtbl.mli b/src/core/CCHashtbl.mli
index fa5f521c..aaa3cf2c 100644
--- a/src/core/CCHashtbl.mli
+++ b/src/core/CCHashtbl.mli
@@ -68,6 +68,14 @@ val to_list : ('a,'b) Hashtbl.t -> ('a * 'b) list
 val of_list : ('a * 'b) list -> ('a,'b) Hashtbl.t
 (** From the given list of bindings, added in order *)
 
+val update : ('a, 'b) Hashtbl.t -> f:('a -> 'b option -> 'b option) -> k:'a -> unit
+(** [update tbl ~f ~k] updates key [k] by calling [f k (Some v)] if
+    [k] was mapped to [v], or [f k None] otherwise; if the call
+    returns [None] then [k] is removed/stays removed, if the call
+    returns [Some v'] then the binding [k -> v'] is inserted
+    using {!Hashtbl.replace}
+    @since NEXT_RELEASE *)
+
 val print : 'a printer -> 'b printer -> ('a, 'b) Hashtbl.t printer
 (** Printer for table
     @since 0.13 *)
@@ -109,6 +117,14 @@ module type S = sig
   val of_list : (key * 'a) list -> 'a t
   (** From the given list of bindings, added in order *)
 
+  val update : 'a t -> f:(key -> 'a option -> 'a option) -> k:key -> unit
+  (** [update tbl ~f ~k] updates key [k] by calling [f k (Some v)] if
+      [k] was mapped to [v], or [f k None] otherwise; if the call
+      returns [None] then [k] is removed/stays removed, if the call
+      returns [Some v'] then the binding [k -> v'] is inserted
+      using {!Hashtbl.replace}
+      @since NEXT_RELEASE *)
+
   val print : key printer -> 'a printer -> 'a t printer
   (** Printer for tables
       @since 0.13 *)