open Maps
open BTL
open RTLcommonaux
open BTLcommonaux
open BTLtypes
open DebugPrint
open PrintBTL

let recompute_inumbs btl entry =
  let btl = reset_visited_ib (reset_visited_ibf btl) in
  let last_used = ref 0 in
  let ibf = get_some @@ PTree.get entry btl in
  let ipos () =
    last_used := !last_used + 1;
    !last_used
  in
  let rec walk ib k =
    (* heuristic: try to explore the lower numbered branch first *)
    let walk_smallest_child s1 s2 ib1 ib2 =
      if s1 < s2 then (
        walk ib1 None;
        walk ib2 None)
      else (
        walk ib2 None;
        walk ib1 None)
    in
    if jump_visit ib then ()
    else
      match ib with
      | BF (Bcall (_, _, _, _, s), iinfo) | BF (Bbuiltin (_, _, _, s), iinfo) ->
          let ib' = (get_some @@ PTree.get s btl).entry in
          walk ib' None;
          iinfo.inumb <- ipos ()
      | BF (Bgoto s, _) ->
          let ib' = (get_some @@ PTree.get s btl).entry in
          walk ib' None
      | BF (Bjumptable (_, tbl), iinfo) ->
          List.iter
            (fun s ->
              let ib' = (get_some @@ PTree.get s btl).entry in
              walk ib' None)
            tbl;
          iinfo.inumb <- ipos ()
      | BF (Btailcall (_, _, _), iinfo) | BF (Breturn _, iinfo) ->
          iinfo.inumb <- ipos ()
      | Bnop None ->
          failwith
            "recompute_inumbs: Bnop None can only be in the right child of \
             Bcond"
      | Bnop (Some iinfo)
      | Bop (_, _, _, iinfo)
      | Bload (_, _, _, _, _, iinfo)
      | Bstore (_, _, _, _, iinfo) ->
          let succ = get_some @@ k in
          walk succ None;
          iinfo.inumb <- ipos ()
      | Bseq (ib1, ib2) -> walk ib1 (Some ib2)
      | Bcond (_, _, BF (Bgoto s1, iinfoL), Bnop None, iinfoF) ->
          iinfoL.visited <- true;
          let ib1 = get_some @@ PTree.get s1 btl in
          let ib2 = get_some @@ k in
          walk_smallest_child (p2i s1) (get_inumb_or_next ib2) ib1.entry ib2;
          iinfoF.inumb <- ipos ()
      | Bcond (_, _, _, _, _) -> failwith "recompute_inumbs: unsupported Bcond"
  in
  walk ibf.entry None;
  btl

let regenerate_btl_tree btl entry =
  let new_entry = ref entry in
  let rec renumber_iblock ib =
    let get_new_succ s =
      let sentry = get_some @@ PTree.get s btl in
      i2p (get_inumb_or_next sentry.entry)
    in
    match ib with
    | BF (Bcall (sign, fn, lr, rd, s), iinfo) ->
        BF (Bcall (sign, fn, lr, rd, get_new_succ s), iinfo)
    | BF (Bbuiltin (sign, fn, lr, s), iinfo) ->
        BF (Bbuiltin (sign, fn, lr, get_new_succ s), iinfo)
    | BF (Bgoto s, iinfo) -> BF (Bgoto (get_new_succ s), iinfo)
    | BF (Bjumptable (arg, tbl), iinfo) ->
        let tbl' = List.map (fun s -> get_new_succ s) tbl in
        BF (Bjumptable (arg, tbl'), iinfo)
    | Bcond (cond, lr, ib1, ib2, iinfo) ->
        Bcond (cond, lr, renumber_iblock ib1, renumber_iblock ib2, iinfo)
    | Bseq (ib1, ib2) -> Bseq (renumber_iblock ib1, renumber_iblock ib2)
    | _ -> ib
  in
  let dm = ref PTree.empty in
  let ord_btl =
    PTree.fold
      (fun ord_btl old_n old_ibf ->
        let ib = renumber_iblock old_ibf.entry in
        let n = get_inumb_or_next ib in
        let n_pos = i2p n in
        let bi = mk_binfo n old_ibf.binfo.s_output_regs old_ibf.binfo.typing in
        let ibf = { entry = ib; input_regs = old_ibf.input_regs; binfo = bi } in
        if old_n = entry then new_entry := n_pos;
        dm := PTree.set old_n n_pos !dm;
        PTree.set n_pos ibf ord_btl)
      btl PTree.empty
  in
  debug "Renumbered BTL with new_entry=%d:\n" (p2i !new_entry);
  print_btl_code stderr ord_btl;
  ((ord_btl, !new_entry), !dm)

let renumber btl entry =
  (*debug_flag := true;*)
  let btl' = recompute_inumbs btl entry in
  (*debug_flag := false;*)
  regenerate_btl_tree btl' entry