let error_prefix = "Clight Stack Size Annotator"
let error = Error.global_error error_prefix


let stack_id_prefix = "__stack_size"
let stack_max_id_prefix = "__stack_size_max"
let stack_incr_prefix = "__stack_size_incr"


let fun_stack_size = Hashtbl.create 10

let add_stack_size s i =
  Hashtbl.add fun_stack_size s i

let get_stack_size s =
  try
    Hashtbl.find fun_stack_size s
  with Not_found ->
    Printf.eprintf
      "get_stack_size: the function %s has not a stack size associated\n%!"
      s;
    exit 1

(* Instrumentation *)

let int_typ = Clight.Tint (Clight.I32, AST.Signed)

let const_int i = Clight.Expr (Clight.Econst_int i, int_typ)

(* Instrument a statement. *)

let rec instrument_body stack_incr stack_size stmt = match stmt with
  | Clight.Sskip | Clight.Sbreak | Clight.Scontinue
  | Clight.Sgoto _ ->
    stmt
  | Clight.Sassign (e1, e2) -> Clight.Sassign (e1, e2)
  | Clight.Scall (eopt, f, args) -> stmt
  | Clight.Ssequence (s1, s2) ->
    Clight.Ssequence (instrument_body stack_incr stack_size s1,
		      instrument_body stack_incr stack_size s2)
  | Clight.Sifthenelse (e, s1, s2) ->
    let s1' = instrument_body stack_incr stack_size s1 in
    let s2' = instrument_body stack_incr stack_size s2 in
    Clight.Sifthenelse (e, s1', s2')
  | Clight.Swhile (e, s) ->
    let s' = instrument_body stack_incr stack_size s in
    Clight.Swhile (e, s')
  | Clight.Sdowhile (e, s) ->
    let s' = instrument_body stack_incr stack_size s in
    Clight.Sdowhile (e, s')
  | Clight.Sfor (s1, e, s2, s3) ->
    let s1' = instrument_body stack_incr stack_size s1 in
    let s2' = instrument_body stack_incr stack_size s2 in
    let s3' = instrument_body stack_incr stack_size s3 in
    Clight.Sfor (s1', e, s2', s3')
  | Clight.Sreturn _ ->
    let fun_typ = Clight.Tfunction ([int_typ], Clight.Tvoid) in
    let f = Clight.Expr (Clight.Evar stack_incr, fun_typ) in
    let ss = Clight.Expr (Clight.Econst_int stack_size, int_typ) in
    let args = [Clight.Expr (Clight.Eunop(Clight.Oneg,ss),int_typ)] in
    Clight.Ssequence (Clight.Scall (None, f, args), stmt)
  | Clight.Sswitch (e, ls) ->
    let ls' = instrument_ls stack_incr stack_size ls in
    Clight.Sswitch (e, ls')
  | Clight.Slabel (lbl, s) ->
    let s' = instrument_body stack_incr stack_size s in
    Clight.Slabel (lbl, s')
  | Clight.Scost (lbl, s) ->
    (* Keep the cost label in the code. *)
    let s' = instrument_body stack_incr stack_size s in
    Clight.Scost (lbl, s')
  (*
    instrument_body stack_incr s
  *)
and instrument_ls stack_incr stack_size = function
  | Clight.LSdefault s ->
    let s' = instrument_body stack_incr stack_size s in
    Clight.LSdefault s'
  | Clight.LScase (i, s, ls) ->
    let s' = instrument_body stack_incr stack_size s in
    let ls' = instrument_ls stack_incr stack_size ls in
    Clight.LScase (i, s', ls')

(* Instrument a function. *)

let instrument_funct stack_incr (id, def) =
  let def = match def with
    | Clight.Internal def ->
      let stack_size = get_stack_size id in
      let body = instrument_body stack_incr
        stack_size def.Clight.fn_body in
      let body =
        let fun_typ = Clight.Tfunction ([int_typ], Clight.Tvoid) in
        let f = Clight.Expr (Clight.Evar stack_incr, fun_typ) in
        let args = [Clight.Expr (Clight.Econst_int stack_size, int_typ)] in
        Clight.Ssequence (Clight.Scall (None, f, args), body) in
      Clight.Internal { def with Clight.fn_body = body }
    | Clight.External _ -> def
  in
  (id, def)

(* This is the declaration of the cost variable. *)

let add_stack_decl stack_id =
  let init = [Clight.Init_int32 0] in
  ((stack_id, init), int_typ)

(* This is the definition of the increment cost function. *)

let stack_incr_def stack_id_max stack_id stack_incr =
  let int_var x = Clight.Expr (Clight.Evar x, int_typ) in
  let param = "incr" in
  let cost = int_var stack_id in
  let cost_max = int_var stack_id_max in
  let increment = int_var param in
  let stack_increment =
    Clight.Expr (Clight.Ebinop (Clight.Oadd, cost, increment), int_typ) in
  let stmt1 = Clight.Sassign (cost, stack_increment) in
  let stack_max =
    Clight.Expr (
      Clight.Econdition(
        Clight.Expr (Clight.Ebinop (Clight.Olt, cost_max, cost), int_typ),
          cost,cost_max), int_typ) in
  let stmt2 = Clight.Sassign (cost_max, stack_max) in
  let cfun =
    { Clight.fn_return = Clight.Tvoid ;
      Clight.fn_params = [(param, int_typ)] ;
      Clight.fn_vars   = [] ;
      Clight.fn_body = Clight.Ssequence(stmt1,stmt2) } in
  (stack_incr, Clight.Internal cfun)

(* Create a fresh uninitialized cost variable for each external function. This
   will be used by the Cost plug-in of the Frama-C platform. *)

let extern_stack_variables make_unique functs =
  let prefix = "_stack_of_" in
  let f (decls, map) (_, def) = match def with
    | Clight.Internal _ -> (decls, map)
    | Clight.External (id, _, _) ->
      let new_var = make_unique (prefix ^ id) in
      (decls @ [add_stack_decl new_var], StringTools.Map.add id new_var map) in
  List.fold_left f ([], StringTools.Map.empty) functs

let save_tmp tmp_file s =
  let cout = open_out tmp_file in
  output_string cout s ;
  flush cout ;
  close_out cout

(** [instrument prog stack_map] instruments the program [prog]. First a fresh
    global variable --- the so-called cost variable --- is added to the program.
    Then, each cost label in the program is replaced by an increment of the cost
    variable, following the mapping [stack_map]. The function also returns the
    name of the cost variable and the name of the cost increment function. *)

let instrument cost_incr p =

  (* Create a fresh 'cost' variable. *)
  let names = ClightAnnotator.names p in
  let make_unique = StringTools.make_unique names in
  let stack_id = make_unique stack_id_prefix in
  let stack_max_id = make_unique stack_max_id_prefix in
  let stack_decl = add_stack_decl stack_id in
  let stack_max_decl = add_stack_decl stack_max_id in

  (* Create a fresh uninitialized global variable for each extern function. *)
  let (extern_stack_decls, extern_stack_variables) =
    extern_stack_variables make_unique p.Clight.prog_funct in

  (* Define an increment function for the cost variable. *)
  let stack_incr =
    StringTools.Gen.fresh_prefix (StringTools.Set.add stack_id names)
      stack_incr_prefix in
  let stack_incr_def = stack_incr_def stack_max_id stack_id stack_incr in

  (* Instrument each function *)
  let prog_funct =
    List.map (fun ((id,_) as a) ->
      if id = cost_incr then a else instrument_funct stack_incr a)
      p.Clight.prog_funct in

  (* Glue all this together. *)
  let prog_vars = stack_decl :: stack_max_decl
    :: extern_stack_decls @ p.Clight.prog_vars in
  let prog_funct = stack_incr_def :: prog_funct in
  let p' =
    { p with Clight.prog_vars = prog_vars ; Clight.prog_funct = prog_funct } in

  (* Save the resulted program. Then re-parse it.
     Indeed, the instrumentation may add side-effects in expressions, which is
     not Clight compliant. Re-parsing the result with CIL will remove these
     side-effects in expressions to obtain a Clight program. *)
  let output = ClightPrinter.print_program p' in
  let res = ClightParser.process (`Source ("annotated", output)) in
  (res, stack_id, stack_max_id, stack_incr, extern_stack_variables)