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 =
  debug "count_cbs\n";
  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 (
      debug "n is %d, add cb at: %d\n" n indexes.(n);
      current_cbs := SI.add indexes.(n) !current_cbs)
    else if is_a_load ib then (
      debug "n is %d, add load at: %d\n" n indexes.(n);
      Hashtbl.add cbs_above indexes.(n) !current_cbs)
  in
  Array.iteri (fun n ib -> update_cbs n ib) bseq;
  (match olast with
  | Some last ->
      debug "last\n";
      update_cbs (Array.length bseq) last
  | None -> ());
  cbs_above

let apply_schedule bseq olast positions =
  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
  debug_flag := true;
  Array.iter (fun i -> debug "%d " i) positions;
  debug "\n";
  Array.iter (fun i -> debug "%d " i) fseq;
  debug "\n";
  Array.iter
    (fun i -> debug "%d " i)
    (Array.init (Array.length positions) (fun i -> i));
  debug "\n";
  let cbs_above_old = count_cbs bseq olast fseq in
  let bseq_new = Array.map (fun i -> bseq.(i)) positions 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
  debug_flag := false;
  build_iblock ibl

let schedule_blk n ibf btl =
  let bseq, olast = flatten_blk_basics ibf in
  match schedule_sequence bseq btl with
  | Some positions ->
      debug "%d," (p2i n);
      Array.iter (fun p -> debug "%d " p) positions;
      debug "\n";
      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 btl in
      let btl' = schedule_blk n ibf code_nt in
      (*debug_flag := true;*)
      if btl != code_exp then (
        debug "#######################################################\n";
        print_btl_code stderr btl;
        debug "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n";
        print_btl_code stderr code_exp);
      (*debug_flag := false;*)
      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'