Improve Header.update

CHANGES.md

@@ -3,6 +3,8 @@
 - cohttp: a change in #694 modified the semantics of Header.replace.
   The semantics change is reverted, and a new Header.update function
   is introduced, following the semantics of Map.update. (#702 @mseri)
+- cohttp: reimplement update to support compilers that are older than
+  OCaml 4.06 (#703 @mseri)
 
 ## v2.5.3 (2020-06-27)
 

cohttp/src/header.ml

@@ -79,24 +79,6 @@ let replace h k v =
   let k = LString.of_string k in
   StringMap.add k [v] h
 
-let update h k f =
-  let k = LString.of_string k in
-  let f v =
-    let v' = match v with
-      | None -> f None
-      | Some l -> 
-        if is_header_with_list_value k then
-          f (Some (String.concat "," l))
-        else f (Some (List.hd l))
-    in match v' with
-    | None -> None
-    | Some s ->
-      if is_header_with_list_value k then
-        Some (String.split_on_char ',' s)
-      else Some [s]
-  in
-  StringMap.update k f h
-
 let get h k =
   let k = LString.of_string k in
   try
@@ -106,6 +88,19 @@ let get h k =
     else Some (List.hd v)
   with Not_found | Failure _ -> None
 
+  let update h k f =
+    let vorig = get h k in
+    let k = LString.of_string k in
+    match f vorig, vorig with
+    | None, _ -> StringMap.remove k h
+    | Some s, Some s' when s == s' -> h
+    | Some s, _ ->
+      let v' =
+        if is_header_with_list_value k then
+          (String.split_on_char ',' s)
+        else [s]
+      in StringMap.add k v' h
+
 let mem h k = StringMap.mem (LString.of_string k) h
 
 let add_unless_exists h k v =

cohttp/src/header.mli

@@ -63,13 +63,12 @@ val replace : t -> string -> string -> t
 (** [update h k f] returns a map containing the same headers as [h],
     except for the header [k]. Depending on the value of [v] where [v] is
     [f (get h k)], the header [k] is added, removed or updated.
-    If [w] is [None], the header is removed if it exists; otherwise,
-    if [w] is [Some z] then [k] is associated to [z] in the resulting headers.
+    If [v] is [None], the header is removed if it exists; otherwise,
+    if [v] is [Some z] then [k] is associated to [z] in the resulting headers.
     If [k] was already associated in [h] to a value that is physically equal
-    to [z], [h] is returned unchanged (the result of the function is then
-    physically equal to [h]). Similarly as for [get], if the header is one
-    of the set of headers defined to have list values, then all of the values are
-    concatenated into a single string separated by commas and passed to [f],
+    to [z], [h] is returned unchanged. Similarly as for [get], if the header is
+    one of the set of headers defined to have list values, then all of the values
+    are concatenated into a single string separated by commas and passed to [f],
     while the return value of [f] is split on commas and associated to [k].
     If it is a singleton header, then the first value is passed to [f] and
     no concatenation is performed, similarly for the return value.

cohttp/test/test_header.ml

@@ -152,6 +152,19 @@ module Updates = struct
     let h1 = H.update h "second" (function | Some _ -> Some "2a" | None -> None) in
     let h2 = H.replace h "second" "2a" in
     Alcotest.(check t_header) "update_existing_header" h1 h2
+
+  let update_headers_if_exists_rm () =
+    let h1 = H.update h "second" (function | Some _ -> None | None -> Some "3") in
+    let h2 = H.remove h "second" in
+    Alcotest.(check t_header) "update_remove_header" h1 h2
+
+  let update_headers_if_absent_add () =
+    let h = H.update h "third" (function | Some _ -> None | None -> Some "3") in
+    Alcotest.(check (option string)) "update_add_new_header" (Some "3") (H.get h "third")
+
+  let update_headers_if_absent_rm () =
+    let h1 = H.update h "third" (function _ -> None) in
+    Alcotest.(check t_header) "update_remove_absent_header" h h1
 
   let update_headers_if_exists_multi () =
     let h1 = H.update h "accept" (function | Some v -> Some ("baz,"^v) | None -> None) in
@@ -519,6 +532,9 @@ Alcotest.run "test_header" [
     "replace absent", `Quick, Updates.replace_headers_if_absent;
     "update existing", `Quick, Updates.update_headers_if_exists;
     "update existing list", `Quick, Updates.update_headers_if_exists_multi;
+    "update add absent", `Quick, Updates.update_headers_if_absent_add;
+    "update rm existing", `Quick, Updates.update_headers_if_exists_rm;
+    "update rm absent", `Quick, Updates.update_headers_if_absent_rm;
     "update absent", `Quick, Updates.update_headers_if_absent;
     "large header", `Slow, large_header;
     "many headers", `Slow, many_headers;