open Maps
open BTL
open BTLtypes
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 apply_schedule bseq olast positions =
  let ibl = Array.to_list (Array.map (fun i -> bseq.(i)) positions) 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 ->
        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 rec do_schedule btl = function
  | [] -> btl
  | (n, ibf) :: blks ->
      let code_exp = expanse n ibf btl in
      let btl' = schedule_blk n ibf code_exp 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'