(** This module gives utility functions for creating terms and
     simplifying them *)

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

(** {Add needed logic function} *)

let get_max_logic_info, add_def_max_logic_info =
  (* TODO the name can collide with other logic function *)
  let module M = State_builder.Option_ref(Cil_datatype.Logic_info)
        (struct
          let name = "max_logic_fun"
          let dependencies = [Annotations.self]
          let kind = `Internal
         end) in
  let create () =
    let x  = Cil_const.make_logic_var "x" Linteger in
    let y  = Cil_const.make_logic_var "y" Linteger in
    let tx = Logic_const.tvar x in let ty = Logic_const.tvar y in
    let body =
      let cond = Logic_const.term (TBinOp (Gt,tx,ty)) Linteger in
      let term = Logic_const.term (Tif (cond,tx,ty)) Linteger in
      term in
    let signature = Larrow ([Linteger;Linteger],Linteger) in
    let li = { l_var_info = Cil_const.make_logic_var "__max" signature;
               l_labels = [];
               l_tparams = [];
               l_type = Some Linteger;
               l_profile = [x;y];
               l_body = LBterm body } in
    li
  in
  (fun () -> M.memo create),
  (fun () ->
    match M.get_option () with
    | None -> ()
    | Some li ->
      Globals.Annotations.add_generated
        (Dfun_or_pred (li, Cil_datatype.Location.unknown)))


(** {2 Simplify casts} *)
(** Casts may get in the way when looking for a loop counter, just remove
    them. *)

let rec remove_casts e = match e.enode with
  | Lval lval ->
    { e with enode = Lval (remove_casts_lval lval) }
  | SizeOfE e ->
    { e with enode = SizeOfE (remove_casts e) }
  | AlignOfE e ->
    { e with enode = AlignOfE (remove_casts e) }
  | UnOp (unop, e, typ) ->
    { e with enode = UnOp (unop, remove_casts e, typ) }
  | BinOp (binop, e1, e2, typ) ->
    let enode =
      BinOp (binop, remove_casts e1, remove_casts e2, typ) in
    { e with enode }
  | CastE (_, e) -> remove_casts e
  | AddrOf lval ->
    { e with enode = AddrOf (remove_casts_lval lval) }
  | StartOf lval ->
    { e with enode = StartOf (remove_casts_lval lval) }
  | Info (e, info) ->
    { e with enode = Info (remove_casts e, info) }
  | _ -> e

and remove_casts_lval (lhost, offset) =
  (remove_casts_lhost lhost, remove_casts_offset offset)

and remove_casts_lhost = function
  | Mem e -> Mem (remove_casts e)
  | lhost -> lhost

and remove_casts_offset = function
  | Field (fieldinfo, offset) ->
    Field (fieldinfo, remove_casts_offset offset)
  | Index (e, offset) ->
    Index (remove_casts e, remove_casts_offset offset)
  | offset -> offset


(** {2 Simplify and convert terms} *)


(*** Exceptions ***)
exception Untranslatable_expr of exp

let print_exception_raise fmt = function
  | Untranslatable_expr exp ->
    Format.fprintf fmt "Can't convert exp %a to term.@."
      Cil.defaultCilPrinter#pExp exp;
  | e -> raise e

let print_exception fmt exn =
  try
    print_exception_raise fmt exn
  with _ -> Format.fprintf fmt "%s" (Printexc.to_string exn)


(** raised when a function call contains something that can't be translated to
    term *)

(** cost of incr_cost (cost of all the instruction until the next one) *)
let rec term_of_exp e =
  let open Cil_types in
  match e.enode with
  | Const (CInt64 _ as c) ->
    Logic_const.term (TConst c) (Ctype (Cil.typeOf e))
  | Lval (Var varinfo, offset) ->
    let logic_var = Cil.cvar_to_lvar varinfo in
    let offset = term_offset_of_offset offset in
    Logic_const.term (TLval (TVar logic_var, offset)) (Ctype (Cil.typeOf e))
  | Info (e, _) -> term_of_exp e
  | UnOp (unop, e, typ) ->
    Logic_const.term (TUnOp (unop, term_of_exp e)) (Ctype typ)
  | BinOp (binop, e1, e2, typ) ->
    Logic_const.term (TBinOp (binop, term_of_exp e1, term_of_exp e2))
      (Ctype typ)
  | CastE (typ, e) ->
    Logic_const.term (TCastE (typ, term_of_exp e)) (Ctype typ)
  | _ -> raise (Untranslatable_expr e)

and term_offset_of_offset = function
  | NoOffset -> TNoOffset
  | Field (fi, offset) ->
    TField(fi, term_offset_of_offset offset)
  | Index (exp, offset) ->
    TIndex(term_of_exp exp, term_offset_of_offset offset)

let term_of_exp e =
  let e = if remove_casts_in_annotations then remove_casts e else e in
  term_of_exp e

(** Same as before but use integer for all the types *)
let term_of_cost_expr e =
  match e.enode with
  | Const (CInt64 _ as c) ->
    Logic_const.term (TConst c) Linteger
  (** TODO: convert the other function use integer when needed *)
  | _ -> raise (Untranslatable_expr e)

(** Operation simplifying (** same as termFoldConst ? *) *)

let make_at label x =
  let open Cil_types in
  match x.term_node with
  | TConst (CInt64 (_,_,_)) -> x
  | _ -> Logic_const.tat (x,label)

(** Use only for cost *)
let make_max x y =
  match x, y with
    (** two constants *)
  | {term_node = TConst (CInt64 (xn,_,_))},
    {term_node = TConst (CInt64 (yn,_,_))} ->
    (Cil.lconstant (My_bigint.max xn yn))
  (*     (\** one null *\) *) (* correct but dangerous *)
  (* | {term_node = TConst (CInt64 (zero,_,_))}, other *)
  (* | other, {term_node = TConst (CInt64 (zero,_,_))} *)
  (*   when My_bigint.is_zero zero -> other *)
    (** generic case *)
  | x, y ->
    Logic_const.term (Tapp (get_max_logic_info (), [], [x;y])) Linteger

let make_max_opt x y =
  match x, y with
    (** one inaccessible *)
  | None, other | other, None -> other
  | Some x, Some y -> Some (make_max x y)

let make_sign = function
    (** two constants *)
  | {term_node = TConst (CInt64 (xn,_,_))} ->
    if My_bigint.ge xn My_bigint.zero
    then (Cil.lconstant (My_bigint.one))
    else (Cil.lconstant (My_bigint.minus_one))
  (*     (\** one null *\) *) (* correct but dangerous *)
  (* | {term_node = TConst (CInt64 (zero,_,_))}, other *)
  (* | other, {term_node = TConst (CInt64 (zero,_,_))} *)
  (*   when My_bigint.is_zero zero -> other *)
    (** generic case *)
  | x ->
    let cond = Logic_const.term (TBinOp (Ge,x,Cil.lzero ())) Linteger in
    let term = Logic_const.term (Tif (cond,
                                      (Cil.lconstant (My_bigint.one)),
                                      (Cil.lconstant (My_bigint.minus_one))))
      Linteger in
    term

let make_unop unop x =
  match unop, x.term_node with
  | Neg, TUnOp(Neg,y) -> y
  | Neg, TConst (CInt64 (xn,_,_)) -> Cil.lconstant (My_bigint.neg xn)
  | _ -> Logic_const.term (TUnOp (unop,x)) Linteger

let make_binop binop x y =
  match binop, x.term_node, y.term_node with
  | PlusA, TConst (CInt64 (xn,_,_)), TConst (CInt64 (yn,_,_)) ->
    Cil.lconstant (My_bigint.add xn yn)
  | PlusA, TConst (CInt64 (zero,_,_)), _ when My_bigint.is_zero zero -> y
  | PlusA, _, TConst (CInt64 (zero,_,_)) when My_bigint.is_zero zero -> x

  | MinusA, TConst (CInt64 (xn,_,_)), TConst (CInt64 (yn,_,_)) ->
    Cil.lconstant (My_bigint.sub xn yn)
  | MinusA, TConst (CInt64 (zero,_,_)), _ when My_bigint.is_zero zero ->
    make_unop Neg y
  | MinusA, _, TConst (CInt64 (zero,_,_)) when My_bigint.is_zero zero -> x
  | Mult, TConst (CInt64 (xn,_,_)), TConst (CInt64 (yn,_,_)) ->
    Cil.lconstant (My_bigint.mul xn yn)
  | Mult, TConst (CInt64 (zero,_,_)), _ when My_bigint.is_zero zero -> x
  | Mult, _, TConst (CInt64 (zero,_,_)) when My_bigint.is_zero zero -> y
  | Mult, TConst (CInt64 (one,_,_)), _ when My_bigint.is_one one -> y
  | Mult, _, TConst (CInt64 (one,_,_)) when My_bigint.is_one one -> x
  | Div, TConst (CInt64 (xn,_,_)), TConst (CInt64 (yn,_,_)) ->
    Cil.lconstant (My_bigint.div xn yn)
  (* | Div, TConst (CInt64 (zero,_,_)), _ when My_bigint.is_zero zero -> x *)
  (* | Div, _, TConst (CInt64 (zero,_,_)) when My_bigint.is_zero zero -> y *)
  | Div, _, TConst (CInt64 (one,_,_)) when My_bigint.is_one one -> x
  | Div, _, TConst (CInt64 (minus_one,_,_))
    when My_bigint.equal My_bigint.minus_one minus_one ->
    make_unop Neg x
  | _ -> Logic_const.term (TBinOp (binop,x,y)) Linteger

let simplify_term_top t =
  match t.term_node with
  | TBinOp(binop,x,y) -> make_binop binop x y
  | TUnOp(unop,x) -> make_unop unop x
  | Tif({term_node =
      TBinOp ((Lt|Le|Gt|Ge|Eq|Ne) as rel,
              {term_node = TConst (CInt64 (xn,_,_))},
              {term_node = TConst (CInt64 (yn,_,_))})},t1,t2) ->
    let b = match rel with
      | Lt -> My_bigint.lt xn yn
      | Gt -> My_bigint.gt xn yn
      | Le -> My_bigint.le xn yn
      | Ge -> My_bigint.ge xn yn
      | Eq -> My_bigint.equal xn yn
      | Ne -> not (My_bigint.equal xn yn)
      | _ -> assert false in
    if b then t1 else t2
  | Tapp(f, [], [{term_node = TConst (CInt64 (xn,_,_))} as xt;
                 {term_node = TConst (CInt64 (yn,_,_))} as yt]) when
      Cil_datatype.Logic_info.equal f (get_max_logic_info ()) ->
    if My_bigint.gt xn yn then xt else yt
  | _ -> t

class simplify_term prj =
object inherit Visitor.frama_c_copy prj as super
       method vterm term = Cil.ChangeDoChildrenPost(term,simplify_term_top)
end

let rec top_const_form t =
  match t.term_node with
  | TConst _ -> t, (Cil.lzero ())
  | TUnOp(unop,x) ->
    let cx,x = top_const_form x in
    begin match unop with
    | Neg -> make_unop unop cx, make_unop unop x
    | _ -> (Cil.lzero ()),
      Logic_const.term (TUnOp (unop,make_binop PlusA cx x)) Linteger end
  | TBinOp(binop,x,y) ->
    let cx,x = top_const_form x in
    let cy,y = top_const_form y in
    begin match binop with
    | PlusA -> make_binop binop cx cy, make_binop binop x y
    | MinusA -> make_binop binop cx cy, make_binop binop x y
    | Mult -> make_binop binop cx cy,
      make_binop PlusA (make_binop binop x y)
        (make_binop PlusA (make_binop binop cx y)
           (make_binop binop cy x))
    | _ -> (Cil.lzero ()),
      Logic_const.term (TBinOp (binop, make_binop PlusA cx x,
                               make_binop PlusA cy y)) Linteger end
  | Tif(t1,t2,t3) ->
    let c1,t1 = top_const_form t1 in
    let c2,t2 = top_const_form t2 in
    let c3,t3 = top_const_form t3 in
    Cil.lzero (),
    Logic_const.term (Tif(make_binop PlusA c1 t1,
                          make_binop PlusA c2 t2,
                          make_binop PlusA c3 t3)) Linteger
  | Tapp(f, labs, l) ->
    Cil.lzero (),
    Logic_const.term
      (Tapp(f, labs, List.map (fun t -> let c,t = top_const_form t in
                                        make_binop PlusA c t) l))
      Linteger
  | _ -> (Cil.lzero ()),t


let top_const_form t =
  let ct,t = top_const_form t in
  make_binop PlusA ct t

let simplify_term term =
  let term =
    Visitor.visitFramacTerm (new simplify_term (Project.current ())) term in
  if put_const_at_the_top_in_annotations then top_const_form term else term


(** {2 Simply labels} *)

class subst_lval subst prj =
object inherit Visitor.frama_c_copy prj as super
       method vterm term =
         let open Cil_types in
         match term.term_node with
         | TLval (TVar {lv_origin = Some vi},TNoOffset)
             when Varinfo.Map.mem vi subst ->
           Cil.ChangeTo (Varinfo.Map.find vi subst)
         | _ -> Cil.DoChildren

end

class subst_with_label stmt_label subst prj  =
object inherit Visitor.frama_c_copy prj as super
       method vterm term =
         let open Cil_types in
         match term.term_node with
         | TLval (TVar {lv_origin = Some vi},TNoOffset) (**TODO other offset *)
             when Varinfo.Map.mem vi subst ->
           Cil.ChangeTo (Varinfo.Map.find vi subst)
         | Tat _ -> Cil.JustCopy (* don't touch if we don't see it *)
         | _ -> Cil.ChangeDoChildrenPost(term,simplify_term_top)

end

class remove_stmt_label stmt_label prj  =
object inherit Visitor.frama_c_copy prj as super
       method vterm term =
         let open Cil_types in
         match term.term_node with
         | Tat(t, (LogicLabel (Some stmt,_)))
             when Stmt.equal stmt stmt_label ->
           Cil.ChangeTo t
         | Tapp(li,l,t) ->
           let rec filter_label l =
             match l with
             | [] -> l
             (** TODO understand what the differences between la and lb *)
             | (( LogicLabel(Some stmt,_),_) | (_,LogicLabel(Some stmt,_)))::ll
                 when Stmt.equal stmt stmt_label -> filter_label ll
             | a::ll ->
               let ll' = filter_label ll in
               if ll == ll' then l else a::ll' in
           let l' = filter_label l in
           if l == l' then Cil.DoChildren
           else Cil.ChangeDoChildrenPost
             ({term with term_node = Tapp(li,l',t)},fun x -> x)
         | _ -> Cil.DoChildren
end


class remove_logic_label lab prj  =
object inherit Visitor.frama_c_copy prj as super
       method vterm term =
         let open Cil_types in
         match term.term_node with
         | Tat(t, label)
             when Logic_label.equal label lab ->
           Cil.ChangeTo t
         | Tapp(li,l,t) ->
           let rec filter_label l =
             match l with
             | [] -> l
             (** TODO understand what the differences between la and lb *)
             | (label1,label2)::ll
                 when Logic_label.equal label1 lab ||
                   Logic_label.equal label2 lab -> filter_label ll
             | a::ll ->
               let ll' = filter_label ll in
               if ll == ll' then l else a::ll' in
           let l' = filter_label l in
           if l == l' then Cil.DoChildren
           else Cil.ChangeDoChildrenPost
             ({term with term_node = Tapp(li,l',t)},fun x -> x)
         | _ -> Cil.DoChildren
end

let remove_logic_label lab cost =
  Visitor.visitFramacTerm
    (new remove_logic_label lab (Project.current ()))
    cost


let rec freevar acc t =
  match t.term_node with
  | TConst _ -> acc
  | TUnOp(_,x) -> freevar acc x
  | TBinOp(_,x,y) -> freevar (freevar acc x) y
  | Tif(t1,t2,t3) -> freevar (freevar (freevar acc t1) t2) t3
  | Tapp(_, _, l) -> List.fold_left freevar acc l
  | Tat(t,_) -> freevar acc t
  | (TLval(TVar lv,_)) -> Cil_datatype.Logic_var.Set.add lv acc
  | _ -> assert false


let tapp li args =
  let cast lv t =
    match lv.lv_type with
    | Ctype typ ->
      (** the main handle cases is t is an integer, lv an is int so we
          need a cast (int)t  *)
      (* if Logic_utils.is_same_type lv.lv_type t.term_type *)
      (* then t *)
      (* else *) Logic_const.term (TCastE(typ,t)) lv.lv_type
    | _  -> t  in
  let args = List.map2 cast li.l_profile args in
  Tapp(li,[],args)