(** Abstract Renderer

    Abstract generic renderer, with tools to ease and unify the output of
    computations on a given program in any chosen format.
 *)

module RawAsm = Asm_info.NoAnnot

type addr_range_t = RawAsm.addr_t * RawAsm.addr_t
type box_data_t = {
  box_file : string;
  box_line : int;
  box_col : int;
  box_class_id : int;
}
type addr_range_tag_t = int
type tag_addr_range_t = TaggedRange of addr_range_tag_t * addr_range_t *
                                       box_data_t

module RangeTagMap = Map.Make(struct
    type t = addr_range_tag_t
    let compare = compare
  end)

type render_data_t = {
  render_prog_path : string ;
  render_prog : RawAsm.asm_info_t ;
  render_boxes : tag_addr_range_t list ;
  render_prev_address : RawAsm.addr_t RawAsm.AddrMap.t ;
}

type multi_render_data_t = render_data_t list

type render_box_event_t =
    BoxStart of addr_range_tag_t * RawAsm.addr_t * box_data_t
  | BoxEnd of addr_range_tag_t * RawAsm.addr_t * box_data_t

type asm_annot_type = {
  events : render_box_event_t list
}
module AnnotAsm = Asm_info.S(struct type instr_annot_t = asm_annot_type end)

(** Make a map mapping addresses to the preceding instruction address appearing
    in `prog` *)
let make_prev_address prog =
  let _, addr_map =
    List.fold_left (fun (prev_addr, cmap) sub ->
        List.fold_left (fun (prev_addr, cmap) instr ->
            (instr.RawAsm.instr_addr,
             RawAsm.AddrMap.add instr.RawAsm.instr_addr prev_addr cmap)
          )
          (prev_addr, cmap)
          sub.RawAsm.sub_asm
      )
      (0, RawAsm.AddrMap.empty)
      prog
  in
  addr_map

(** Initializes a rendering data structure on the given program *)
let init_render_data prog path =
  {
    render_prog_path = path ;
    render_prog = prog ;
    render_boxes = [] ;
    render_prev_address = make_prev_address prog ;
  }

let fresh_box_id =
  let next_id = ref 0 in
  fun () ->
    let out = !next_id in
    incr next_id;
    out

(** [add_box render_data range] adds a box around an address range, inclusive
    of both bounds, in [render_data] and returns the new render data and the
    newly inserted element's id *)
let add_box render_data range opt_data : render_data_t * int =
  let elt_id = fresh_box_id () in

  let data = (match opt_data with
      | Some data -> data
      | None -> {
          box_file = "";
          box_line = 0;
          box_col = 0;
          box_class_id = 0;
        })
  in
  {
    render_data with
    render_boxes =
      (TaggedRange(elt_id, range, data)) :: render_data.render_boxes
  }, elt_id


(** Same as [add_box], with an included start, excluded end for range. *)
let add_box_excl render_data (addr_beg, addr_end) opt_data =
  let lower_than_end addr = addr <= addr_end in
  (match RawAsm.AddrMap.find_last_opt
           lower_than_end render_data.render_prev_address with
  | None ->
    Format.eprintf "Box end address %x not found, ignoring box.@."
      addr_end ;
    raise Not_found
  | Some (_, end_bound) -> add_box render_data (addr_beg, end_bound) opt_data
  )

(** [to_file renderer render_data path] renders the given [render_data] to a
    file at [path] using the provided specific renderer *)
let to_file renderer render_data path =
  let handle = open_out path in
  render_data
  |> renderer
  |> output_string handle ;
  close_out handle

(** [to_string renderer multi_render_data] renders the given
    [multi_render_data] to a returned string using the provided specific
    renderer.  This function is provided for unification and clarity, but is an
    alias for [renderer].
*)
let to_string renderer multi_render_data =
  renderer multi_render_data

let extract_event_addr box_event = match box_event with
  | BoxStart(_, addr, _) -> addr
  | BoxEnd(_, addr, _) -> addr

(** [boxes_to_events render_data] converts the list of boxes of `render_data`
    to a list of events, that is, an address-ordered list of
    `render_box_event_t`
*)
let boxes_to_events render_data =
  let compare_events e1 e2 =
    match (compare (extract_event_addr e1) (extract_event_addr e2)) with
    | 0 -> (match e1, e2 with
        | BoxEnd(_), BoxStart(_) -> -1
        | BoxStart(_), BoxEnd(_) -> 1
        | BoxStart(id1, _, _), BoxStart(id2, _, _) -> compare id1 id2
        | BoxEnd(id1, _, _), BoxEnd(id2, _, _) -> compare id1 id2
      )
    | n -> n
  in

  let unordered_events = List.fold_left
      (fun accu (TaggedRange(elt_id, (range_beg, range_end), data)) ->
         (BoxStart(elt_id, range_beg, data))
         :: (BoxEnd(elt_id, range_end, data))
         :: accu
      )
      []
      render_data.render_boxes
  in
  let ordered_events = List.sort
      compare_events
      unordered_events
  in
  ordered_events

(** Transforms `render_data` into an `asm_info_t` whose asm lines are annotated
    with events, for ease of rendering *)
let render_data_to_annotated_asm render_data : AnnotAsm.asm_info_t =
  let render_events = boxes_to_events render_data in

  (* Splits `events` into two lists: a list of events whose addresses match
      `addr` and the remaining tail list. *)
  let extract_current_events events addr =
    let reached_stop event = (extract_event_addr event) > addr in
    let rec fold_until accu l = match l with
      | (hd::_) as cur_l when reached_stop hd -> (List.rev accu), cur_l
      | hd::tl -> fold_until (hd::accu) tl
      | [] -> (List.rev accu), []
    in
    fold_until [] events
  in

  let map_asm render_events lst =
    let asm, leftover_events = List.fold_left
      (fun (accu, render_events) asm_instr ->
         let current_events, future_events =
           extract_current_events render_events RawAsm.(asm_instr.instr_addr)
         in
         (
           {
             AnnotAsm.instr_addr = RawAsm.(asm_instr.instr_addr) ;
             AnnotAsm.instr_bytes = RawAsm.(asm_instr.instr_bytes) ;
             AnnotAsm.instr_asm =  RawAsm.(asm_instr.instr_asm) ;
             AnnotAsm.instr_annot = {
               events = current_events
             }
           }::accu
         ), future_events
      )
      ([], render_events)
      lst
    in
    (List.rev asm, leftover_events)
  in
  let annotated_asm, _ =
    List.fold_left (fun (accu, render_events) sub ->
        let mapped_asm, n_render_events =
          map_asm render_events RawAsm.(sub.sub_asm) in
        (
          {
            AnnotAsm.sub_section = RawAsm.(sub.sub_section) ;
            AnnotAsm.sub_name = RawAsm.(sub.sub_name) ;
            AnnotAsm.sub_addr = RawAsm.(sub.sub_addr) ;
            AnnotAsm.sub_asm = mapped_asm ;
          } :: accu
        ), n_render_events
      )
      ([], render_events)
      render_data.render_prog
  in
  List.rev annotated_asm