open Types
open LlbcAst
include Charon.LlbcAstUtils
open Collections

module FunIdOrderedType : OrderedType with type t = fun_id = struct
  type t = fun_id

  let compare = compare_fun_id
  let to_string = show_fun_id
  let pp_t = pp_fun_id
  let show_t = show_fun_id
end

module FunIdMap = Collections.MakeMap (FunIdOrderedType)
module FunIdSet = Collections.MakeSet (FunIdOrderedType)

let body_as_body = Charon.LlbcAstUtils.body_as_structured

let body_as_body_exn file line f =
  match body_as_body f with
  | Some body -> body
  | None -> Errors.craise_opt_span file line None "Not a LLBC body"

let body_is_known (b : body) : bool = Option.is_some (body_as_body b)

let body_as_target_dispatch (b : body) :
    (string * Types.fun_decl_ref) list option =
  match b with
  | TargetDispatchBody targets -> Some targets
  | _ -> None

let body_is_target_dispatch (b : body) : bool =
  Option.is_some (body_as_target_dispatch b)

(** A function body is translatable if it is either a structured body (normal
    case) or a target dispatch body (multi-target). *)
let body_is_translatable (b : body) : bool =
  body_is_known b || body_is_target_dispatch b

let lookup_fun_sig (fun_id : fun_id) (fun_decls : fun_decl FunDeclId.Map.t) :
    bound_fun_sig =
  match fun_id with
  | FRegular id ->
      let fun_decl = FunDeclId.Map.find id fun_decls in
      bound_fun_sig_of_decl fun_decl
  | FBuiltin aid -> Builtin.get_builtin_fun_sig aid

(** Return the opaque declarations found in the crate, which are also *not
    builtin*.

    [filter_builtin]: if [true], do not consider as opaque the external
    definitions that we will map to definitions from the standard library.

    Remark: the list of functions also contains the list of opaque global
    bodies. *)
let crate_get_opaque_non_builtin_decls (k : crate) (filter_builtin : bool)
    (type_decls : type_decl TypeDeclId.Map.t)
    (fun_decls : fun_decl FunDeclId.Map.t) : type_decl list * fun_decl list =
  let open ExtractBuiltin in
  let ctx = Charon.NameMatcher.ctx_from_crate k in
  let is_opaque_fun (d : fun_decl) : bool =
    (not (body_is_known d.body))
    && ((not filter_builtin)
       || (not
             (NameMatcherMap.mem ctx d.item_meta.name (builtin_globals_map ())))
          && not (NameMatcherMap.mem ctx d.item_meta.name (builtin_funs_map ()))
       )
  in
  let is_opaque_type (d : type_decl) : bool =
    d.kind = Opaque
    && ((not filter_builtin)
       || not (NameMatcherMap.mem ctx d.item_meta.name (builtin_types_map ())))
  in
  (* Note that by checking the function bodies we also the globals *)
  ( List.filter is_opaque_type (TypeDeclId.Map.values type_decls),
    List.filter is_opaque_fun (FunDeclId.Map.values fun_decls) )

(** Return true if the crate contains opaque declarations, ignoring the builtin
    definitions. *)
let crate_has_opaque_non_builtin_decls (k : crate) (filter_builtin : bool)
    (type_decls : type_decl TypeDeclId.Map.t)
    (fun_decls : fun_decl FunDeclId.Map.t) : bool =
  crate_get_opaque_non_builtin_decls k filter_builtin type_decls fun_decls
  <> ([], [])

(** Strip trailing [PeTarget] elements from a name.

    Multi-target extraction appends [PeTarget] to per-target function names.
    This element doesn't participate in pattern matching (the pattern generator
    skips it), so we strip it before calling [name_to_pattern] to avoid
    triggering its roundtrip assertion. *)
let strip_target_suffix (n : name) : name =
  match List.rev n with
  | Types.PeTarget _ :: rest -> List.rev rest
  | _ -> n

let strip_target_or_instantiated_suffix (n : name) : name =
  let rec strip_all (n : name) : name =
    match n with
    | (PeTarget _ | PeInstantiated _) :: rest -> strip_all rest
    | _ -> List.rev n
  in
  strip_all (List.rev n)

(** Extract and strip any trailing [PeTarget] element from a name, returning the
    cleaned name and an optional target suffix string (with [-] replaced by
    [_]). *)
let extract_target_suffix (name : name) : name * string option =
  match Collections.List.last name with
  | PeTarget target ->
      let target = String.concat "_" (String.split_on_char '-' target) in
      (Collections.List.prefix (List.length name - 1) name, Some target)
  | _ -> (name, None)

let add_target_suffix (name : name) (target_suffix : string option) : name =
  match target_suffix with
  | None -> name
  | Some target -> name @ [ PeTarget target ]

let name_to_pattern (span : Meta.span option) (ctx : Charon.NameMatcher.ctx)
    (c : Charon.NameMatcher.to_pat_config) (n : name) =
  let n = strip_target_suffix n in
  if !Config.fail_hard then Charon.NameMatcher.name_to_pattern ctx c n
  else
    try Charon.NameMatcher.name_to_pattern ctx c n
    with Not_found | Assert_failure _ ->
      [%craise_opt_span] span "Could not convert the name to a pattern"

let name_with_crate_to_pattern_string (span : Meta.span option)
    (crate : LlbcAst.crate) (n : Types.name) : string =
  let mctx = Charon.NameMatcher.ctx_from_crate crate in
  let c : Charon.NameMatcher.to_pat_config =
    {
      tgt = TkPattern;
      use_trait_decl_refs = Config.match_patterns_with_trait_decl_refs;
    }
  in
  let pat = name_to_pattern span mctx c n in
  Charon.NameMatcher.pattern_to_string { tgt = TkPattern } pat

let name_with_generics_to_pattern (span : Meta.span option)
    (ctx : Charon.NameMatcher.ctx) (c : Charon.NameMatcher.to_pat_config)
    (params : generic_params) (n : Charon.Types.name) (args : generic_args) =
  if !Config.fail_hard then
    Charon.NameMatcher.name_with_generics_to_pattern ctx c params n args
  else
    try Charon.NameMatcher.name_with_generics_to_pattern ctx c params n args
    with Not_found | Assert_failure _ ->
      [%craise_opt_span] span "Could not convert the name to a pattern"

let name_with_generics_crate_to_pattern_string (span : Meta.span option)
    (crate : LlbcAst.crate) (n : Types.name) (params : Types.generic_params)
    (args : Types.generic_args) : string =
  let mctx = Charon.NameMatcher.ctx_from_crate crate in
  let c : Charon.NameMatcher.to_pat_config =
    {
      tgt = TkPattern;
      use_trait_decl_refs = Config.match_patterns_with_trait_decl_refs;
    }
  in
  let pat = name_with_generics_to_pattern span mctx c params n args in
  Charon.NameMatcher.pattern_to_string { tgt = TkPattern } pat

let trait_impl_with_crate_to_pattern_string (span : Meta.span option)
    (crate : LlbcAst.crate) (trait_decl : LlbcAst.trait_decl)
    (trait_impl : LlbcAst.trait_impl) : string =
  name_with_generics_crate_to_pattern_string span crate
    trait_decl.item_meta.name trait_decl.generics trait_impl.impl_trait.generics

(** Return true if the statement contains an instruction which breaks the
    control flow, at the exception of panics (that is: a break, a continue or a
    return) *)
let statement_has_break_continue_return (st : statement) : bool =
  let visitor =
    object
      inherit [_] iter_statement
      method! visit_Break _ _ = raise Utils.Found
      method! visit_Continue _ _ = raise Utils.Found
      method! visit_Return _ = raise Utils.Found
    end
  in
  try
    visitor#visit_statement () st;
    false
  with Utils.Found -> true

(** Return true if the block contains a statement which breaks the control flow,
    at the exception of panics (that is: a break, a continue or a return) *)
let block_has_break_continue_return (st : block) : bool =
  let visitor =
    object
      inherit [_] iter_statement
      method! visit_Break _ _ = raise Utils.Found
      method! visit_Continue _ _ = raise Utils.Found
      method! visit_Return _ = raise Utils.Found
    end
  in
  try
    visitor#visit_block () st;
    false
  with Utils.Found -> true

(** Compute the size of a function body - we count the number of statements and
    blocks *)
let compute_body_size (b : body) : int =
  let size = ref 0 in
  let incr () = size := !size + 1 in
  let visitor =
    object
      inherit [_] iter_statement as super

      method! visit_statement env st =
        incr ();
        super#visit_statement env st

      method! visit_block env st =
        incr ();
        super#visit_block env st
    end
  in
  let () =
    match b with
    | StructuredBody body -> visitor#visit_block () body.body
    | _ -> ()
  in
  !size

(** Compute the size of a function - we count the number of statements and
    blocks *)
let compute_fun_decl_size (f : fun_decl) : int = compute_body_size f.body
