open Parameters
open Cil_types
open Simplify_terms
module Varinfo = Cil_datatype.Varinfo
module Stmt = Cil_datatype.Stmt
module Term = Cil_datatype.Term
module Logic_label = Cil_datatype.Logic_label


(** This module defines the main analysis of the plug-in. Its actions are:
    - build the CFG of the program;
    - initialize the static environment of analysis (parameters of the
    functions, number of loops, etc);
    - compute the cost of each function depending on the costs of the others;
    - try to solve the inequations formed from the previous step so as to obtain
      an independent cost for each function;
    - save the results;
    - add the annotations on the program. *)


(** Function computing cost of special instruction  *)

type env =
  { env_fun_cost : term Varinfo.Hashtbl.t;
    env_cost_incr_name : string; (* TODO put the varinfo *)
    env_cost_var_max : logic_var;
    env_cost_var : logic_var;
  }


let cost_loop env kern_fun stmt lab_before_loop cost_body =
  (** invariant cost_max <= max(\old(cost_max), \old(cost) + [cost_body])
      invariant cost = \old(cost)
  *)
  (** first *)
  let cost_var_max = Logic_const.tvar env.env_cost_var_max in
  let old_cost_max = make_at lab_before_loop cost_var_max in
  let cost_var = Logic_const.tvar env.env_cost_var in
  let old_cost = make_at lab_before_loop cost_var in
  let cost_iterations = make_binop PlusA old_cost cost_body in
  let invariant = if Term.equal old_cost cost_iterations
    then Logic_const.prel (Req, cost_var_max, old_cost_max)
    else
      let new_cost = make_max old_cost_max cost_iterations in
      let new_cost = simplify_term new_cost in
      Logic_const.prel (Rle, cost_var_max, new_cost) in
  let annot =
    Logic_const.new_code_annotation
      (AInvariant ([stack_cost_behavior],true,invariant)) in
  Annotations.add kern_fun stmt [Ast.self] (User annot);
  (** second *)
  let invariant = Logic_const.prel (Req, cost_var, old_cost) in
  let annot =
    Logic_const.new_code_annotation
      (AInvariant ([stack_cost_behavior],true,invariant)) in
  Annotations.add kern_fun stmt [Ast.self] (User annot);
  cost_body

(** computing cost of fundec which contain a call to funvar *)
let cost_of_fun_call env funvar args =
  (** can't fail since we use the callgraph *)
  let cost = Varinfo.Hashtbl.find env.env_fun_cost funvar in
  (* match cost.term_node with *)
  (* | TConst (CInt64 (_,_,_)) -> cost *)
  (* | _ -> *)
  (** arguments of the function *)
  let formals = Cil.getFormalsDecl funvar in
(*  (** The real arguments are binded using a let and the formal argument are
      replaced by fresh variable *)
  let letvars = List.map
    (fun vi -> Cil.make_temp_logic_var (Ctype vi.vtype))
    formals in
  let subst = List.fold_left2 (fun map vi v ->
    Varinfo.Map.add vi v map) Varinfo.Map.empty formals letvars in
  let cost =
    Visitor.visitFramacTerm (new subst_lval subst (Project.current ()))
      cost in
  let cost = List.fold_left2 (fun tlet v arg ->
    let logic_info = { l_var_info = v; l_type = None; l_tparams = [];
                       l_labels = []; l_profile = [];
                       l_body = LBnone; } in
    let arg = term_of_exp arg in
    let tlet =
      Logic_const.term (Tlet (logic_info, arg))
        tlet.term_type in
    tlet)
    cost letvars args in*)
  let subst = List.fold_left2 (fun map vi arg ->
    try
      Varinfo.Map.add vi (term_of_exp arg) map
    with Untranslatable_expr _ -> map (** TODO verify that this formal doesn't
                                          appear in the cost of the function *)
  ) Varinfo.Map.empty formals args in
  let cost = remove_logic_label Logic_const.pre_label cost in
  let cost =
    Visitor.visitFramacTerm (new subst_lval subst (Project.current ()))
      cost in
  cost


(** We go through the cfg we assume that:
    - all functions called have a cost already computed
    - no implicit loop (goto), all are explicit Loop
    - all loops have an invariant (given by the user or previously computed for
    for-loops)

    We iterate on the cfg except when arriving on a loop where we compute the
    body alone and after that go to the next statements (the break statement)

   break_statement in which statement to stop the sum (for loop)
 *)


let rec find_variant = function
  | [] -> None
  | {annot_content = AVariant (t,_)}::_ -> Some t
  | _::l -> find_variant l

(** TODO: cache the statement already done after if-then-else
    (otherwise exponential) *)
(** return None if no path arrive to a good_statement or don't go through
    a bad_statetement *)
let rec cfg_sum env kern_fun bad_stmts stop_stmts stmt =
  if !debug then
    Format.printf "sid: %i,@ bad_stmts: %a,@ stop_stmts: %a,@\n"
      stmt.sid
      Stmt.Set.pretty bad_stmts
      Stmt.Set.pretty stop_stmts;
  (** statement in last statement *)
  if Stmt.Set.mem stmt stop_stmts
  then Some (Cil.lzero ())
  else if Stmt.Set.mem stmt bad_stmts
  then None
  else
  (** successors *)
  let succs_bad_stmts =
    match stmt.skind with
    | Loop (_,_,_,_,_) -> Stmt.Set.add stmt bad_stmts
    | _ -> bad_stmts in
  (** Cost of the successors of the instruction *)
  let cost_succs =
    cfg_sum_succs env kern_fun succs_bad_stmts stop_stmts
      stmt.succs in
  match cost_succs with
  | None -> None (** can't go in a good state so count nothing *)
  | Some cost_succs ->
  (** normalize the cost of the instruction regarding the labels *)
  let labels = NormAtLabels.labels_stmt_pre stmt in
  (** Cost of the instruction *)

  let sum =
    match stmt.skind with
    (* (\** call to cost_incr *\) *)
    (** cost_stack_incr is not an usual function *)
    | Instr
        (Call (None,
                         { enode =
                             Lval (Var var, _)} ,
                         [cost], _))
        when var.vname = env.env_cost_incr_name ->
      make_binop PlusA cost_succs
        (make_max (Cil.lzero ()) (term_of_cost_expr cost))
    (** function call *)
    | Instr
        (Call (_, { enode =
            Lval (Var var, _)},
                         args,_)) ->
      let cost_instr = NormAtLabels.preproc_annot_term labels
        (cost_of_fun_call env var args) in
      make_max cost_succs cost_instr
    (** loops *)
    | Loop (_,_,_,_,_) ->
      (** the first one perhaps not needed *)
      let bad_stmts = Stmt.Set.union bad_stmts stop_stmts in
      let stop_stmts = Stmt.Set.singleton stmt in
      let cost_body =
        cfg_sum_succs env kern_fun bad_stmts stop_stmts stmt.succs in
      let cost_body = match cost_body with
        | None -> (* no cycle in fact *) Cil.lzero ()
        | Some cost -> simplify_term cost in
      if stmt.labels = [] then
        stmt.labels <- (** ugly but... *)
          Label("__stmt_"^(string_of_int stmt.sid),Cil.builtinLoc,false)::
          stmt.labels;
      let label_before_loop = StmtLabel (ref stmt) in
      let cost_instr =
        cost_loop env kern_fun stmt label_before_loop cost_body in
      let cost_instr = NormAtLabels.preproc_annot_term labels cost_instr in
      make_max cost_instr cost_succs
    | _ -> cost_succs in
  (** try to propagate \at(i,Label) *)
  let sum = match stmt.skind with
    | Instr (Set ((Var vi,_),e,_)) ->
      begin try
        let t = term_of_exp e in
        let subst = Varinfo.Map.singleton vi t in
        let sum =
          Visitor.visitFramacTerm
            (new subst_with_label stmt subst (Project.current ()))
            sum in
        sum
      with exn ->
        if !debug then Format.eprintf "can't convert exp %a to term: %a.@."
          Cil.defaultCilPrinter#pExp e print_exception exn;
        sum end
      | _ -> sum
  in
  (** remove the at that point to this statement we don't know the value at
      this point perhaps we will know before ( use Here?? ) *)
  let sum =
    Visitor.visitFramacTerm
      (new remove_stmt_label stmt (Project.current ()))
      sum in
  Some sum

and cfg_sum_succs env kern_fun bad_stmts stop_stmts stmts =
  let fold cost stmt =
    let cost_succ = cfg_sum env kern_fun bad_stmts stop_stmts stmt in
    make_max_opt cost cost_succ in
  List.fold_left fold None stmts

let make_leq_cost env sum =
    (** ensures cost_max <= max(\old(cost_max), [sum] + cost) *)
    let cost_var_max = Logic_const.tvar env.env_cost_var_max in
    let cost_var = Logic_const.tvar env.env_cost_var in
    let old_cost = Logic_const.told cost_var in
    let old_cost_max = Logic_const.told cost_var_max in
    let expr = make_binop PlusA sum old_cost in
    let expr = make_max old_cost_max expr in
    let post = Logic_const.prel (Rle, cost_var_max, expr) in
    let post = Logic_const.new_predicate post in
    post

let make_eq_cost env =
  (** ensures cost = old cost *)
  let cost_var = Logic_const.tvar env.env_cost_var in
  let old_cost = Logic_const.told cost_var in
  let post = Logic_const.prel (Req, cost_var, old_cost) in
  let post = Logic_const.new_predicate post in
  post


let fun_sum env vi =
  let kernel_fun = Globals.Functions.get vi in
  let fst_stmt = Kernel_function.find_first_stmt kernel_fun in
  let last_stmt = Kernel_function.find_return kernel_fun in
  let bad_stmts = Stmt.Set.empty in
  let stop_stmts = Stmt.Set.singleton last_stmt in
  let sum = cfg_sum env kernel_fun bad_stmts stop_stmts fst_stmt in
  let sum = match sum with
    | None -> assert false (* Oups return not achievable?? *)
          (* Infinity!! *)
    | Some sum ->
      let sum = simplify_term sum in
      let sum = Logic_const.told sum in (** add old and push it *)
      let labels = NormAtLabels.labels_fct_post in
      let sum = NormAtLabels.preproc_annot_term labels sum in
      sum in
  if !debug then
    Format.printf "@[<hov 3>sum of %a: %a@]@."
      Cil.defaultCilPrinter#pVar vi
      Cil.defaultCilPrinter#pTerm sum;
(** add postcondition about cost *)
  let mod_funspec funspec =
    let post = make_leq_cost env sum in
    let post2 = make_eq_cost env in
    let behavior =
      Cil.mk_behavior ~name:stack_cost_behavior ~post_cond:[Normal, post;
                                                      Normal, post2] () in
    {funspec with spec_behavior = behavior::funspec.spec_behavior;}
  in
  Kernel_function.set_spec kernel_fun mod_funspec;
  sum


let initialize _tmp_prefix _fname
    cost_var cost_var_max cost_incr _extern_costs =
  (** TODO? use extern_costs for initializing env_fun_cost ? *)
  let callgraph = Callgraph.computeGraph (Ast.get ()) in
  if !debug then
    Printf.printf "callgraph size: %a\n%!" Callgraph.printGraph callgraph;
  let env = { env_fun_cost = Varinfo.Hashtbl.create 10;
              env_cost_incr_name = cost_incr;
              env_cost_var_max = cost_var_max;
              env_cost_var = cost_var
            } in
  let rec iter : 'a. 'a -> Callgraph.callnode -> unit = fun _ callnode ->
    match callnode.Callgraph.cnInfo with
    | Callgraph.NIVar (vi,{contents = true }) ->
      if not (Varinfo.Hashtbl.mem env.env_fun_cost vi) then begin
        Inthash.iter iter callnode.Callgraph.cnCallees;
        let sum = fun_sum env vi in
        Varinfo.Hashtbl.add env.env_fun_cost vi sum
      end
    | _ -> () (** put an undefined constant here for top *)
  in
  (** add the cost of cost_incr *)
  let cost_incr = Hashtbl.find callgraph cost_incr in
  let cost_incr = match cost_incr.Callgraph.cnInfo with
    | Callgraph.NIVar (vi,_) -> vi
    | _ -> assert false (* cost_incr must be at least declared *) in
  let arg_cost_incr = List.hd (Cil.getFormalsDecl cost_incr) in
  let sum_cost_incr = Logic_const.tvar (Cil.cvar_to_lvar arg_cost_incr) in
  let post_cost_incr = make_leq_cost env sum_cost_incr in
  let post_cost_incr2 =
    (** ensures cost = cost + incr *)
    let cost_var = Logic_const.tvar env.env_cost_var in
    let old_cost = Logic_const.told cost_var in
    let expr = make_binop PlusA sum_cost_incr old_cost in
    let post = Logic_const.prel (Req, cost_var, expr) in
    let post = Logic_const.new_predicate post in
    post in
  Kernel_function.set_spec (Globals.Functions.get cost_incr) (fun funspec ->
    let assigns =
      Writes [Logic_const.(new_identified_term (tvar cost_var_max)), FromAny;
              Logic_const.(new_identified_term (tvar cost_var)), FromAny] in
    let behavior =
      Cil.mk_behavior ~name:stack_cost_behavior
        ~post_cond:[Normal, post_cost_incr;Normal, post_cost_incr2]
        ~assigns () in
    {funspec with spec_behavior = behavior::funspec.spec_behavior;} );
  Varinfo.Hashtbl.add env.env_fun_cost cost_incr sum_cost_incr;
  (** make the other functions *)
  Hashtbl.iter iter callgraph

(** Save file *)
let save_file fname =
  Kernel.CodeOutput.set fname ;
  File.pretty_ast ()

(*** Main ***)

exception NoStackInst

let cost ((fname, _), _, _cost_time, cost_stack) =
  match cost_stack with
  | None -> ()
  | Some (cost_stack_max, cost_stack_id) ->
    try
      if !debug then Printf.printf "Initialize stack size\n%!" ;
      let cost_var_max =
        try
          Cil.cvar_to_lvar (Globals.Vars.find_from_astinfo
                              cost_stack_max.Cerco.cost_id VGlobal)
        with _ -> raise NoStackInst in
      let cost_var =
        Cil.cvar_to_lvar (Globals.Vars.find_from_astinfo
                            cost_stack_id VGlobal) in
      initialize "__tmp" fname
        cost_var cost_var_max
        cost_stack_max.Cerco.cost_incr
        cost_stack_max.Cerco.extern_costs;
    with
    | NoStackInst ->
      Format.eprintf "No stack information found.\n%!"
    | e -> Format.eprintf "** ERROR: %a.@." print_exception_raise e