(* *************************************************************)
(*                                                             *)
(*             The Compcert verified compiler                  *)
(*                                                             *)
(*           David Monniaux     CNRS, VERIMAG                  *)
(*                                                             *)
(*  Copyright VERIMAG. All rights reserved.                    *)
(*  This file is distributed under the terms of the INRIA      *)
(*  Non-Commercial License Agreement.                          *)
(*                                                             *)
(* *************************************************************)

open CSE3analysis
open Maps
open HashedSet
open Camlcoq
open Coqlib
   
let flatten_eq eq =
  ((P.to_int eq.eq_lhs), eq.eq_op, List.map P.to_int eq.eq_args);;

let imp_add_i_j s i j =
  s := PMap.set i (PSet.add j (PMap.get i !s)) !s;;

let string_of_chunk = function
    | AST.Mint8signed -> "int8signed"
    | AST.Mint8unsigned -> "int8unsigned"
    | AST.Mint16signed -> "int16signed"
    | AST.Mint16unsigned -> "int16unsigned"
    | AST.Mint32 -> "int32"
    | AST.Mint64 -> "int64"
    | AST.Mfloat32 -> "float32"
    | AST.Mfloat64 -> "float64"
    | AST.Many32 -> "any32"
    | AST.Many64 -> "any64";;

let print_reg channel i =
  Printf.fprintf channel "r%d" i;;

let print_eq channel (lhs, sop, args) =
  match sop with
  | SOp op ->
     Printf.printf "%a = %a" print_reg lhs (PrintOp.print_operation print_reg) (op, args)
  | SLoad(chunk, addr) ->
     Printf.printf "%a = %s @ %a" print_reg lhs (string_of_chunk chunk)
       (PrintOp.print_addressing print_reg) (addr, args);;

let pp_intset oc s =
  Printf.fprintf oc "{ ";
  List.iter (fun i -> Printf.fprintf oc "%d; " (P.to_int i)) (PSet.elements s);
  Printf.fprintf oc "}";;

let pp_rhs oc (sop, args) =
  match sop with
  | SOp op -> PrintOp.print_operation PrintRTL.reg oc (op, args)
  | SLoad(chunk, addr) ->
     Printf.fprintf oc "%s[%a]"
       (PrintAST.name_of_chunk chunk)
         (PrintOp.print_addressing PrintRTL.reg) (addr, args);;

let pp_eq oc eq =
  Printf.fprintf oc "x%d = %a" (P.to_int eq.eq_lhs)
    pp_rhs (eq.eq_op, eq.eq_args);;

let pp_P oc x = Printf.fprintf oc "%d" (P.to_int x)
              
let pp_option pp oc = function
  | None -> output_string oc "none"
  | Some x -> pp oc x;;

let is_trivial eq =
  (eq.eq_op = SOp Op.Omove) && (eq.eq_args = [eq.eq_lhs]);;

let rec pp_list separator pp_item chan = function
  | [] -> ()
  | [h] -> pp_item chan h
  | h::t ->
     pp_item chan h;
     output_string chan separator;
     pp_list separator pp_item chan t;;

let pp_set separator pp_item chan s =
  pp_list separator pp_item chan (PSet.elements s);;

let pp_equation hints chan x =
  match PTree.get x hints.hint_eq_catalog with
  | None -> output_string chan "???"
  | Some eq ->
     print_eq chan  (P.to_int eq.eq_lhs, eq.eq_op, List.map P.to_int eq.eq_args);;

let pp_relation hints chan rel =
  pp_set "; " (pp_equation hints) chan rel;;

let pp_relation_b hints chan = function
  | None -> output_string chan "bot"
  | Some rel -> pp_relation hints chan rel;;

let pp_results f (invariants : RB.t PMap.t) hints chan =
  let max_pc = P.to_int (RTL.max_pc_function f) in
  for pc=max_pc downto 1
  do
    Printf.fprintf chan "%d: %a\n\n" pc
      (pp_relation_b hints) (PMap.get (P.of_int pc) invariants)
  done

module IntSet=Set.Make(struct type t=int let compare = ( - ) end);;

let rec union_list prev = function
  | [] -> prev
  | h::t -> union_list (RB.lub prev h) t;;

let rb_glb (x : RB.t) (y : RB.t) : RB.t =
  match x, y with
  | None, _ | _, None -> None
  | (Some x'), (Some y') -> Some (RELATION.glb x' y');;

let refine_invariants
      (nodes : RTL.node list)
      (entrypoint : RTL.node)
      (successors : RTL.node -> RTL.node list)
      (predecessors : RTL.node -> RTL.node list)
      (tfr : RTL.node -> RB.t -> RB.t) (invariants0 : RB.t PMap.t) =
  let todo = ref IntSet.empty
  and invariants = ref invariants0 in  
  let add_todo (pc : RTL.node) =
    todo := IntSet.add (P.to_int pc) !todo in 
  let update_node (pc : RTL.node) =
    (if !Clflags.option_debug_compcert > 9
     then Printf.printf "updating node %d\n" (P.to_int pc));
    if not (peq pc entrypoint)
    then
      let cur = PMap.get pc !invariants in
      let nxt = union_list RB.bot
                  (List.map
                     (fun pred_pc->
                       rb_glb cur
                         (tfr pred_pc (PMap.get pred_pc !invariants)))
                     (predecessors pc)) in
      if not (RB.beq cur nxt)
      then
        begin
          (if !Clflags.option_debug_compcert > 4
           then Printf.printf "refining CSE3 node %d\n" (P.to_int pc));
          List.iter add_todo (successors pc)
        end in
  (List.iter add_todo nodes);
  while not (IntSet.is_empty !todo) do
    let nxt = IntSet.max_elt !todo in
    todo := IntSet.remove nxt !todo;
    update_node (P.of_int nxt)
  done;
  !invariants;;

let get_default default x ptree =
  match PTree.get x ptree with
  | None -> default
  | Some y -> y;;

let refine_analysis ctx tenv
      (f : RTL.coq_function) (invariants0 : RB.t PMap.t) =
  let succ_map = RTL.successors_map f in
  let succ_f x = get_default [] x succ_map in
  let pred_map = Kildall.make_predecessors f.RTL.fn_code RTL.successors_instr in
  let pred_f x = get_default [] x pred_map in
  let tfr = apply_instr' ctx tenv f.RTL.fn_code in
  refine_invariants
    (List.map fst (PTree.elements f.RTL.fn_code))
    f.RTL.fn_entrypoint succ_f pred_f tfr invariants0;;
  
let preanalysis (tenv : typing_env) (f : RTL.coq_function) =
  let cur_eq_id = ref 0
  and cur_catalog = ref PTree.empty
  and eq_table = Hashtbl.create 100
  and rhs_table = Hashtbl.create 100
  and cur_kill_reg = ref (PMap.init PSet.empty)
  and cur_kill_mem = ref PSet.empty
  and cur_kill_store = ref PSet.empty
  and cur_moves = ref (PMap.init PSet.empty) in
  let eq_find_oracle node eq =
    assert (not (is_trivial eq));
    let o = Hashtbl.find_opt eq_table (flatten_eq eq) in
    (if o = None then failwith "eq_find_oracle");
    (if !Clflags.option_debug_compcert > 5
     then Printf.printf "@%d: eq_find %a -> %a\n" (P.to_int node)
            pp_eq eq (pp_option pp_P) o);
    o
  and rhs_find_oracle node sop args =
    let o =
      match Hashtbl.find_opt rhs_table (sop, List.map P.to_int args) with
      | None -> PSet.empty
      | Some s -> s in
    (if !Clflags.option_debug_compcert > 5
     then Printf.printf "@%d: rhs_find %a = %a\n"
            (P.to_int node) pp_rhs (sop, args) pp_intset o);
    o in
  let mutating_eq_find_oracle node eq : P.t option =
    let (flat_eq_lhs, flat_eq_op, flat_eq_args) as flat_eq = flatten_eq eq in
    let o =
    match Hashtbl.find_opt eq_table flat_eq with
    | Some x ->
       Some x
    | None ->
       (* TODO print_eq stderr flat_eq; *)
       incr cur_eq_id;
       let id = !cur_eq_id in
       let coq_id = P.of_int id in
       begin
         Hashtbl.add eq_table flat_eq coq_id;
         (cur_catalog := PTree.set coq_id eq !cur_catalog);
         Hashtbl.add rhs_table (flat_eq_op, flat_eq_args)
           (PSet.add coq_id
              (match Hashtbl.find_opt rhs_table (flat_eq_op, flat_eq_args) with
               | None -> PSet.empty
               | Some s -> s));
         List.iter
           (fun reg -> imp_add_i_j cur_kill_reg reg coq_id)
           (eq.eq_lhs :: eq.eq_args);
         (if eq_depends_on_mem eq
          then cur_kill_mem := PSet.add coq_id !cur_kill_mem);
         (if eq_depends_on_store eq
          then cur_kill_store := PSet.add coq_id !cur_kill_store);
         (match eq.eq_op, eq.eq_args with
          | (SOp Op.Omove), [rhs] -> imp_add_i_j cur_moves eq.eq_lhs coq_id
          | _, _ -> ());
         Some coq_id
       end
    in
    (if !Clflags.option_debug_compcert > 5
     then Printf.printf "@%d: mutating_eq_find %a -> %a\n" (P.to_int node)
      pp_eq eq (pp_option pp_P) o);    
    o
  in
  let ctx = { eq_catalog     = (fun eq_id -> PTree.get eq_id !cur_catalog);
              eq_find_oracle = mutating_eq_find_oracle;
              eq_rhs_oracle  = rhs_find_oracle ;
              eq_kill_reg    = (fun reg -> PMap.get reg !cur_kill_reg);
              eq_kill_mem    = (fun () -> !cur_kill_mem);
              eq_kill_store  = (fun () -> !cur_kill_store);
              eq_moves       = (fun reg -> PMap.get reg !cur_moves)
            } in
  match internal_analysis ctx tenv f
  with None -> failwith "CSE3analysisaux analysis failed, try re-running with -fno-cse3"
     | Some invariants ->
        let invariants' =
          if ! Clflags.option_fcse3_refine
          then refine_analysis ctx tenv f invariants
          else invariants
        and hints = { hint_eq_catalog    = !cur_catalog;
                      hint_eq_find_oracle= eq_find_oracle;
                      hint_eq_rhs_oracle = rhs_find_oracle } in
        (if !Clflags.option_debug_compcert > 1
         then pp_results f invariants' hints stdout);
        invariants', hints
;;