open Maps
open BTL
open BTLtypes
open Machine
open DebugPrint
open PrintBTL
open Registers
open RTLcommonaux
open BTLcommonaux
open ExpansionOracle
open PrepassSchedulingOracle

let flatten_blk_basics ibf =
  let ib = ibf.entry in
  let last = ref None in
  let rec traverse_blk ib =
    match ib with
    | BF (_, _) ->
        last := Some ib;
        []
    | Bseq ((Bcond (_, _, _, _, iinfo) as ib1), ib2) -> (
        match iinfo.opt_info with
        | Some _ -> [ ib1 ] @ traverse_blk ib2
        | None ->
            last := Some ib;
            [])
    | Bseq (ib1, ib2) -> traverse_blk ib1 @ traverse_blk ib2
    | _ -> [ ib ]
  in
  let ibl = traverse_blk ib in
  (Array.of_list ibl, !last)

let is_a_cb = function Bcond _ -> true | _ -> false

let is_a_load = function Bload _ -> true | _ -> false

module SI = Set.Make (Int)

let find_array arr n =
  let index = ref None in
  (try
     Array.iteri
       (fun i n' ->
         match !index with
         | Some _ -> raise Exit
         | None -> if n = n' then index := Some i)
       arr
   with Exit -> ());
  get_some @@ !index

let count_cbs bseq olast indexes =
  let current_cbs = ref SI.empty in
  let cbs_above = Hashtbl.create 100 in
  let update_cbs n ib =
    print_btl_inst stderr ib;
    if is_a_cb ib then current_cbs := SI.add indexes.(n) !current_cbs
    else if is_a_load ib then Hashtbl.add cbs_above indexes.(n) !current_cbs
  in
  Array.iteri (fun n ib -> update_cbs n ib) bseq;
  (match olast with
  | Some last -> update_cbs (Array.length bseq) last
  | None -> ());
  cbs_above

let apply_schedule bseq olast positions =
  let bseq_new = Array.map (fun i -> bseq.(i)) positions in
  (if !config.has_non_trapping_loads then
   let fmap n = find_array positions n in
   let seq = Array.init (Array.length positions) (fun i -> i) in
   let fseq = Array.map fmap seq in
   let cbs_above_old = count_cbs bseq olast fseq in
   let cbs_above_new = count_cbs bseq_new olast seq in
   Array.iteri
     (fun n ib ->
       let n' = fseq.(n) in
       match ib with
       | Bload (t, a, b, c, d, e) ->
           let set_old = Hashtbl.find cbs_above_old n' in
           let set_new = Hashtbl.find cbs_above_new n' in
           if SI.subset set_old set_new then
             if get_some @@ e.opt_info then
               bseq_new.(n') <- Bload (AST.TRAP, a, b, c, d, e)
           else assert !config.has_non_trapping_loads
       | _ -> ())
     bseq);
  let ibl = Array.to_list bseq_new in
  let rec build_iblock = function
    | [] -> failwith "build_iblock: empty list"
    | [ ib ] -> ( match olast with Some last -> Bseq (ib, last) | None -> ib)
    | ib1 :: ib2 :: k -> Bseq (ib1, build_iblock (ib2 :: k))
  in
  build_iblock ibl

(** the useful one. Returns a hashtable with bindings of shape
 ** [(r,(t, n)], where [r] is a pseudo-register (Registers.reg),
 ** [t] is its class (according to [typing]), and [n] the number of
 ** times it's referenced as an argument in instructions of [seqa] ;
 ** and an arrray containg the list of regs referenced by each
 ** instruction, with a boolean to know whether it's as arg or dest *)
let reference_counting (seqa : (iblock * Regset.t) array) (out_regs : Regset.t)
    (typing : RTLtyping.regenv) :
    (reg, int * int) Hashtbl.t * (reg * bool) list array =
  let retl = Hashtbl.create 42 in
  let retr = Array.make (Array.length seqa) [] in
  (* retr.(i) : (r, b) -> (r', b') -> ...
   * where b = true if seen as arg, false if seen as dest
   *)
  List.iter
    (fun reg ->
      Hashtbl.add retl reg (Machregsaux.class_of_type (typing reg), 1))
    (Regset.elements out_regs);
  let add_reg reg =
    match Hashtbl.find_opt retl reg with
    | Some (t, n) -> Hashtbl.add retl reg (t, n + 1)
    | None -> Hashtbl.add retl reg (Machregsaux.class_of_type (typing reg), 1)
  in
  let map_true = List.map (fun r -> (r, true)) in
  Array.iteri
    (fun i (ins, _) ->
      match ins with
      | Bop (_, args, dest, _) | Bload (_, _, _, args, dest, _) ->
          List.iter add_reg args;
          retr.(i) <- (dest, false) :: map_true args
      | Bcond (_, args, _, _, _) ->
          List.iter add_reg args;
          retr.(i) <- map_true args
      | Bstore (_, _, args, src, _) ->
          List.iter add_reg args;
          add_reg src;
          retr.(i) <- (src, true) :: map_true args
      | BF (Bcall (_, fn, args, dest, _), _) ->
          List.iter add_reg args;
          retr.(i) <-
            (match fn with
            | Datatypes.Coq_inl reg ->
                add_reg reg;
                (dest, false) :: (reg, true) :: map_true args
            | _ -> (dest, false) :: map_true args)
      | BF (Btailcall (_, fn, args), _) ->
          List.iter add_reg args;
          retr.(i) <-
            (match fn with
            | Datatypes.Coq_inl reg ->
                add_reg reg;
                (reg, true) :: map_true args
            | _ -> map_true args)
      | BF (Bbuiltin (_, args, dest, _), _) ->
          let rec bar = function
            | AST.BA r ->
                add_reg r;
                retr.(i) <- (r, true) :: retr.(i)
            | AST.BA_splitlong (hi, lo) | AST.BA_addptr (hi, lo) ->
                bar hi;
                bar lo
            | _ -> ()
          in
          List.iter bar args;
          let rec bad = function
            | AST.BR r -> retr.(i) <- (r, false) :: retr.(i)
            | AST.BR_splitlong (hi, lo) ->
                bad hi;
                bad lo
            | _ -> ()
          in
          bad dest
      | BF (Bjumptable (reg, _), _) | BF (Breturn (Some reg), _) ->
          add_reg reg;
          retr.(i) <- [ (reg, true) ]
      | _ -> ())
    seqa;
  (* print_string "mentions\n";
   * Array.iteri (fun i l ->
   *     print_int i;
   *     print_string ": [";
   *     List.iter (fun (r, b) ->
   *         print_int (Camlcoq.P.to_int r);
   *         print_string ":";
   *         print_string (if b then "a:" else "d");
   *         if b then print_int (snd (Hashtbl.find retl r));
   *         print_string ", "
   *       ) l;
   *     print_string "]\n";
   *     flush stdout;
   *   ) retr; *)
  (retl, retr)

let get_live_regs_entry seqa ibf btl =
  if !Clflags.option_debug_compcert > 6 then debug_flag := true;
  let ret =
    Array.fold_right
      (fun (ins, liveins) regset_i ->
        let regset = Regset.union liveins regset_i in
        match ins with
        | Bnop _ -> regset
        | Bop (_, args, dest, _) | Bload (_, _, _, args, dest, _) ->
            List.fold_left
              (fun set reg -> Regset.add reg set)
              (Regset.remove dest regset)
              args
        | Bstore (_, _, args, src, _) ->
            List.fold_left
              (fun set reg -> Regset.add reg set)
              (Regset.add src regset) args
        | BF (Bcall (_, fn, args, dest, _), _) ->
            List.fold_left
              (fun set reg -> Regset.add reg set)
              ((match fn with
               | Datatypes.Coq_inl reg -> Regset.add reg
               | Datatypes.Coq_inr _ -> fun x -> x)
                 (Regset.remove dest regset))
              args
        | BF (Btailcall (_, fn, args), _) ->
            List.fold_left
              (fun set reg -> Regset.add reg set)
              (match fn with
              | Datatypes.Coq_inl reg -> Regset.add reg regset
              | Datatypes.Coq_inr _ -> regset)
              args
        | BF (Bbuiltin (_, args, dest, _), _) ->
            List.fold_left
              (fun set arg ->
                let rec add reg set =
                  match reg with
                  | AST.BA r -> Regset.add r set
                  | AST.BA_splitlong (hi, lo) | AST.BA_addptr (hi, lo) ->
                      add hi (add lo set)
                  | _ -> set
                in
                add arg set)
              (let rec rem dest set =
                 match dest with
                 | AST.BR r -> Regset.remove r set
                 | AST.BR_splitlong (hi, lo) -> rem hi (rem lo set)
                 | _ -> set
               in
               rem dest regset)
              args
        | BF (Bjumptable (reg, _), _) | BF (Breturn (Some reg), _) ->
            Regset.add reg regset
        | Bcond (_, args, _, _, _) ->
            List.fold_left (fun set reg -> Regset.add reg set) regset args
        | _ -> regset)
      seqa ibf.binfo.s_output_regs
  in
  debug "live in regs: ";
  print_regset ret;
  debug "\n";
  debug_flag := false;
  ret

let schedule_blk n ibf btl =
  if not !Clflags.option_fprepass then btl
  else
    let bseq, olast = flatten_blk_basics ibf in
    let seqa =
      Array.mapi
        (fun i inst -> (inst, get_liveins inst))
        bseq
    in
    let live_regs_entry = get_live_regs_entry seqa ibf btl in
    let refcount =
      reference_counting seqa ibf.binfo.s_output_regs ibf.binfo.typing
    in
    match
      schedule_sequence seqa btl live_regs_entry ibf.binfo.typing refcount
    with
    | Some positions ->
        let new_ib = apply_schedule bseq olast positions in
        let new_ibf =
          { entry = new_ib; binfo = ibf.binfo; input_regs = ibf.input_regs }
        in
        PTree.set n new_ibf btl
    | None -> btl

let turn_all_loads_nontrap n ibf btl =
  if not !config.has_non_trapping_loads || not !Clflags.option_fnontrap_loads then btl
  else
    let rec traverse_rec ib =
      match ib with
      | Bseq (ib1, ib2) -> Bseq (traverse_rec ib1, traverse_rec ib2)
      | Bload (t, a, b, c, d, e) ->
        let _opt_info =
          match t with
          | AST.TRAP -> Some true
          | AST.NOTRAP -> Some false in
        e.opt_info <- _opt_info;
        Bload (AST.NOTRAP, a, b, c, d, e)
      | _ -> ib
    in
    let ib' = traverse_rec ibf.entry in
    let ibf' =
      { entry = ib'; input_regs = ibf.input_regs; binfo = ibf.binfo }
    in
    PTree.set n ibf' btl

let compute_liveins n ibf btl =
  let rec traverse_rec ib =
    match ib with
    | Bseq (ib1, ib2) ->
        traverse_rec ib1;
        traverse_rec ib2
    | BF (Bgoto succ, iinfo)
    | BF (Bcall (_, _, _, _, succ), iinfo)
    | BF (Bbuiltin (_, _, _, succ), iinfo) ->
        let lives = (get_some @@ PTree.get succ btl).input_regs in
        iinfo.liveins <- lives
    | BF (Bjumptable (_, tbl), iinfo) ->
        List.iter
          (fun ex ->
            let lives = (get_some @@ PTree.get ex btl).input_regs in
            iinfo.liveins <- Regset.union iinfo.liveins lives)
          tbl
    | Bcond (_, _, BF (Bgoto succ, _), Bnop None, iinfo) -> (
        match iinfo.opt_info with
        | Some predb ->
            assert (predb = false);
            let lives = (get_some @@ PTree.get succ btl).input_regs in
            iinfo.liveins <- lives
        | None -> ())
    | _ -> ()
  in
  traverse_rec ibf.entry

let rec do_schedule btl = function
  | [] -> btl
  | (n, ibf) :: blks ->
      compute_liveins n ibf btl;
      let code_exp = expanse n ibf btl in
      let ibf_exp = get_some @@ PTree.get n code_exp in
      let code_nt = turn_all_loads_nontrap n ibf_exp code_exp in
      let ibf_nt = get_some @@ PTree.get n code_nt in
      let btl' = schedule_blk n ibf_nt code_nt in
      (*debug_flag := true;*)
      print_btl_code stderr btl;
      print_btl_code stderr btl';
      (*debug_flag := false;*)
      do_schedule btl' blks

let btl_scheduler f =
  let btl = f.fn_code in
  let elts = PTree.elements btl in
  find_last_reg elts;
  let btl' = do_schedule btl elts in
  btl'