(** Matches lines of assembly according to their DWARF line info *)

type blockid_t = int

module BlockidMap = Map.Make(struct
    type t = blockid_t
    let compare = compare
  end)

type asm_block_info = {
  asm_bi_dwarfhead : Dwarf.line_number_header ; (* FIXME overkill to carry
                                                   around? *)
  asm_bi_dwarfline : Dwarf.line_number_registers ;
  asm_bi_blockid : blockid_t ;
}

(** Generates a fresh block id *)
let fresh_blockid =
  let next_id = ref 0 in
  fun () ->
    let out = !next_id in
    incr next_id ;
    out

(** Gets the file path for a given file id and header data *)
let extract_file_name file_id files_list =
  let file_data = List.nth files_list (file_id - 1) in
  file_data.Dwarf.lnfe_path
  |> List.to_seq
  |> String.of_seq


(** Gets the file path for the given blk, instead of its id *)
let get_file_name blk = Dwarf.(
    let file_id = blk.asm_bi_dwarfline.lnr_file
                  |> Z.to_int in
    extract_file_name file_id blk.asm_bi_dwarfhead.lnh_file_names
  )

(** Compares two block infos wrt. class equivalence *)
let asm_bi_compare blk1 blk2 =
  let select_file blk = get_file_name blk in
  let select_line blk = Dwarf.(blk.asm_bi_dwarfline.lnr_line) in
  let select_col blk = Dwarf.(blk.asm_bi_dwarfline.lnr_column) in

  match compare (select_file blk1) (select_file blk2) with
  | 0 -> (match compare (select_line blk1) (select_line blk2) with
      | 0 -> compare (select_col blk1) (select_col blk2)
      | n -> n)
  | n -> n

module BoxRootSet = Set.Make(struct
    type t = asm_block_info
    let compare = asm_bi_compare
  end)

type matcher_state_t = {
  ms_block_infos : asm_block_info BlockidMap.t ;
  ms_block_classes : blockid_t BlockidMap.t ;
  ms_block_roots : BoxRootSet.t;
}

let empty_state = {
  ms_block_infos = BlockidMap.empty ;
  ms_block_classes = BlockidMap.empty ;
  ms_block_roots = BoxRootSet.empty ;
}

(** Add a new block to the state, returning the new state and the class id of
    the added block *)
let add_block_bi state blockinfo =
  match BoxRootSet.find_opt blockinfo state.ms_block_roots with
  | None ->
    {
      ms_block_infos = BlockidMap.add
          blockinfo.asm_bi_blockid
          blockinfo
          state.ms_block_infos ;
      ms_block_classes = BlockidMap.add
          blockinfo.asm_bi_blockid
          blockinfo.asm_bi_blockid
          state.ms_block_classes ;
      ms_block_roots = BoxRootSet.add blockinfo state.ms_block_roots ;
    }, blockinfo.asm_bi_blockid
  | Some root_block ->
    { state with
      ms_block_infos = BlockidMap.add
          blockinfo.asm_bi_blockid
          blockinfo
          state.ms_block_infos ;
      ms_block_classes = BlockidMap.add
          blockinfo.asm_bi_blockid
          root_block.asm_bi_blockid
          state.ms_block_classes ;
    }, root_block.asm_bi_blockid

(** Same as [add_block_bi], but takes simply a [Dwarf.line_number_registers] as 
    argument and generates a fresh id by itself.
    WARNING! Do not mix automatically and manually generated ids.
*)
let add_block state lnh lnr =
  add_block_bi state {
    asm_bi_dwarfhead = lnh ;
    asm_bi_dwarfline = lnr ;
    asm_bi_blockid = fresh_blockid ()
  }

(** Gets the class ID of a previously added blockinfo. If the blockinfo was
    never added, raises Not_found. *)
let block_class state blockinfo =
  BlockidMap.find blockinfo.asm_bi_blockid state.ms_block_classes