Merge pull request #220 from rgrinberg/simplify-response

Simplification/Deforestation

opium/src/import.ml

@@ -16,6 +16,28 @@ module List = struct
       | Some _ as result -> result
       | None -> find_map ~f l)
   ;;
+
+  let replace_or_add ~f to_add l =
+    let rec aux acc l found =
+      match l with
+      | [] -> rev (if not found then to_add :: acc else acc)
+      | el :: rest ->
+        if f el to_add then aux (to_add :: acc) rest true else aux (el :: acc) rest found
+    in
+    aux [] l false
+  ;;
 end
 
-module String = StringLabels
+module String = struct
+  include StringLabels
+
+  let rec check_prefix s ~prefix len i =
+    i = len || (s.[i] = prefix.[i] && check_prefix s ~prefix len (i + 1))
+  ;;
+
+  let is_prefix s ~prefix =
+    let len = length s in
+    let prefix_len = length prefix in
+    len >= prefix_len && check_prefix s ~prefix prefix_len 0
+  ;;
+end

opium/src/request.ml

@@ -102,25 +102,10 @@ let cookies ?signed_with t =
   Cookie.cookies_of_headers ?signed_with (t.headers |> Headers.to_list)
 ;;
 
-let replace_or_add_to_list ~f to_add l =
-  let found = ref false in
-  let rec aux acc l =
-    match l with
-    | [] -> if not !found then to_add :: acc |> List.rev else List.rev acc
-    | el :: rest ->
-      if f el to_add
-      then (
-        found := true;
-        aux (to_add :: acc) rest)
-      else aux (el :: acc) rest
-  in
-  aux [] l
-;;
-
 let add_cookie ?sign_with (k, v) t =
   let cookies = cookies t in
   let cookies =
-    replace_or_add_to_list
+    List.replace_or_add
       ~f:(fun (k2, _v2) _ -> String.equal k k2)
       ( k
       , match sign_with with
@@ -142,8 +127,8 @@ let add_cookie_unless_exists ?sign_with (k, v) t =
 let remove_cookie key t =
   let cookie_header =
     cookies t
-    |> List.filter ~f:(fun (k, _) -> not (String.equal k key))
-    |> List.map ~f:Cookie.make
+    |> List.filter_map ~f:(fun (k, v) ->
+           if not (String.equal k key) then Some (Cookie.make (k, v)) else None)
     |> Cookie.to_cookie_header
   in
   add_header_or_replace cookie_header t
@@ -158,8 +143,7 @@ let to_multipart_form_data
   =
   match t.meth, content_type t with
   | `POST, Some content_type
-    when String.length content_type > 30
-         && String.sub content_type ~pos:0 ~len:30 = "multipart/form-data; boundary=" ->
+    when String.is_prefix content_type ~prefix:"multipart/form-data; boundary=" ->
     let open Lwt.Syntax in
     let body = t.body |> Body.copy |> Body.to_stream in
     let* result = Multipart_form_data.parse ~stream:body ~content_type ~callback in
@@ -178,19 +162,15 @@ let to_multipart_form_data_exn ?callback t =
 
 let find_in_query key query =
   query
-  |> List.find_opt ~f:(fun (k, _) -> k = key)
-  |> Option.map (fun (_, r) -> r)
+  |> List.assoc_opt key
   |> fun opt ->
-  Option.bind opt (fun x ->
-      try Some (List.hd x) with
-      | Not_found -> None)
+  Option.bind opt (function
+      | [] -> None
+      | x :: _ -> Some x)
 ;;
 
 let find_list_in_query key query =
-  query
-  |> List.find_all ~f:(fun (k, _) -> k = key)
-  |> List.map ~f:(fun (_, v) -> v)
-  |> List.concat
+  query |> List.concat_map ~f:(fun (k, v) -> if k = key then v else [])
 ;;
 
 let urlencoded key t =

opium/src/response.ml

@@ -43,35 +43,17 @@ let add_headers_unless_exists hs t =
 let remove_header key t = { t with headers = Headers.remove t.headers key }
 
 let cookie ?signed_with key t =
-  let cookie_opt =
-    headers "Set-Cookie" t
-    |> List.map ~f:(fun v -> Cookie.of_set_cookie_header ?signed_with ("Set-Cookie", v))
-    |> List.find_opt ~f:(function
-           | Some Cookie.{ value = k, _; _ } when String.equal k key -> true
-           | _ -> false)
-  in
-  Option.bind cookie_opt (fun x -> x)
+  headers "Set-Cookie" t
+  |> List.find_map ~f:(fun v ->
+         match Cookie.of_set_cookie_header ?signed_with ("Set-Cookie", v) with
+         | Some (Cookie.{ value = k, _; _ } as c) when String.equal k key -> Some c
+         | _ -> None)
 ;;
 
 let cookies ?signed_with t =
   headers "Set-Cookie" t
-  |> List.map ~f:(fun v -> Cookie.of_set_cookie_header ?signed_with ("Set-Cookie", v))
-  |> List.filter_map ~f:(fun x -> x)
-;;
-
-let replace_or_add_to_list ~f to_add l =
-  let found = ref false in
-  let rec aux acc l =
-    match l with
-    | [] -> if not !found then to_add :: acc |> List.rev else List.rev acc
-    | el :: rest ->
-      if f el to_add
-      then (
-        found := true;
-        aux (to_add :: acc) rest)
-      else aux (el :: acc) rest
-  in
-  aux [] l
+  |> List.filter_map ~f:(fun v ->
+         Cookie.of_set_cookie_header ?signed_with ("Set-Cookie", v))
 ;;
 
 let add_cookie ?sign_with ?expires ?scope ?same_site ?secure ?http_only value t =
@@ -89,15 +71,12 @@ let add_cookie_or_replace ?sign_with ?expires ?scope ?same_site ?secure ?http_on
     |> Cookie.to_set_cookie_header
   in
   let headers =
-    replace_or_add_to_list
+    List.replace_or_add
       ~f:(fun (k, v) _ ->
         match k, v with
         | k, v
           when String.equal (String.lowercase_ascii k) "set-cookie"
-               && String.length v > String.length (fst value)
-               && String.equal
-                    (StringLabels.sub v ~pos:0 ~len:(String.length (fst value)))
-                    (fst value) -> true
+               && String.is_prefix v ~prefix:(fst value) -> true
         | _ -> false)
       cookie_header
       (Headers.to_list t.headers)
@@ -212,8 +191,8 @@ let to_json t =
 let to_plain_text t = Body.copy t.body |> Body.to_string
 
 let sexp_of_t { version; status; reason; headers; body; env } =
-  let open Sexplib0.Sexp_conv in
-  let open Sexplib0.Sexp in
+  let open Sexp_conv in
+  let open Sexp in
   List
     [ List [ Atom "version"; Version.sexp_of_t version ]
     ; List [ Atom "status"; Status.sexp_of_t status ]
@@ -238,5 +217,5 @@ let http_string_of_t t =
     t.body
 ;;
 
-let pp fmt t = Sexplib0.Sexp.pp_hum fmt (sexp_of_t t)
-let pp_hum fmt t = Format.fprintf fmt "%s\n%!" (http_string_of_t t)
+let pp fmt t = Sexp.pp_hum fmt (sexp_of_t t)
+let pp_hum fmt t = Format.fprintf fmt "%s@." (http_string_of_t t)