open RTLpath
open RTL
open Maps
open RTLpathLivegenaux
open Registers
open Camlcoq

type superblock = {
  instructions: P.t array; (* pointers to code instructions *)
  (* each predicted Pcb has its attached liveins *)
  (* This is indexed by the pc value *)
  liveins: Regset.t PTree.t;
  (* Union of the input_regs of the last successors *)
  output_regs: Regset.t;
  typing: RTLtyping.regenv
}

let print_instructions insts code =
  if (!debug_flag) then begin
    dprintf "[ ";
    List.iter (
      fun n -> (PrintRTL.print_instruction stdout (P.to_int n, get_some @@ PTree.get n code))
    ) insts; dprintf "]"
  end

let print_superblock sb code =
  let insts = sb.instructions in
  let li = sb.liveins in
  let outs = sb.output_regs in
  begin
    dprintf "{ instructions = "; print_instructions (Array.to_list insts) code; dprintf "\n";
    dprintf "  liveins = "; print_ptree_regset li; dprintf "\n";
    dprintf "  output_regs = "; print_regset outs; dprintf "}"
  end

let print_superblocks lsb code =
  let rec f = function
    | [] -> ()
    | sb :: lsb -> (print_superblock sb code; dprintf ",\n"; f lsb)
  in begin
    dprintf "[\n";
    f lsb;
    dprintf "]"
  end

(* Adapted from backend/PrintRTL.ml: print_function *)
let print_code code = let open PrintRTL in let open Printf in
  if (!debug_flag) then begin
    fprintf stdout "{\n";
    let instrs =
      List.sort
        (fun (pc1, _) (pc2, _) -> compare pc2 pc1)
        (List.rev_map
          (fun (pc, i) -> (P.to_int pc, i))
          (PTree.elements code)) in
    List.iter (print_instruction stdout) instrs;
    fprintf stdout "}"
  end

let print_arrayp arr = begin
  dprintf "[| ";
  Array.iter (fun n -> dprintf "%d, " (P.to_int n)) arr;
  dprintf "|]"
end

let get_superblocks code entry pm typing =
  let visited = ref (PTree.map (fun n i -> false) code) in
  let rec get_superblocks_rec pc =
    let liveins = ref (PTree.empty) in
    let rec follow pc n =
      let inst = get_some @@ PTree.get pc code in
      if (n == 0) then begin
        (match (non_predicted_successors inst) with
          | [pcout] ->
              let live = (get_some @@ PTree.get pcout pm).input_regs in
                liveins := PTree.set pc live !liveins
          | _ -> ());
        ([pc], successors_inst inst)
      end else
        let nexts_from_exit = match (non_predicted_successors inst) with
          | [pcout] -> 
              let live = (get_some @@ PTree.get pcout pm).input_regs in begin
                liveins := PTree.set pc live !liveins;
                [pcout]
              end
          | [] -> []
          | _ -> failwith "Having more than one non_predicted_successor is not handled"
        in match (predicted_successor inst) with
          | None -> failwith "Incorrect path"
          | Some succ ->
              let (insts, nexts) = follow succ (n-1) in (pc :: insts, nexts_from_exit @ nexts)
    in if (get_some @@ PTree.get pc !visited) then []
    else begin
      visited := PTree.set pc true !visited;
      let pi = get_some @@ PTree.get pc pm in
      let (insts, nexts) = follow pc (Camlcoq.Nat.to_int pi.psize) in
      let superblock = { instructions = Array.of_list insts; liveins = !liveins;
        output_regs = pi.output_regs; typing = typing } in
      superblock :: (List.concat @@ List.map get_superblocks_rec nexts)
    end
  in let lsb = get_superblocks_rec entry in begin
    (* debug_flag := true; *)
    dprintf "Superblocks identified:"; print_superblocks lsb code; dprintf "\n";
    (* debug_flag := false; *)
    lsb
end

(* TODO David *)
let schedule_superblock sb code =
  if not !Clflags.option_fprepass
  then sb.instructions
  else
    let old_flag = !debug_flag in
    debug_flag := true;
    print_endline "ORIGINAL SUPERBLOCK";
    print_superblock sb code;
    debug_flag := old_flag;
    let nr_instr = Array.length sb.instructions in
    let trailer_length =
      match PTree.get (sb.instructions.(nr_instr-1)) code with
      | None -> 0
      | Some ii ->
         match predicted_successor ii with
         | Some _ -> 0
         | None -> 1 in
    match PrepassSchedulingOracle.schedule_sequence
            (Array.map (fun i ->
                 (match PTree.get i code with
                 | Some ii -> ii
                 | None -> failwith "RTLpathScheduleraux.schedule_superblock"),
                 (match PTree.get i sb.liveins with
                 | Some s -> s
                 | None -> Regset.empty))
               (Array.sub sb.instructions 0 (nr_instr-trailer_length))) with
    | None -> sb.instructions
    | Some order ->
       let ins' =
         Array.append 
           (Array.map (fun i -> sb.instructions.(i)) order)
           (Array.sub sb.instructions (nr_instr-trailer_length) trailer_length) in
       Printf.printf "REORDERED SUPERBLOCK %d\n" (Array.length ins');
       debug_flag := true;
       print_instructions (Array.to_list ins') code;
       debug_flag := old_flag;
       flush stdout;
       assert ((Array.length sb.instructions) = (Array.length ins'));
       (*sb.instructions; *)
       ins';;

  (* stub2: reverse function *)
  (*
  let reversed = Array.of_list @@ List.rev @@ Array.to_list (sb.instructions) in
  let tmp = reversed.(0) in
  let last_index = Array.length reversed - 1 in
  begin
    reversed.(0) <- reversed.(last_index);
    reversed.(last_index) <- tmp;
    reversed
  end *)
  (* stub: identity function *)

(**
 * Perform basic checks on the new order :
 * - must have the same length as the old order
 * - non basic instructions (call, tailcall, return, jumptable, non predicted CB) must not move
 *)
let check_order code old_order new_order = begin
  assert ((Array.length old_order) == (Array.length new_order));
  let length = Array.length new_order in
  if length > 0 then
    let last_inst = Array.get old_order (length - 1) in
    let instr = get_some @@ PTree.get last_inst code in
    match predicted_successor instr with
    | None ->
        if (last_inst != Array.get new_order (length - 1)) then
          failwith "The last instruction of the superblock is not basic, but was moved"
    | _ -> ()
end

type sinst =
  (* Each middle instruction has a direct successor *)
  (* A Smid can be the last instruction of a superblock, but a Send cannot be moved *)
  | Smid of RTL.instruction * node
  | Send of RTL.instruction

let rinst_to_sinst inst =
  match inst with
  | Inop n -> Smid(inst, n)
  | Iop (_,_,_,n) -> Smid(inst, n)
  | Iload (_,_,_,_,_,n) -> Smid(inst, n)
  | Istore (_,_,_,_,n) -> Smid(inst, n)
  | Icond (_,_,n1,n2,p) -> (
      match p with
      | Some true -> Smid(inst, n1)
      | Some false -> Smid(inst, n2)
      | None -> Send(inst)
    )
  | Icall _ | Ibuiltin _ | Ijumptable _ | Itailcall _ | Ireturn _ -> Send(inst)

let change_predicted_successor s = function
  | Smid(i, n) -> Smid(i, s)
  | Send _ -> failwith "Called change_predicted_successor on Send. Are you trying to move a non-basic instruction in the middle of the block?"

(* Forwards the successor changes into an RTL instruction *)
let sinst_to_rinst = function
  | Smid(inst, s) -> (
      match inst with
      | Inop n -> Inop s
      | Iop (a,b,c,n) -> Iop (a,b,c,s)
      | Iload (a,b,c,d,e,n) -> Iload (a,b,c,d,e,s)
      | Istore (a,b,c,d,n) -> Istore (a,b,c,d,s)
      | Icond (a,b,n1,n2,p) -> (
        match p with
        | Some true -> Icond(a, b, s, n2, p)
        | Some false -> Icond(a, b, n1, s, p)
        | None -> failwith "Non predicted Icond as a middle instruction!"
        )
      | _ -> failwith "That instruction shouldn't be a middle instruction"
      )
  | Send i -> i

let apply_schedule code sb new_order =
  let tc = ref code in
  let old_order = sb.instructions in
  begin
    check_order code old_order new_order;
    Array.iteri (fun i n' ->
      let inst' = get_some @@ PTree.get n' code in
      let iend = Array.length old_order - 1 in
      let new_inst =
        if (i == iend) then
          let final_inst_node = Array.get old_order iend in
          let sinst' = rinst_to_sinst inst' in
          match sinst' with
          (* The below assert fails if a Send is in the middle of the original superblock *)
          | Send i -> (assert (final_inst_node == n'); i)
          | Smid _ ->
              let final_inst = get_some @@ PTree.get final_inst_node code in
              match rinst_to_sinst final_inst with
              | Smid (_, s') -> sinst_to_rinst @@ change_predicted_successor s' sinst'
              | Send _ -> assert(false) (* should have failed earlier *)
        else
          sinst_to_rinst
          (* this will fail if the moved instruction is a Send *)
          @@ change_predicted_successor (Array.get old_order (i+1))
          @@ rinst_to_sinst inst'
      in tc := PTree.set (Array.get old_order i) new_inst !tc
    ) new_order;
    !tc
  end

  (*
let main_successors = function
| Inop n | Iop (_,_,_,n) | Iload (_,_,_,_,_,n) | Istore (_,_,_,_,n) -> [n]
| Icall (_,_,_,_,n) | Ibuiltin (_,_,_,n) -> [n]
| Icond (_,_,n1,n2,p) -> (
    match p with
    | Some true -> [n1; n2]
    | Some false ->  [n2; n1]
    | None -> [n1; n2] )
| Ijumptable _ | Itailcall _ | Ireturn _ -> []

let change_predicted_successor i s = match i with
  | Itailcall _ | Ireturn _ -> failwith "Wrong instruction (5) - Tailcalls and returns should not be moved in the middle of a superblock"
  | Ijumptable _ -> failwith "Wrong instruction (6) - Jumptables should not be moved in the middle of a superblock"
  | Inop n -> Inop s
  | Iop (a,b,c,n) -> Iop (a,b,c,s)
  | Iload (a,b,c,d,e,n) -> Iload (a,b,c,d,e,s)
  | Istore (a,b,c,d,n) -> Istore (a,b,c,d,s)
  | Icall (a,b,c,d,n) -> Icall (a,b,c,d,s)
  | Ibuiltin (a,b,c,n) -> Ibuiltin (a,b,c,s)
  | Icond (a,b,n1,n2,p) -> (
      match p with
      | Some true -> Icond (a,b,s,n2,p)
      | Some false -> Icond (a,b,n1,s,p)
      | None -> failwith "Predicted a successor for an Icond with p=None - unpredicted CB should not be moved in the middle of the superblock"
      )

let rec change_successors i = function
  | [] -> (
      match i with
      | Itailcall _ | Ireturn _ -> i
      | _ -> failwith "Wrong instruction (1)")
  | [s] -> (
      match i with
      | Inop n -> Inop s
      | Iop (a,b,c,n) -> Iop (a,b,c,s)
      | Iload (a,b,c,d,e,n) -> Iload (a,b,c,d,e,s)
      | Istore (a,b,c,d,n) -> Istore (a,b,c,d,s)
      | Icall (a,b,c,d,n) -> Icall (a,b,c,d,s)
      | Ibuiltin (a,b,c,n) -> Ibuiltin (a,b,c,s)
      | Ijumptable (a,[n]) -> Ijumptable (a,[s])
      | Icond (a,b,n1,n2,p) -> (
          match p with
          | Some true -> Icond (a,b,s,n2,p)
          | Some false -> Icond (a,b,n1,s,p)
          | None -> failwith "Icond Wrong instruction (2) ; should not happen?"
        )
      | _ -> failwith "Wrong instruction (2)")
  | [s1; s2] -> (
      match i with
      | Icond (a,b,n1,n2,p) -> Icond (a,b,s1,s2,p)
      | Ijumptable (a, [n1; n2]) -> Ijumptable (a, [s1; s2])
      | _ -> change_successors i [s1])
  | ls -> (
      match i with
      | Ijumptable (a, ln) -> begin
          assert ((List.length ln) == (List.length ls));
          Ijumptable (a, ls)
        end
      | _ -> failwith "Wrong instruction (4)")


let apply_schedule code sb new_order =
  let tc = ref code in
  let old_order = sb.instructions in
  let last_node = Array.get old_order (Array.length old_order - 1) in
  let last_successors = main_successors
    @@ get_some @@ PTree.get last_node code in
  begin
    check_order code old_order new_order;
    Array.iteri (fun i n' ->
      let inst' = get_some @@ PTree.get n' code in
      let new_inst =
        if (i == (Array.length old_order - 1)) then
          change_successors inst' last_successors
        else
          change_predicted_successor inst' (Array.get old_order (i+1))
      in tc := PTree.set (Array.get old_order i) new_inst !tc
    ) new_order;
    !tc
  end
*)

let rec do_schedule code = function
  | [] -> code
  | sb :: lsb ->
      let schedule = schedule_superblock sb code in
      let new_code = apply_schedule code sb schedule in
      begin
        (* debug_flag := true; *)
        dprintf "Old Code: "; print_code code;
        dprintf "\nSchedule to apply: "; print_arrayp schedule;
        dprintf "\nNew Code: "; print_code new_code;
        dprintf "\n";
        (* debug_flag := false; *)
        do_schedule new_code lsb
      end

let get_ok r = match r with Errors.OK x -> x | _ -> failwith "Did not get OK"

let scheduler f =
  let code = f.fn_RTL.fn_code in
  let id_ptree = PTree.map (fun n i -> n) (f.fn_path) in
  let entry = f.fn_RTL.fn_entrypoint in
  let pm = f.fn_path in
  let typing = get_ok @@ RTLtyping.type_function f.fn_RTL in
  let lsb = get_superblocks code entry pm typing in
  begin
    (* debug_flag := true; *)
    dprintf "Pathmap:\n"; dprintf "\n";
    print_path_map pm;
    dprintf "Superblocks:\n";
    print_superblocks lsb code; dprintf "\n";
    (* debug_flag := false; *)
    let tc = do_schedule code lsb in
    (((tc, entry), pm), id_ptree)
  end