open AST
open Maps
open Registers
open BTL
open BTLtypes
open DebugPrint
open PrintBTL
open RTLcommonaux
open InstructionScheduler
open PrepassSchedulingOracleDeps

let use_alias_analysis () = false

let build_constraints_and_resources (opweights : opweights) insts btl =
  let last_reg_reads : int list PTree.t ref = ref PTree.empty
  and last_reg_write : (int * int) PTree.t ref = ref PTree.empty
  and last_mem_reads : int list ref = ref []
  and last_mem_write : int option ref = ref None
  and last_branch : int option ref = ref None
  and last_non_pipelined_op : int array =
    Array.make opweights.nr_non_pipelined_units (-1)
  and latency_constraints : latency_constraint list ref = ref []
  and resources = ref [] in
  let add_constraint instr_from instr_to latency =
    assert (instr_from <= instr_to);
    assert (latency >= 0);
    if instr_from = instr_to then
      if latency = 0 then ()
      else
        failwith "PrepassSchedulingOracle.get_dependencies: negative self-loop"
    else
      latency_constraints :=
        { instr_from; instr_to; latency } :: !latency_constraints
  and get_last_reads reg =
    match PTree.get reg !last_reg_reads with Some l -> l | None -> []
  in
  let add_input_mem i =
    if not (use_alias_analysis ()) then (
      (* Read after write *)
      (match !last_mem_write with None -> () | Some j -> add_constraint j i 1);
      last_mem_reads := i :: !last_mem_reads)
  and add_output_mem i =
    if not (use_alias_analysis ()) then (
      (* Write after write *)
      (match !last_mem_write with None -> () | Some j -> add_constraint j i 1);
      (* Write after read *)
      List.iter (fun j -> add_constraint j i 0) !last_mem_reads;
      last_mem_write := Some i;
      last_mem_reads := [])
  and add_input_reg i reg =
    (* Read after write *)
    (match PTree.get reg !last_reg_write with
    | None -> ()
    | Some (j, latency) -> add_constraint j i latency);
    last_reg_reads := PTree.set reg (i :: get_last_reads reg) !last_reg_reads
  and add_output_reg i latency reg =
    (* Write after write *)
    (match PTree.get reg !last_reg_write with
    | None -> ()
    | Some (j, _) -> add_constraint j i 1);
    (* Write after read *)
    List.iter (fun j -> add_constraint j i 0) (get_last_reads reg);
    last_reg_write := PTree.set reg (i, latency) !last_reg_write;
    last_reg_reads := PTree.remove reg !last_reg_reads
  in
  let add_input_regs i regs = List.iter (add_input_reg i) regs
  and irreversible_action i =
    match !last_branch with None -> () | Some j -> add_constraint j i 1
  in
  let set_branch i =
    irreversible_action i;
    last_branch := Some i
  and add_non_pipelined_resources i resources =
    Array.iter2
      (fun latency last ->
        if latency >= 0 && last >= 0 then add_constraint last i latency)
      resources last_non_pipelined_op;
    Array.iteri
      (fun rsc latency -> if latency >= 0 then last_non_pipelined_op.(rsc) <- i)
      resources
  in
  Array.iteri
    (fun i inst ->
      match inst with
      | Bnop _ ->
          let rs = Array.map (fun _ -> 0) opweights.pipelined_resource_bounds in
          resources := rs :: !resources
      | Bop (op, lr, rd, _) ->
          add_non_pipelined_resources i
            (opweights.non_pipelined_resources_of_op op (List.length lr));
          if Op.is_trapping_op op then irreversible_action i;
          add_input_regs i lr;
          add_output_reg i (opweights.latency_of_op op (List.length lr)) rd;
          let rs = opweights.resources_of_op op (List.length lr) in
          resources := rs :: !resources
      | Bload (trap, chk, addr, lr, rd, _) ->
          if trap = TRAP then irreversible_action i;
          add_input_mem i;
          add_input_regs i lr;
          add_output_reg i
            (opweights.latency_of_load trap chk addr (List.length lr))
            rd;
          let rs = opweights.resources_of_load trap chk addr (List.length lr) in
          resources := rs :: !resources
      | Bstore (chk, addr, lr, src, _) ->
          irreversible_action i;
          add_input_regs i lr;
          add_input_reg i src;
          add_output_mem i;
          let rs = opweights.resources_of_store chk addr (List.length lr) in
          resources := rs :: !resources
      | Bcond (cond, lr, BF (Bgoto s, _), ibnot, _) ->
          (* TODO gourdinl test with/out this line *)
          let live = (get_some @@ PTree.get s btl).input_regs in
          add_input_regs i (Regset.elements live);
          set_branch i;
          add_input_mem i;
          add_input_regs i lr;
          let rs = opweights.resources_of_cond cond (List.length lr) in
          resources := rs :: !resources
      | Bcond (_, _, _, _, _) ->
          failwith "get_simple_dependencies: invalid Bcond"
      | BF (_, _) -> failwith "get_simple_dependencies: BF"
      | Bseq (_, _) -> failwith "get_simple_dependencies: Bseq")
    insts;
  (!latency_constraints, Array.of_list (List.rev !resources))

let define_problem (opweights : opweights) ibf btl =
  let simple_deps, resources =
    build_constraints_and_resources opweights ibf btl
  in
  {
    max_latency = -1;
    resource_bounds = opweights.pipelined_resource_bounds;
    instruction_usages = resources;
    latency_constraints = simple_deps;
  }

let zigzag_scheduler problem early_ones =
  let nr_instructions = get_nr_instructions problem in
  assert (nr_instructions = Array.length early_ones);
  match list_scheduler problem with
  | Some fwd_schedule ->
      let fwd_makespan = fwd_schedule.(Array.length fwd_schedule - 1) in
      let constraints' = ref problem.latency_constraints in
      Array.iteri
        (fun i is_early ->
          if is_early then
            constraints' :=
              {
                instr_from = i;
                instr_to = nr_instructions;
                latency = fwd_makespan - fwd_schedule.(i);
              }
              :: !constraints')
        early_ones;
      validated_scheduler reverse_list_scheduler
        { problem with latency_constraints = !constraints' }
  | None -> None

let prepass_scheduler_by_name name problem insts =
  match name with
  | "zigzag" ->
      let early_ones =
        Array.map
          (fun inst ->
            match inst with Bcond (_, _, _, _, _) -> true | _ -> false)
          insts
      in
      zigzag_scheduler problem early_ones
  | _ -> scheduler_by_name name problem

let schedule_sequence insts btl =
  let opweights = OpWeights.get_opweights () in
  try
    if Array.length insts <= 1 then None
    else
      let nr_instructions = Array.length insts in
      let problem = define_problem opweights insts btl in
      match
        prepass_scheduler_by_name !Clflags.option_fprepass_sched problem insts
      with
      | None ->
          Printf.printf "no solution in prepass scheduling\n";
          None
      | Some solution ->
          let positions = Array.init nr_instructions (fun i -> i) in
          Array.sort
            (fun i j ->
              let si = solution.(i) and sj = solution.(j) in
              if si < sj then -1 else if si > sj then 1 else i - j)
            positions;
          Some positions
  with Failure s ->
    Printf.printf "failure in prepass scheduling: %s\n" s;
    None

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 =
  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 btl' = schedule_blk n ibf btl in
      do_schedule btl' blks

let btl_scheduler f =
  let btl = f.fn_code in
  (*debug_flag := true;*)
  let btl' = do_schedule btl (PTree.elements btl) in
  debug "Scheduled BTL Code:\n";
  print_btl_code stderr btl';
  (*debug_flag := false;*)
  btl'