Source file vif_route.ml

(* Part of this code is based on the furl project and
   Copyright (c) 2015 Gabriel Radanne <drupyo@zoho.com>
   SPDX-License-Identifier: MIT
   Copyright (c) 2025 Romain Calascibetta <romain.calascibetta@gmail.com>
*)

let src = Logs.Src.create "vif.r"

module Log = (val Logs.src_log src : Logs.LOG)

module Ext = struct
  let slash = Re.char '/'
  let comma = Re.char ','
  let amper = Re.char '&'
  let question_mark = Re.char '?'

  (** -?[0-9]+( .[0-9]* )? *)
  let float =
    let open Re in
    seq [ opt (char '-'); rep1 digit; opt (seq [ char '.'; rep digit ]) ]

  (** -?[0-9]+ *)
  let arbitrary_int =
    let open Re in
    seq [ opt (char '-'); rep1 digit ]

  (** true|false *)
  let bool =
    let open Re in
    alt [ str "true"; str "false" ]

  (** Non empty list of safe chars *)
  let string component =
    let open Re in
    match component with
    | `Path -> rep1 @@ compl [ slash; question_mark ]
    | `Query_value -> rep1 @@ compl [ set "&;+," ]

  (** Separated by , or by / *)
  let list ?m ~component n re =
    let open Re in
    match component with
    | `Path -> repn (seq [ slash; re ]) n m
    | `Query_value ->
        if n = 0 then alt [ epsilon; seq [ re; repn (seq [ comma; re ]) 0 m ] ]
        else seq [ re; repn (seq [ comma; re ]) (n - 1) m ]

  let query_sep ~any =
    if not any then amper
    else
      let open Re in
      seq [ amper; rep @@ seq [ rep1 @@ compl [ amper ]; amper ] ]
end

module Utils = struct
  let map_snd f (x, y) = (x, f y)

  let rec intersperse sep = function
    | [] -> []
    | [ x ] -> [ x ]
    | h :: t -> h :: sep :: intersperse sep t

  (** Offset of [el] in [l], given the function count. Used to get the first
      regexp group at a given place. *)
  let find_idx count el l =
    let rec aux el i = function
      | [] -> raise Not_found
      | x :: l' -> if x == el then i else aux el (i + count el) l'
    in
    aux el 0 l

  (* Invariants:
     - [l_before] is included in [l_after].
     - No duplicates (see note on {!find_idx}).
  *)

  (** if [l' ∈ l] then [build_permutation offset count l l'] builds a mapping:
      index in [l => offset in l']. Offsets are computed respecting [offset] and
      [count]. *)
  let build_permutation offset count l_before l_after =
    let t = Array.make (List.length l_before) 0 in
    l_before
    |> List.iteri (fun i x ->
        let j = find_idx count x l_after in
        t.(i) <- offset + j);
    t
end

open Tyre.Internal

let sort_query l = List.sort (fun (x, _) (y, _) -> compare (x : string) y) l

type 'a re_atom = 'a Tyre.Internal.wit

(** Top level atoms are specialized for path and query, see documentation. *)
let re_atom re = Tyre.Internal.build re

let re_atom_path : type e a. int -> (e, a) raw -> int * a re_atom * Re.t list =
  let open Re in
  fun i -> function
    | Rep e ->
        let _, w, re = re_atom 1 e in
        ( i + 1
        , Rep (i, w, Re.compile re)
        , [ group @@ Ext.list ~component:`Path 0 @@ no_group re ] )
    | Opt e ->
        let i', w, re = re_atom i e in
        let id, re = mark re in
        (i', Opt (id, w), [ alt [ epsilon; seq [ Ext.slash; re ] ] ])
    | e ->
        let i', w, re = re_atom i e in
        (i', w, [ Ext.slash; re ])

let re_atom_query : type e a. int -> (e, a) raw -> int * a re_atom * Re.t =
  let open Re in
  fun i -> function
    | Rep e ->
        let _, w, re = re_atom 1 e in
        ( i + 1
        , Rep (i, w, Re.compile re)
        , group @@ Ext.list ~component:`Query_value 0 @@ no_group re )
    | e ->
        let i', w, re = re_atom i e in
        (i', w, re)

type ('f, 'r) re_path =
  | Start : ('r, 'r) re_path
  | PathAtom : ('f, 'a -> 'r) re_path * 'a re_atom -> ('f, 'r) re_path

let rec re_path : type e r f.
    int -> (e, f, r) Vif_uri.path -> bool * int * (f, r) re_path * Re.t list =
  let open Re in
  fun i -> function
    | Host str ->
        let re = Re.str (Uri.pct_encode ~component:`Host str) in
        (true, i, Start, [ re ])
    | Rel ->
        (* NOTE(dinosaure): a relative route must match {i regardless} of the
           host. As the host (when available) is prepended to the string we try
           to match (see {!val:prepare_target}), we let the route optionally
           consume such a host prefix (anything up to the first ['/']). This is
           what allows [rel] and [host] routes to coexist in the same dispatch
           table. *)
        let re = opt (rep1 (compl [ Ext.slash ])) in
        (false, i, Start, [ re ])
    | Path_const (p, s) ->
        let with_host, i', p, re = re_path i p in
        (with_host, i', p, str s :: Ext.slash :: re)
    | Path_atom (p, a) ->
        let with_host, i', wp, rp = re_path i p in
        let i'', wa, ra = re_atom_path i' @@ from_t a in
        (with_host, i'', PathAtom (wp, wa), List.rev_append ra rp)

type ('fu, 'ret) re_query =
  | Nil : ('r, 'r) re_query
  | Any : ('r, 'r) re_query
  | Cons : 'a re_atom * ('f, 'r) re_query -> ('a -> 'f, 'r) re_query

let rec collect_re_query : type e r f.
       (e, f, r) Vif_uri.query
    -> int * (f, r) re_query * bool * (string * (Re.t * int)) list = function
  | Nil -> (0, Nil, false, [])
  | Any -> (0, Any, true, [])
  | Query_atom (s, a, q) ->
      let grps, wa, ra = re_atom_query 0 @@ from_t a in
      let total_grps, wq, b_any, rq = collect_re_query q in
      let total_grps = total_grps + grps in
      (total_grps, Cons (wa, wq), b_any, (s, (ra, grps)) :: rq)

let rec shift_lits : type a. int -> a re_atom -> a re_atom =
 fun shift -> function
  | Lit i -> Lit (i + shift)
  | Conv (x, f) -> Conv (shift_lits shift x, f)
  | Opt (m, x) -> Opt (m, shift_lits shift x)
  | Alt (m, x1, x2) -> Alt (m, shift_lits shift x1, shift_lits shift x2)
  | Seq (x1, x2) -> Seq (shift_lits shift x1, shift_lits shift x2)
  | Rep (i, x, r) -> Rep (shift + i, x, r)
  | Map (x, f) -> Map (shift_lits shift x, f)
  | Either (m, a, b) -> Either (m, shift_lits shift a, shift_lits shift b)

let rec permut_query : type r f.
    int -> int array -> (r, f) re_query -> (r, f) re_query =
 fun n permutation -> function
  | Nil -> Nil
  | Any -> Any
  | Cons (wa, wq) ->
      let shift = permutation.(n) in
      let wa = shift_lits shift wa in
      Cons (wa, permut_query (n + 1) permutation wq)

let re_query current_idx q =
  let grps, wq, b, rql = collect_re_query q in
  let rel = sort_query rql in
  let p = Utils.build_permutation current_idx (fun (_, (_, i)) -> i) rql rel in
  let wq = permut_query 0 p wq in
  (grps, wq, b, rel)

type ('f, 'r) re_url =
  | ReUrl : bool * ('f, 'x) re_path * ('x, 'r) re_query -> ('f, 'r) re_url

let re_url : type e f r.
    int -> (e, f, r) Vif_uri.t -> int * (f, r) re_url * Re.t =
 fun i -> function
  | Url (slash, p, q) -> begin
      let end_path =
        match slash with
        | No_slash -> Re.epsilon
        | Slash -> Re.char '/'
        | Maybe_slash -> Re.(opt @@ char '/')
      in
      let with_host, idx, wp, rp = re_path i p in
      match q with
      | Nil ->
          let re = Re.seq (List.rev (end_path :: rp)) in
          (idx, ReUrl (with_host, wp, Nil), re)
      | Any ->
          let end_re = Re.(opt @@ seq [ Re.char '?'; rep any ]) in
          let re = Re.seq (List.rev_append rp [ end_path; end_re ]) in
          (idx, ReUrl (with_host, wp, Nil), re)
      | _ ->
          let grps, wq, any_query, rel = re_query idx q in
          let query_sep = Ext.query_sep ~any:any_query in
          let add_around_query =
            if not any_query then fun x -> x else fun l -> Re.(rep any) :: l
          in
          let fn l (s, (re, _)) = Re.seq [ Re.str (s ^ "="); re ] :: l in
          let re =
            rel
            |> List.fold_left fn []
            |> Utils.intersperse query_sep
            |> add_around_query
            |> List.rev
            |> add_around_query
          in
          let re =
            Re.seq @@ List.rev_append rp (end_path :: Re.char '?' :: re)
          in
          (idx + grps, ReUrl (with_host, wp, wq), re)
    end

let get_re url =
  let _, _, re = re_url 1 url in
  re

(** {3 Extraction.} *)

exception Tyre_exn of exn
(* NOTE(dinosaure): the goal of this exception is to dispatch correctly errors from
   [conv] values which may fail and exception from [Tyre.Internal.extract]. *)

(** Extracting atom is just a matter of following the witness. We just need to
    take care of counting where we are in the matching groups. *)
let extract_atom ~original rea s =
  try extract ~original rea s with exn -> raise (Tyre_exn exn)

(** Since path is in reversed order, we proceed by continuation. *)
let rec extract_path : type f x r.
    original:string -> (f, x) re_path -> Re.Group.t -> (x -> r) -> f -> r =
 fun ~original wp subs k ->
  match wp with
  | Start -> k
  | PathAtom (rep, rea) ->
      let v = extract_atom ~original rea subs in
      let k fn = k (fn v) in
      extract_path ~original rep subs k

(** Query are in the right order, we can proceed in direct style. *)
let rec extract_query : type x r.
    original:string -> (x, r) re_query -> Re.Group.t -> x -> r =
 fun ~original wq subs f ->
  match wq with
  | Nil -> f
  | Any -> f
  | Cons (rea, req) ->
      let v = extract_atom ~original rea subs in
      extract_query ~original req subs (f v)

let extract_url : type r f.
    original:string -> (f, r) re_url -> Re.Group.t -> f -> r =
 fun ~original (ReUrl (_with_host, wp, wq)) subs fn ->
  let k = extract_query ~original wq subs in
  let k = extract_path ~original wp subs k in
  k fn

let prepare_target ?host target =
  let target =
    match String.index_opt target '?' with
    | None -> target
    | Some _ ->
        let uri = Uri.of_string target in
        uri
        |> Uri.query
        |> sort_query
        |> Uri.with_query uri
        |> Uri.path_and_query
  in
  match host with
  | Some host -> Uri.pct_encode ~component:`Host host ^ target
  | None -> target

let extract url =
  let _idx, (ReUrl (with_host, _, _) as re_url), re = re_url 1 url in
  let re = Re.(compile @@ whole_string re) in
  fun ~fn ?host uri ->
    let host = if with_host then host else None in
    let str = prepare_target ?host (Uri.path_and_query uri) in
    let subs = Re.exec re str in
    extract_url ~original:str re_url subs fn

(** {4 Multiple match} *)

type ('socket, 'fu, 'return) req =
  | Request :
      Vif_method.t option * ('c, 'a) Vif_type.t
      -> ('socket, ('socket, 'c, 'a) Vif_request.t -> 'r, 'r) req

type ('socket, 'r) t =
  | Route :
      ('socket, 'f, 'x) req * ('e, 'x, 'r) Vif_uri.t * 'f
      -> ('socket, 'r) t

let route req t f = Route (req, t, f)

type ('socket, 'r) re_ex =
  | ReEx :
      ('socket, 'f, 'x) req * 'f * Re.Mark.t * ('x, 'r) re_url
      -> ('socket, 'r) re_ex

(* It's important to keep the order here, since Re will choose
   the first regexp if there is ambiguity.
*)
let rec build_info_list : type s r.
       (Vif_method.t option -> bool)
    -> int
    -> (s, r) t list
    -> Re.t list * (s, r) re_ex list =
 fun p idx -> function
  | [] -> ([], [])
  | Route ((Request (meth, _) as req), url, fn) :: l when p meth ->
      let idx, re_url, re = re_url idx url in
      let rel, wl = build_info_list p idx l in
      let id, re = Re.mark re in
      (re :: rel, ReEx (req, fn, id, re_url) :: wl)
  | Route (Request _, _, _) :: l -> build_info_list p idx l

let build_info_list p l =
  let rel, wl = build_info_list p 1 l in
  (Re.(compile @@ whole_string @@ alt rel), wl)

let build_info : type s r.
       (s, r) t list
    -> (Re.re * (s, r) re_ex list) Vif_method.Map.t
       * (Re.re * (s, r) re_ex list) =
 fun l ->
  (* First figure out what methods the routes match *)
  (* We abuse Vif_method.Map as a set *)
  let fn : type s r. 'acc -> (s, r) t -> 'acc =
   fun acc r ->
    match r with
    | Route (Request (None, _), _, _) -> acc
    | Route (Request (Some meth, _), _, _) -> Vif_method.Map.add meth () acc
  in
  let methods = List.fold_left fn Vif_method.Map.empty l in
  let methods =
    Vif_method.Map.mapi
      (fun meth () ->
        build_info_list
          (function None -> true | Some meth' -> Vif_method.equal meth meth')
          l)
      methods
  and jokers = build_info_list Option.is_none l in
  (methods, jokers)

type 'socket request = {
    extract:
      'c 'a.
         Vif_method.t option
      -> ('c, 'a) Vif_type.t
      -> ('socket, 'c, 'a) Vif_request.t option
}

let rec find_and_trigger : type socket r.
       original:string
    -> socket request
    -> Re.Group.t
    -> (socket, r) re_ex list
    -> r =
 fun ~original e subs -> function
  | [] -> raise Not_found
  | ReEx (Request (meth, c), fn, id, re_url) :: l ->
      if Re.Mark.test subs id then
        begin match e.extract meth c with
        | None -> find_and_trigger ~original e subs l
        | Some v ->
            begin try extract_url ~original re_url subs (fn v)
            with Tyre_exn exn ->
              Log.debug (fun m ->
                  m "route converter raised exception: %a" Fmt.exn exn);
              find_and_trigger ~original e subs l
            end
        end
      else find_and_trigger ~original e subs l

let match_ (methods, jokers) meth s =
  match Vif_method.Map.find_opt meth methods with
  | Some (re, wl) ->
      begin match Re.exec_opt re s with
      | None -> None
      | Some subs -> Some (subs, wl)
      end
  | None ->
      begin match Re.exec_opt (fst jokers) s with
      | None -> None
      | Some subs -> Some (subs, snd jokers)
      end

let dispatch : type s r c.
       default:((s, c, string) Vif_request.t -> string -> r)
    -> (s, r) t list
    -> meth:Vif_method.t
    -> request:s request
    -> ?host:string
    -> string
    -> r =
 fun ~default l ->
  let info = build_info l in
  fun ~meth ~request:e ?host target ->
    let str = prepare_target ?host target in
    match match_ info meth str with
    | None -> default (Option.get (e.extract None Any)) str
    | Some (subs, wl) ->
        begin try find_and_trigger ~original:str e subs wl
        with Not_found -> default (Option.get (e.extract None Any)) str
        end