Source file salto_id.ml

(** @author Pierre Lermusiaux <pierre.lermusiaux@inria.fr>
    Copyright © Inria 2022-2023
*)

open Location

let star_ident_regex = Str.regexp {|\*\([_A-Za-z][_'A-Za-z]+\)\*|}

module OCamlPath = Path

let hash_fold_ident acc id = Base.Hash.fold_string acc (Ident.unique_name id)

module Id = struct
  type t =
    | FromIdent of Ident.t
    | FromNamedId of string Location.loc * Ident.t
    | Internal of string * int

  let equal (id1 : t) (id2 : t) : bool =
    match (id1, id2) with
    | ( (FromIdent ident1 | FromNamedId (_, ident1)),
        (FromIdent ident2 | FromNamedId (_, ident2)) ) ->
        Ident.same ident1 ident2
    | Internal (name1, ref1), Internal (name2, ref2) ->
        String.equal name1 name2 && ref1 = ref2
    | _, _ -> false

  let compare (id1 : t) (id2 : t) : int =
    match (id1, id2) with
    | ( (FromIdent ident1 | FromNamedId (_, ident1)),
        (FromIdent ident2 | FromNamedId (_, ident2)) ) ->
        Ident.compare ident1 ident2
    | Internal (name1, ref1), Internal (name2, ref2) ->
        let n = String.compare name1 name2 in
        if n <> 0 then n else Int.compare ref1 ref2
    | _, Internal _ -> -1
    | Internal _, _ -> 1

  (** reference to a counter for internal identifiers, used to guarantee
    uniqueness *)
  let internal_counter = ref 0

  let internal (name : string) : t =
    incr internal_counter;
    Internal (name, !internal_counter)

  let from_ident (id : Ident.t) : t = FromIdent id

  let from_named_id (name : string loc) (id : Ident.t) : t =
    if Str.string_match star_ident_regex name.txt 0 && Location.is_none name.loc
    then FromNamedId ({ name with txt = Str.matched_group 1 name.txt }, id)
    else FromNamedId (name, id)

  let name : t -> string = function
    | FromIdent id -> Ident.name id
    | FromNamedId (name, _) ->
        if Str.string_match star_ident_regex name.txt 0 then
          Str.matched_group 1 name.txt
        else name.txt
    | Internal (name, stamp) -> Format.sprintf "%s_%d" name stamp

  let pp fmt id = Format.pp_print_string fmt (name id)

  let name_loc : t -> string loc = function
    | FromIdent id -> Location.mknoloc (Ident.name id)
    | FromNamedId (name, _) ->
        if Str.string_match star_ident_regex name.txt 0 then
          { name with txt = Str.matched_group 1 name.txt }
        else name
    | Internal (name, stamp) ->
        Location.mknoloc (Format.sprintf "%s_%d" name stamp)

  let lident_loc : t -> Longident.t loc = function
    | FromIdent id -> Location.mknoloc (Longident.Lident (Ident.name id))
    | FromNamedId (name, _) ->
        if Str.string_match star_ident_regex name.txt 0 then
          { name with txt = Lident (Str.matched_group 2 name.txt) }
        else { name with txt = Lident name.txt }
    | Internal (name, stamp) ->
        Location.mknoloc (Longident.Lident (Format.sprintf "%s_%d" name stamp))

  let ocaml_ident = function
    | FromIdent id | FromNamedId (_, id) -> id
    | Internal _ ->
        invalid_arg
          "Id.ocaml_ident: internal identifier does not have an OCaml \
           identifier"

  let is_internal = function Internal _ -> true | _ -> false

  let is_persistent = function
    | FromIdent id | FromNamedId (_, id) -> Ident.persistent id
    | Internal _ -> false

  let hash_fold s =
    let open Base.Hash in
    function
    | FromIdent id | FromNamedId (_, id) -> fold_int (hash_fold_ident s id) 1
    | Internal (name, i) -> fold_int (fold_string (fold_int s i) name) 2
end

module Longident = struct
  include Longident

  let head_lid (id : Id.t) : Longident.t =
    let name = Id.name id in
    if Str.string_match star_ident_regex name 0 then
      Lident (Str.matched_group 1 name)
    else Lident name

  let rec strip : Longident.t -> Longident.t =
    let rec strip0 = function
      | Lident name ->
          if Str.string_match star_ident_regex name 0 then
            let name = Str.matched_group 1 name in
            if name = "predef" then None else Some (Lident name)
          else Some (Lident name)
      | Ldot (path, name) -> begin
          match strip0 path with
          | None -> Some (Lident name)
          | Some path -> Some (Ldot (path, name))
        end
      | Lapply (lid1, lid2) ->
          let ( let* ) = Option.bind in
          let* lid1 = strip0 lid1 in
          let* lid2 = strip0 lid2 in
          Some (Lapply (lid1, lid2))
    in
    function
    | Lident name ->
        if Str.string_match star_ident_regex name 0 then
          Lident (Str.matched_group 1 name)
        else Lident name
    | Ldot (path, name) -> begin
        match strip0 path with
        | None -> Lident name
        | Some path -> Ldot (path, name)
      end
    | Lapply (lid1, lid2) ->
        let lid1 = strip lid1 and lid2 = strip lid2 in
        Lapply (lid1, lid2)

  let rec hash_fold acc =
    let open Base.Hash in
    function
    | Longident.Lident s -> fold_int (fold_string acc s) 0
    | Longident.Ldot (p, s) -> fold_int (hash_fold (fold_string acc s) p) 1
    | Longident.Lapply (p1, p2) -> hash_fold (hash_fold acc p1) p2

  let rec compare (lid1 : Longident.t) (lid2 : Longident.t) =
    match (lid1, lid2) with
    | Lident s1, Lident s2 -> String.compare s1 s2
    | Lident _, (Ldot _ | Lapply _) -> -1
    | (Ldot _ | Lapply _), Lident _ -> 1
    | Ldot (l1, s1), Ldot (l2, s2) ->
        let n = String.compare s1 s2 in
        if n <> 0 then n else compare l1 l2
    | Ldot _, Lapply _ -> -1
    | Lapply _, Ldot _ -> 1
    | Lapply (l1l, l1r), Lapply (l2l, l2r) ->
        let n = compare l1l l1r in
        if n <> 0 then n else compare l2l l2r

  let rec pp fmt =
    let open Format in
    function
    | Longident.Lident s -> pp_print_string fmt s
    | Ldot (x, s) -> fprintf fmt "%a.%s" pp x s
    | Lapply (x1, x2) -> fprintf fmt "%a(%a)" pp x1 pp x2
end

module Path = struct
  type flat =
    | Ident of Id.t
    | Dot of Id.t * flat
  type path =
    | Appl of proj * proj
    | Flat of flat
  and proj = Proj of path * string list

  type 'a t = proj Location.loc * 'a

  let rec compare_dots path names1 names2 =
    match (path, names2) with
    | _, [] -> 1
    | Ident id, name :: names2 ->
        let h = String.compare (Id.name id) name in
        if h <> 0 then h else List.compare String.compare names1 names2
    | Dot (id, path), name :: names2 ->
        let h = String.compare (Id.name id) name in
        if h <> 0 then h else compare_dots path names1 names2

  let rec compare_flat path1 names1 path2 names2 =
    match (path1, path2) with
    | Ident id1, Ident id2 ->
        let h = Id.compare id1 id2 in
        if h <> 0 then h else List.compare String.compare names1 names2
    | Dot (id1, path1), Dot (id2, path2) ->
        let h = Id.compare id1 id2 in
        if h <> 0 then h else compare_flat path1 names1 path2 names2
    | Dot (id1, path1), Ident id2 ->
        let h = Id.compare id1 id2 in
        if h <> 0 then h else compare_dots path1 names1 names2
    | Ident id1, Dot (id2, path2) ->
        let h = Id.compare id1 id2 in
        if h <> 0 then h else -compare_dots path2 names2 names1

  let rec compare_proj (Proj (path1, names1)) (Proj (path2, names2)) =
    compare_path path1 names1 path2 names2

  and compare_path path1 names1 path2 names2 =
    match (path1, path2) with
    | Flat _, Appl _ -> 1
    | Flat path1, Flat path2 -> compare_flat path1 names1 path2 names2
    | Appl (fun1, arg1), Appl (fun2, arg2) ->
        let h = List.compare String.compare names1 names2 in
        if h <> 0 then h
        else
          let h = compare_proj fun1 fun2 in
          if h <> 0 then h else compare_proj arg1 arg2
    | _ -> -1

  let compare (p1, _) (p2, _) = compare_proj p1.txt p2.txt

  type typed_path = Types.type_expr t

  let get_desc ((_, desc) : 'a t) : 'a = desc

  let void (path, _) = (path, ())

  let rec equal_dots path names1 names2 =
    match (path, names2) with
    | _, [] -> false
    | Ident id, name :: names2 ->
        String.equal (Id.name id) name && List.equal String.equal names1 names2
    | Dot (id, path), name :: names2 ->
        String.equal (Id.name id) name && equal_dots path names1 names2

  let rec equal_flat path1 names1 path2 names2 =
    match (path1, path2) with
    | Ident id1, Ident id2 ->
        Id.equal id1 id2 && List.equal String.equal names1 names2
    | Dot (id1, path1), Dot (id2, path2) ->
        Id.equal id1 id2 && equal_flat path1 names1 path2 names2
    | Dot (id1, path1), Ident id2 ->
        Id.equal id1 id2 && equal_dots path1 names1 names2
    | Ident id1, Dot (id2, path2) ->
        Id.equal id1 id2 && equal_dots path2 names2 names1

  let rec equal_proj (Proj (path1, names1)) (Proj (path2, names2)) =
    equal_path path1 names1 path2 names2

  and equal_path path1 names1 path2 names2 =
    match (path1, path2) with
    | Flat path1, Flat path2 -> equal_flat path1 names1 path2 names2
    | Appl (fun1, arg1), Appl (fun2, arg2) ->
        List.equal String.equal names1 names2
        && equal_proj fun1 fun2 && equal_proj arg1 arg2
    | _ -> false

  let equal (p1, _) (p2, _) = equal_proj p1.txt p2.txt

  let equal_id ({ txt = p; _ }, _) y =
    match p with Proj (Flat (Ident x), []) -> Id.equal x y | _ -> false

  let rec pp_flat fmt = function
    | Ident id -> Id.pp fmt id
    | Dot (id, path) -> Format.fprintf fmt "%a.%a" Id.pp id pp_flat path

  let rec pp_proj fmt = function
    | Proj (path, []) -> Format.fprintf fmt "%a" pp_path path
    | Proj (path, names) ->
        let open Format in
        let pp_sep fmt () = pp_print_char fmt '.' in
        let pp_names = pp_print_list ~pp_sep pp_print_string in
        fprintf fmt "%a.%a" pp_path path pp_names names

  and pp_path fmt = function
    | Appl (path1, path2) ->
        Format.fprintf fmt "%a(%a)" pp_proj path1 pp_proj path2
    | Flat path -> pp_flat fmt path

  let pp fmt (path, _) = pp_proj fmt path.txt

  let from_path : OCamlPath.t -> 'a -> Location.t -> 'a t =
    let rec from_path names : OCamlPath.t -> proj = function
      | Pident id -> Proj (Flat (Ident (Id.from_ident id)), names)
      | Pdot (path, name) -> from_path (name :: names) path
      | Papply (path1, path2) ->
          Proj (Appl (from_path [] path1, from_path [] path2), names)
    in
    fun path desc loc -> (Location.mkloc (from_path [] path) loc, desc)

  let from_ident (id : Id.t) (fields : string list) (desc : 'a) : 'a t =
    ({ (Id.name_loc id) with txt = Proj (Flat (Ident id), fields) }, desc)

  let append (id : Id.t) : 'a t -> 'a t = function
    | { txt = Proj (Flat path, names); loc }, desc ->
        ({ txt = Proj (Flat (Dot (id, path)), names); loc }, desc)
    | _ -> invalid_arg "Path.append: cannot append on an applicative path"

  let[@tail_mod_cons] rec split_flat (names : string list) = function
    | Ident id -> Id.name id :: names
    | Dot (id, path) -> Id.name id :: split_flat names path

  let rec split_proj (Proj (path, names)) acc = split_path names acc path

  and split_path names acc = function
    | Appl (path1, path2) -> split_proj path2 acc |> split_proj path1
    | Flat (Ident head) -> (head, names) :: acc
    | Flat (Dot (head, path)) -> (head, split_flat names path) :: acc

  let split (path, _) = split_proj path.txt []

  let head = function
    | { txt = Proj (Flat (Ident head | Dot (head, _)), _); _ }, _ -> head
    | _ ->
        invalid_arg "Path.head: cannot return the head of an applicative path"

  let rec fold_flat f v = function
    | Ident id -> f v (Id.name id)
    | Dot (id, path) -> fold_flat f (f v (Id.name id)) path

  let rec fold_proj ~init ~proj ~apply (Proj (path, names)) =
    List.fold_left proj (fold_path ~init ~proj ~apply path) names

  and fold_path ~init ~proj ~apply = function
    | Appl (path1, path2) ->
        apply
          (fold_proj ~init ~proj ~apply path1)
          (fold_proj ~init ~proj ~apply path2)
    | Flat (Ident id) -> init id
    | Flat (Dot (id, path)) -> fold_flat proj (init id) path

  let fold ~init ~proj ~apply (path, _) = fold_proj ~init ~proj ~apply path.txt

  let name_loc (path, _) =
    { path with txt = Format.asprintf "%a" pp_proj path.txt }

  let rec lident_of_flat lid : flat -> Longident.t = function
    | Ident id -> Ldot (lid, Id.name id)
    | Dot (id, path) -> lident_of_flat (Ldot (lid, Id.name id)) path

  let rec lident_of_proj : proj -> Longident.t =
    let ldot lid name : Longident.t = Ldot (lid, name) in
    function
    | Proj (Flat (Ident id), head :: names)
      when String.equal (Id.name id) "*predef*" ->
        List.fold_left ldot (Lident head) names
    | Proj (Flat (Dot (id, path)), names)
      when String.equal (Id.name id) "*predef*" ->
        let lid =
          match path with
          | Ident head -> Longident.head_lid head
          | Dot (head, path) -> lident_of_flat (Longident.head_lid head) path
        in
        List.fold_left ldot lid names
    | Proj (path, names) -> List.fold_left ldot (lident_of_path path) names

  and lident_of_path : path -> Longident.t = function
    | Appl (path1, path2) -> Lapply (lident_of_proj path1, lident_of_proj path2)
    | Flat (Ident head) -> Longident.head_lid head
    | Flat (Dot (head, path)) -> lident_of_flat (Longident.head_lid head) path

  let lident_loc (path, _) : Longident.t loc =
    { path with txt = lident_of_proj path.txt }
end

let stdlib_path : Path.path =
  Flat (Ident (Id.from_ident (Ident.create_persistent "Stdlib")))

let raise_path : Primitive.description Path.t =
  ( Location.mknoloc @@ Path.Proj (stdlib_path, [ "raise" ]),
    Primitive.simple ~name:"%raise" ~arity:1 ~alloc:true )

let reraise_path : Primitive.description Path.t =
  ( Location.mknoloc @@ Path.Proj (stdlib_path, [ "raise_notrace" ]),
    Primitive.simple ~name:"%reraise" ~arity:1 ~alloc:true )

let succ_path : Primitive.description Path.t =
  ( Location.mknoloc @@ Path.Proj (stdlib_path, [ "succ" ]),
    Primitive.simple ~name:"%succint" ~arity:1 ~alloc:true )

let pred_path : Primitive.description Path.t =
  ( Location.mknoloc @@ Path.Proj (stdlib_path, [ "pred" ]),
    Primitive.simple ~name:"%predint" ~arity:1 ~alloc:true )

let ge_path : Primitive.description Path.t =
  ( Location.mknoloc @@ Path.Proj (stdlib_path, [ ">=" ]),
    Primitive.simple ~name:"greaterequal" ~arity:1 ~alloc:true )

let le_path : Primitive.description Path.t =
  ( Location.mknoloc @@ Path.Proj (stdlib_path, [ "<=" ]),
    Primitive.simple ~name:"%lessequal" ~arity:1 ~alloc:true )

module Salto_env = struct
  include Map.Make (Id)

  let pp pp_elt fmt env =
    if is_empty env then Format.fprintf fmt "[]"
    else
      Format.fprintf fmt "@[[@[<hv 1> %a@]@ ]@]"
        (Format.pp_print_list
           ~pp_sep:(fun fmt () -> Format.fprintf fmt ",@ ")
           (fun fmt (x, a) ->
             Format.fprintf fmt "@[<hv 2>%a ->@ %a@]" Id.pp x pp_elt a ) )
        (bindings env)

  let hash_fold hash_fold_val acc m =
    fold (fun x v acc -> Id.hash_fold (hash_fold_val acc v) x) m acc
end

module Idents = Set.Make (Id)