(* *************************************************************)
(*                                                             *)
(*             The Compcert verified compiler                  *)
(*                                                             *)
(*           Sylvain Boulmé  Grenoble-INP, VERIMAG             *)
(*                                                             *)
(*  Copyright VERIMAG. All rights reserved.                    *)
(*  This file is distributed under the terms of the INRIA      *)
(*  Non-Commercial License Agreement.                          *)
(*                                                             *)
(* *************************************************************)

(*

This file implements the [branch_target] oracle that identifies "nop" branches in a LTL function,
and computes their target node with the distance (ie the number of cummulated nops) toward this target.

See [Tunneling.v]

*)

open Coqlib
open LTL
open Maps
open Camlcoq

let limit_tunneling = None (* for debugging: [Some x] limit the number of iterations *)
let debug_flag = ref false
let final_dump = false   (* set to true to have a more verbose debugging *)

let debug fmt =
  if !debug_flag then Printf.eprintf fmt
  else Printf.ifprintf stderr fmt

exception BugOnPC of int

(* type of labels in the cfg *)
type label = int * P.t

(* instructions under analyzis *)
type simple_inst = (* a simplified view of LTL instructions *)
  LBRANCH of node
| LCOND of node * node
| OTHER
and node = {
    lab : label;
    mutable inst: simple_inst;
    mutable link: node; (* link in the union-find: itself for non "nop"-nodes, target of the "nop" otherwise *)
    mutable dist: int;
    mutable tag: int
  }

(* type of the (simplified) CFG *)
type cfg = {
    nodes: (int, node) Hashtbl.t;
    mutable rems: node list; (* remaining conditions that may become lbranch or not *)
    mutable num_rems: int;
    mutable iter_num: int (* number of iterations in elimination of conditions *)
  }

let lab_i (n: node): int = fst n.lab
let lab_p (n: node): P.t = snd n.lab

let rec target c n = (* inspired from the "find" of union-find algorithm *)
  match n.inst with
  | LCOND(s1,s2) ->
     if n.link != n
     then update c n
     else if n.tag < c.iter_num then (
       (* we try to change the condition ... *)
       n.tag <- c.iter_num; (* ... but at most once by iteration *)
       let ts1 = target c s1 in
       let ts2 = target c s2 in
       if ts1 == ts2 then (n.link <- ts1; ts1) else n
     ) else n
  | _ ->
     if n.link != n
     then update c n
     else n
and update c n =
  let t = target c n.link in
  n.link <- t; t

let get_node c p =
  let li = P.to_int p in
  try
    Hashtbl.find c.nodes li
  with
    Not_found ->
      let rec n = { lab = (li, p); inst = OTHER; link = n ; dist = 0; tag = 0 }  in
      Hashtbl.add c.nodes li n;
      n

let set_branch c p s =
  let li = P.to_int p in
  try
    let n = Hashtbl.find c.nodes li in
    n.inst <- LBRANCH s;
    n.link <- target c s
  with
    Not_found ->
      let n = { lab = (li,p); inst = LBRANCH s; link = target c s; dist = 0; tag = 0 } in
      Hashtbl.add c.nodes li n


(* build [c.nodes] and accumulate in [acc] conditions at beginning of LTL basic-blocks *)
let build_simplified_cfg c acc pc bb =
  match bb with
  | Lbranch s :: _ ->
     let ns = get_node c s in
     set_branch c pc ns;
     acc
  | Lcond (_, _, s1, s2, _) :: _ ->
     c.num_rems <- c.num_rems + 1;
     let ns1 = get_node c s1 in
     let ns2 = get_node c s2 in
     let npc = get_node c pc in
     npc.inst <- LCOND(ns1, ns2);
     npc::acc
  | _ -> acc

(* try to change a condition into a branch
[acc] is the current accumulator of conditions to consider in the next iteration of repeat_change_cond
*)
let try_change_cond c acc pc =
  match pc.inst with
  | LCOND(s1,s2) ->
     let ts1 = target c s1 in
     let ts2 = target c s2 in
     if ts1 == ts2 then (
       pc.link <- ts1;
       c.num_rems <- c.num_rems - 1;
       acc
     ) else
       pc::acc
  | _  -> raise (BugOnPC (lab_i pc)) (* LCOND expected *)

(* repeat [try_change_cond] until no condition is changed into a branch *)
let rec repeat_change_cond c =
  c.iter_num <- c.iter_num + 1;
  debug "++ Tunneling.branch_target %d: remaining number of conds to consider = %d\n" (c.iter_num) (c.num_rems);
  let old =  c.num_rems in
  c.rems <- List.fold_left (try_change_cond c) [] c.rems;
  let curr = c.num_rems in
  let continue =
    match limit_tunneling with
    | Some n -> curr < old && c.iter_num < n
    | None -> curr < old
  in
  if continue
  then repeat_change_cond c


(* compute the final distance of each nop nodes to its target *)
let undef_dist = -1
let self_dist = undef_dist-1
let rec dist n =
  if n.dist = undef_dist
  then (
    n.dist <- self_dist; (* protection against an unexpected loop in the data-structure *)
    n.dist <-
      (match n.inst with
       | OTHER -> 0
       | LBRANCH p -> 1 + dist p
       | LCOND (p1,p2) -> 1 + (max (dist p1) (dist p2)));
    n.dist
  ) else if n.dist=self_dist then raise (BugOnPC (lab_i n))
    else n.dist

let final_export f c =
  let count = ref 0 in
  let filter_nops_init_dist _ n acc =
    let tn = target c n in
    if tn == n
    then (
      n.dist <- 0; (* force [n] to be a base case in the recursion of [dist] *)
      acc
    ) else (
      n.dist <- undef_dist; (* force [dist] to compute the actual [n.dist] *)
      count := !count+1;
      n::acc
    )
  in
  let nops = Hashtbl.fold filter_nops_init_dist c.nodes [] in
  let res = List.fold_left (fun acc n -> PTree.set (lab_p n) (lab_p n.link, Z.of_uint (dist n)) acc) PTree.empty nops in
  debug "* Tunneling.branch_target: final number of eliminated nops = %d\n" !count;
  res

(*********************************************)
(*** START: printing and debugging functions *)

let string_of_labeli nodes ipc =
  try
    let pc = Hashtbl.find nodes ipc in
    if pc.link == pc
    then Printf.sprintf "(Target@%d)" (dist pc)
    else Printf.sprintf "(Nop %d @%d)" (lab_i pc.link) (dist pc)
  with
    Not_found -> ""

let print_bblock c println (pc, bb) =
  match bb with
  | Lbranch s::_ -> (if println then debug "\n"); debug "%d:Lbranch %d %s\n" pc (P.to_int s) (string_of_labeli c.nodes pc); false
  | Lcond (_, _, s1, s2, _)::_ -> (if println then debug "\n"); debug "%d:Lcond (%d,%d) %s\n" pc (P.to_int s1) (P.to_int s2) (string_of_labeli c.nodes pc); false
  | _ -> debug "%d " pc; true


let print_cfg f c  =
  let a = Array.of_list (PTree.fold (fun acc pc bb -> (P.to_int pc,bb)::acc) f.fn_code []) in
  Array.fast_sort (fun (i1,_) (i2,_) -> i2 - i1) a;
  let ep = P.to_int f.fn_entrypoint in
  debug "entrypoint: %d %s\n" ep (string_of_labeli c.nodes ep);
  let println = Array.fold_left (print_bblock c) false a in
  (if println then debug "\n");debug "remaining cond:";
  List.iter (fun n -> debug "%d " (lab_i n)) c.rems;
  debug "\n"

(*************************************************************)
(* Copy-paste of the extracted code of the verifier          *)
(* with [raise (BugOnPC (P.to_int pc))] instead of [Error.*] *)

let get td pc =
  match PTree.get pc td with
  | Some p -> let (t0, d) = p in (t0, d)
  | None -> (pc, Z.of_uint 0)

let check_bblock td pc bb =
  match PTree.get pc td with
  | Some p ->
    let (tpc, dpc) = p in
    let dpc0 = dpc in
    (match bb with
     | [] ->
       raise (BugOnPC (P.to_int pc))
     | i :: _ ->
       (match i with
        | Lbranch s ->
          let (ts, ds) = get td s in
          if peq tpc ts
          then if zlt ds dpc0
               then ()
               else raise (BugOnPC (P.to_int pc))
          else raise (BugOnPC (P.to_int pc))
        | Lcond (_, _, s1, s2, _) ->
          let (ts1, ds1) = get td s1 in
          let (ts2, ds2) = get td s2 in
          if peq tpc ts1
          then if peq tpc ts2
               then if zlt ds1 dpc0
                    then if zlt ds2 dpc0
                         then ()
                         else raise (BugOnPC (P.to_int pc))
                    else raise (BugOnPC (P.to_int pc))
               else raise (BugOnPC (P.to_int pc))
          else raise (BugOnPC (P.to_int pc))
        | _ ->
          raise (BugOnPC (P.to_int pc))))
  | None -> ()

(** val check_code : coq_UF -> code -> unit res **)

let check_code td c =
  PTree.fold (fun _ pc bb -> check_bblock td pc bb) c (())

(*** END: copy-paste & debugging functions *******)

let branch_target f =
  debug "* Tunneling.branch_target: starting on a new function\n";
  if limit_tunneling <> None then debug "* WARNING: limit_tunneling <> None\n";
  let c = { nodes = Hashtbl.create 100; rems = []; num_rems = 0; iter_num = 0 } in
  c.rems <- PTree.fold (build_simplified_cfg c) f.fn_code [];
  repeat_change_cond c;
  let res = final_export f c in
  if !debug_flag then (
    try
      check_code res f.fn_code;
      if final_dump then print_cfg f c;
    with e -> (
      print_cfg f c;
      check_code res f.fn_code
    )
  );
  res