open Maps
open BTL
open BTLtypes
open Machine
open DebugPrint
open PrintBTL
open RTLcommonaux
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.pcond 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
             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

let schedule_blk n ibf btl =
  if not !Clflags.option_fprepass then btl
  else
    let bseq, olast = flatten_blk_basics ibf in
    match schedule_sequence bseq btl 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 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) -> 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 rec do_schedule btl = function
  | [] -> btl
  | (n, ibf) :: blks ->
      let code_exp = expanse n ibf btl in
      let code_nt = turn_all_loads_nontrap n ibf code_exp in
      let btl' = schedule_blk n ibf code_nt in
      do_schedule btl' blks

let btl_scheduler f =
  let btl = f.fn_code in
  (*debug_flag := true;*)
  let elts = PTree.elements btl in
  find_last_reg elts;
  let btl' = do_schedule btl elts in
  debug "Scheduled BTL Code:\n";
  print_btl_code stderr btl';
  (*debug_flag := false;*)
  btl'