Type State Pattern in OCaml

response.ml

module Response : sig
  type start
  type headers
  type _ t

  type content_type =
      | Json

  val start : start t
  val status : int -> start t -> headers t
  val header : string -> string -> headers t -> headers t
  val content_type : content_type -> headers t -> headers t
  val body : Bytes.t -> headers t -> headers t
  val to_string : headers t -> string
end = struct
  type start
  type headers

  type content_type =
      | Json

  type state = {
    code : int;
    headers : (string * string) list;
    body : Bytes.t option;
  }

  type _ t = Start : start t | Headers : state -> headers t

  let start = Start
  let status code (res : start t) = Headers { code; headers = []; body = None }
  let body payload (Headers state) = Headers { state with body = Some payload }

  let header k v (Headers state) =
    Headers { state with headers = (k, v) :: state.headers }

  let content_type ty (Headers state) =
    let v = match ty with Json -> "application/json" in
    Headers { state with headers = ("Content-Type", v) :: state.headers }

  let to_string (Headers { code; headers; body }) =
    let headers =
      headers
      |> List.map (fun (k, v) -> Printf.sprintf "%s: %s" k v)
      |> String.concat ";\n"
    in
    let str_of_body bytes =
      Bytes.(cat bytes (of_string "\n") |> to_string)
    in
    let body = Option.fold ~none:"" ~some:str_of_body body in
    Printf.sprintf {|HTTP/2 %d
%s
%s|} code headers body
end

let response =
  Response.(
    start
    |> status 200
    |> content_type Json
    |> header "X-Unexpected" "tada"
    |> status 404 (* Fails type check, because only one status line is allowed. *)
    |> body Bytes.(of_string {|{"msg":"Hello, world!"}|})
    |> to_string
    |> print_string)