open DebugPrint
open Machine
open RTLpathLivegenaux
open RTLpath
open RTLpathCommon
open RTL
open Maps
open Registers
open ExpansionOracle

let config = Machine.config

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

let print_superblocks lsb code =
  let rec f = function
    | [] -> ()
    | sb :: lsb -> (print_superblock sb code; debug ",\n"; f lsb)
  in begin
    debug "[\n";
    f lsb;
    debug "]"
  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;
        s_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; *)
    debug "Superblocks identified:"; print_superblocks lsb code; debug "\n";
    (* debug_flag := false; *)
    lsb
end

(** the useful one. Returns a hashtable with bindings of shape
 ** [(r,(t, n)], where [r] is a pseudo-register (Registers.reg),
 ** [t] is its class (according to [typing]), and [n] the number of
 ** times it's referenced as an argument in instructions of [seqa] ;
 ** and an arrray containg the list of regs referenced by each
 ** instruction, with a boolean to know whether it's as arg or dest *)
let reference_counting (seqa : (instruction * Regset.t) array)
      (out_regs : Registers.Regset.t) (typing : RTLtyping.regenv) :
      (Registers.reg, int * int) Hashtbl.t *
        (Registers.reg * bool) list  array =
  let retl = Hashtbl.create 42 in
  let retr = Array.make (Array.length seqa) [] in
  (* retr.(i) : (r, b) -> (r', b') -> ...
   * where b = true if seen as arg, false if seen as dest
   *)
  List.iter (fun reg ->
      Hashtbl.add retl
        reg (Machregsaux.class_of_type (typing reg), 1)
    ) (Registers.Regset.elements out_regs);
  let add_reg reg = 
    match Hashtbl.find_opt retl reg with
    | Some (t, n) -> Hashtbl.add retl reg (t, n+1)
    | None -> Hashtbl.add retl reg (Machregsaux.class_of_type
                                    (typing reg), 1)
  in
  let map_true = List.map (fun r -> r, true) in
  Array.iteri (fun i  (ins, _) ->
      match ins with
      | Iop(_,args,dest,_) | Iload(_,_,_,args,dest,_) ->
         List.iter (add_reg) args;
         retr.(i) <- (dest, false)::(map_true args)
      | Icond(_,args,_,_,_) ->
         List.iter (add_reg) args;
         retr.(i) <- map_true args
      | Istore(_,_,args,src,_) ->
         List.iter (add_reg) args;
         add_reg src;
         retr.(i) <- (src, true)::(map_true args)
      | Icall(_,fn,args,dest,_) ->
         List.iter (add_reg) args;
         retr.(i) <- (match fn with
                     | Datatypes.Coq_inl reg ->
                        add_reg reg;
                        (dest,false)::(reg, true)::(map_true args)
                     | _ -> (dest,false)::(map_true args))

      | Itailcall(_,fn,args) ->
         List.iter (add_reg) args;
         retr.(i) <- (match fn with
                     | Datatypes.Coq_inl reg ->
                        add_reg reg;
                        (reg, true)::(map_true args)
                     | _ -> map_true args)
      | Ibuiltin(_,args,dest,_) ->
         let rec bar = function
           | AST.BA r -> add_reg r;
                        retr.(i) <- (r, true)::retr.(i)
           | AST.BA_splitlong  (hi, lo) | AST.BA_addptr (hi, lo) ->
              bar hi; bar lo
           | _ -> ()
         in
         List.iter (bar) args;
         let rec bad = function
           | AST.BR r -> retr.(i) <- (r, false)::retr.(i)
           | AST.BR_splitlong (hi, lo) ->
              bad hi; bad lo
           | _ -> ()
         in
         bad dest;
      | Ijumptable (reg,_) | Ireturn (Some reg) ->
         add_reg reg;
         retr.(i) <- [reg, true]
      | _ -> ()
    ) seqa;
  (* print_string "mentions\n";
   * Array.iteri (fun i l ->
   *     print_int i;
   *     print_string ": [";
   *     List.iter (fun (r, b) ->
   *         print_int (Camlcoq.P.to_int r);
   *         print_string ":";
   *         print_string (if b then "a:" else "d");
   *         if b then print_int (snd (Hashtbl.find retl r));
   *         print_string ", "
   *       ) l;
   *     print_string "]\n";
   *     flush stdout;
   *   ) retr; *)
  retl, retr

  
let get_live_regs_entry (sb : superblock) code =
  (if !Clflags.option_debug_compcert > 6
   then debug_flag := true);
  debug "getting live regs for superblock:\n";
  print_superblock sb code;
  debug "\n";
  let seqa = Array.map (fun i ->
                 (match PTree.get i code with
                  | Some ii -> ii
                  | None -> failwith "RTLpathScheduleraux.get_live_regs_entry"
                 ),
                 (match PTree.get i sb.liveins with
                  | Some s -> s
                  | None -> Regset.empty))
               sb.instructions in
  let ret = 
    Array.fold_right (fun (ins, liveins) regset_i ->
        let regset = Registers.Regset.union liveins regset_i in
        match ins with
        | Inop _ -> regset
        | Iop (_, args, dest, _)
        | Iload (_, _, _, args, dest, _) ->
           List.fold_left (fun set reg -> Registers.Regset.add reg set)
             (Registers.Regset.remove dest regset) args
        | Istore (_, _, args, src, _) ->
           List.fold_left (fun set reg -> Registers.Regset.add reg set)
             (Registers.Regset.add src regset) args
        | Icall (_, fn, args, dest, _) ->
           List.fold_left (fun set reg -> Registers.Regset.add reg set)
             ((match fn with
               | Datatypes.Coq_inl reg -> (Registers.Regset.add reg)
               | Datatypes.Coq_inr _ -> (fun x -> x))
                (Registers.Regset.remove dest regset))
             args
        | Itailcall (_, fn, args) ->
           List.fold_left (fun set reg -> Registers.Regset.add reg set)
             (match fn with
              | Datatypes.Coq_inl reg -> (Registers.Regset.add reg regset)
              | Datatypes.Coq_inr _ -> regset)
             args
        | Ibuiltin (_, args, dest, _) ->
           List.fold_left (fun set arg ->
               let rec add reg set =
                 match reg with
                 | AST.BA r -> Registers.Regset.add r set
                 | AST.BA_splitlong (hi, lo)
                 | AST.BA_addptr (hi, lo) -> add hi (add lo set)
                 | _ -> set
             in add arg set)
             (let rec rem dest set =
                match dest with
                | AST.BR r -> Registers.Regset.remove r set
                | AST.BR_splitlong (hi, lo) -> rem hi (rem lo set)
                | _ -> set
             in rem dest regset)
             args
        | Icond (_, args, _, _, _) ->
           List.fold_left (fun set reg ->
               Registers.Regset.add reg set)
             regset args
        | Ijumptable (reg, _)
          | Ireturn (Some reg) ->
           Registers.Regset.add reg regset
        | _ -> regset
      ) seqa sb.s_output_regs
  in debug "live in regs: ";
     print_regset ret;
     debug "\n";
     debug_flag := false;
     ret
  
(* 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
    debug "hello\n";
    let live_regs_entry = get_live_regs_entry sb code in
    let seqa = 
      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)) in
    match PrepassSchedulingOracle.schedule_sequence
            seqa
            live_regs_entry
            sb.typing
            (reference_counting seqa sb.s_output_regs sb.typing) 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 is_a_cb = function Icond _ -> true | _ -> false
let is_a_load = function Iload _ -> true | _ -> false

let find_array arr n =
  let index = ref None in
  begin
    Array.iteri (fun i n' ->
      if n = n' then
        match !index with
        | Some _ -> failwith "More than one element present"
        | None -> index := Some i
    ) arr;
    !index
  end

let rec hashedset_from_list = function
  | [] -> HashedSet.PSet.empty
  | n::ln -> HashedSet.PSet.add n (hashedset_from_list ln)

let hashedset_map f hs = hashedset_from_list @@ List.map f @@ HashedSet.PSet.elements hs

let apply_schedule code sb new_order =
  let tc = ref code in
  let old_order = sb.instructions in
  let count_cbs order code =
    let current_cbs = ref HashedSet.PSet.empty in
    let cbs_above = ref PTree.empty in
    Array.iter (fun n ->
      let inst = get_some @@ PTree.get n code in
      if is_a_cb inst then current_cbs := HashedSet.PSet.add n !current_cbs
      else if is_a_load inst then cbs_above := PTree.set n !current_cbs !cbs_above
    ) order;
    !cbs_above
  in let fmap n =
    let index = get_some @@ find_array new_order n in
    old_order.(index)
  in begin
    check_order code old_order new_order;
    (* First pass - modify the positions, nothing else *)
    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;
    (* Second pass - turn the loads back into trapping when it was not needed *)
    (* 1) We remember which CBs are "above" a given load *)
    let cbs_above = count_cbs old_order code in
    (* 2) We do the same for new_order *)
    let cbs_above' = count_cbs (Array.map fmap new_order) !tc in
    (* 3) We examine each load, turn it back into trapping if cbs_above is included in cbs_above' *)
    Array.iter (fun n ->
      let n' = fmap n in
      let inst' = get_some @@ PTree.get n' !tc in
      match inst' with
      | Iload (t,a,b,c,d,s) ->
          let pset = hashedset_map fmap @@ get_some @@ PTree.get n cbs_above in
          let pset' = get_some @@ PTree.get n' cbs_above' in
          if HashedSet.PSet.is_subset pset pset' then tc := PTree.set n' (Iload (AST.TRAP,a,b,c,d,s)) !tc
          else assert !config.has_non_trapping_loads
      | _ -> ()
    ) old_order;
    !tc
  end

let turn_all_loads_nontrap sb code =
  if not !config.has_non_trapping_loads then code
  else begin
    let code' = ref code in
    Array.iter (fun n ->
      let inst = get_some @@ PTree.get n code in
      match inst with
      | Iload (t,a,b,c,d,s) -> code' := PTree.set n (Iload (AST.NOTRAP,a,b,c,d,s)) !code'
      | _ -> ()
    ) sb.instructions;
    !code'
  end

let rec do_schedule code pm = function
  | [] -> (code, pm)
  | sb :: lsb ->
      (*debug_flag := true;*)
      let (code_exp, pm) = expanse sb code pm in
      (*debug_flag := false;*)
      (* Trick: instead of turning loads into non trap as needed..
       * First, we turn them all into non-trap.
       * Then, we turn back those who didn't need to be turned, into TRAP again
       * This is because the scheduler (rightfully) refuses to schedule ahead of a branch
       * operations that might trap *)
      let code' = turn_all_loads_nontrap sb code_exp in
      let schedule = schedule_superblock sb code' in
      let new_code = apply_schedule code' sb schedule in
      begin
        (*debug_flag := true;*)
        if code != code_exp then (
        debug "Old Code: "; print_code code;
        debug "Exp Code: "; print_code code_exp);
        debug "\nSchedule to apply: "; print_arrayp schedule;
        debug "\nNew Code: "; print_code new_code;
        debug "\n";
        do_schedule new_code pm 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; *)
    debug "Pathmap:\n"; debug "\n";
    print_path_map pm;
    debug "Superblocks:\n";
    (*print_code code; flush stdout; flush stderr;*)
    (*debug_flag := false;*)
    (*print_superblocks lsb code; debug "\n";*)
    find_last_node_reg (PTree.elements code);
    let (tc, pm) = do_schedule code pm lsb in
    (((tc, entry), pm), id_ptree)
  end