(** This module provides functions to manipulate parameterized arithmetic
    expressions (see the ArithSig module) when instanciaded on a specific
    type. *)


module Make (S : ArithSig.S) = struct

  include ArithSig.P

  let compare = compare S.compare

  type t = S.t tt

  let fill_subs v subs = ArithSig.fill_subs v subs

  let subs v = ArithSig.subs v

  let fold f v = ArithSig.fold f v


  let f_to_string v subs_res = match v, subs_res with
    | A a, _ -> "(A " ^ (S.to_string a) ^ ")"
    | Int i, _ -> string_of_int i
    | Var x, _ -> x
    | UnOp (unop, _), v :: _ -> (string_of_unop unop) ^ "(" ^ v ^ ")"
    | BinOp (binop, _, _), v1 :: v2 :: _ ->
      "(" ^ v1 ^ ")" ^ (string_of_binop binop) ^ "(" ^ v2 ^ ")"
    | Cond (_, rel, _, _, _), t1 :: t2 :: tif :: telse :: _ ->
      Printf.sprintf "(%s %s %s)? (%s) : (%s)"
	t1 (string_of_relation rel) t2 tif telse
    | UnOp _, _ | BinOp _, _ | Cond _, _ ->
      assert false (* wrong number of arguments *)

  let to_string v = fold f_to_string v


  let is_supported_rel = function
    | Cil_types.Lt | Cil_types.Gt | Cil_types.Le | Cil_types.Ge -> true
    | _ -> false

  let rel_of_cil_binop = function
    | Cil_types.Lt -> Lt
    | Cil_types.Gt -> Gt
    | Cil_types.Le -> Le
    | Cil_types.Ge -> Ge
    | _ -> raise (Failure "Arith.rel_of_cil_binop")


  let cil_binop_of_rel = function
    | Lt -> Cil_types.Lt
    | Gt -> Cil_types.Gt
    | Le -> Cil_types.Le
    | Ge -> Cil_types.Ge

  let cil_unop_of_unop = function
    | Neg -> Cil_types.Neg

  let cil_binop_of_binop = function
    | Add -> Cil_types.PlusA
    | Sub -> Cil_types.MinusA
    | Mul -> Cil_types.Mult
    | Div -> Cil_types.Div
    | Mod -> Cil_types.Mod
    (* No direct translation. Handle this case in the calling function. *)
    | Max -> assert false


  let integer_term term = Logic_const.term term Cil_types.Linteger

  let tinteger i =
    let cint64 = Cil_types.CInt64 (My_bigint.of_int i, Cil_types.IInt, None) in
    let iterm = Cil_types.TConst cint64 in
    integer_term iterm

  let f_to_cil_term v subs_res = match v, subs_res with
    | A a, _ ->  S.to_cil_term a
    | Int i, _ -> tinteger i
    | Var x, _ ->
      Logic_const.tvar (Cil_const.make_logic_var x Cil_types.Linteger)
    | UnOp (unop, _), t :: _ ->
      integer_term (Cil_types.TUnOp (cil_unop_of_unop unop, t))
    | BinOp (Max, _, _), t1 :: t2 :: _ ->
      let test = integer_term (Cil_types.TBinOp (Cil_types.Ge, t1, t2)) in
      integer_term (Cil_types.Tif (test, t1, t2))
    | BinOp (binop, _, _), t1 :: t2 :: _ ->
      integer_term (Cil_types.TBinOp (cil_binop_of_binop binop, t1, t2))
    | Cond (_, rel, _, _, _), t1 :: t2 :: tif :: telse :: _ ->
      let test =
	integer_term (Cil_types.TBinOp (cil_binop_of_rel rel, t1, t2)) in
      integer_term (Cil_types.Tif (test, tif, telse))
    | _ -> assert false (* wrong number of arguments *)

  let to_cil_term v = fold f_to_cil_term v


  let cil_rel_of_rel = function
    | Lt -> Cil_types.Rlt
    | Gt -> Cil_types.Rgt
    | Le -> Cil_types.Rle
    | Ge -> Cil_types.Rge


  let of_int i = Int i
  let of_var x = Var x

  let apply_unop unop v = UnOp (unop, v)
  let apply_binop binop v1 v2 = BinOp (binop, v1, v2)

  let int_op_of_rel = function
    | Lt -> (<)
    | Le -> (<=)
    | Gt -> (>)
    | Ge -> (>=)

  (* The functions below allows to build arithmetic expressions with their usual
     operators (add, sub, etc) and with a partial simplification. *)

  let rec neg = function
    | A a -> S.neg a
    | Int i -> Int (-i)
    | UnOp (Neg, v) -> v
    | BinOp (Add, v1, v2) -> add (neg v1) (neg v2)
    | BinOp (Sub, v1, v2) -> minus (neg v1) (neg v2)
    | v -> apply_unop Neg v

  and add v1 v2 = match v1, v2 with
    | Int i1, Int i2 -> Int (i1 + i2)
    | Int 0, v | v, Int 0 -> v
    | BinOp (Add, Int i1, v), Int i2
    | BinOp (Add, v, Int i1), Int i2
    | Int i1, BinOp (Add, Int i2, v)
    | Int i1, BinOp (Add, v, Int i2) -> add (Int (i1 + i2)) v
    | BinOp (Sub, Int i1, v), Int i2
    | Int i1, BinOp (Sub, Int i2, v) -> minus (Int (i1 + i2)) v
    | BinOp (Sub, v, Int i1), Int i2
    | Int i1, BinOp (Sub, v, Int i2) -> add (Int (i1 - i2)) v
    | A a, v -> S.addl a v
    | v, A a -> S.addr v a
    | _ -> apply_binop Add v1 v2

  and minus v1 v2 = match v1, v2 with
    | A a, v -> S.minusl a v
    | v, A a -> S.minusr v a
    | Int i1, Int i2 -> Int (i1 - i2)
    | _, Int 0 -> v1
    | Int 0, _ -> neg v2
    | _ -> apply_binop Sub v1 v2

  and mul v1 v2 = match v1, v2 with
    | A a, v -> S.mull a v
    | v, A a -> S.mulr v a
    | Int i1, Int i2 -> Int (i1 * i2)
    | Int 1, v | v, Int 1 -> v
    | Int 0, _ | _, Int 0 -> Int 0
    | Int (-1), v | v, Int (-1) -> neg v
    | _ -> apply_binop Mul v1 v2

  and div v1 v2 = match v1, v2 with
    | _, Int 0 -> raise (Invalid_argument "Arith.div")
    | A a, v -> S.divl a v
    | v, A a -> S.divr v a
    | Int i1, Int i2 -> Int (i1 / i2)
    | Int 1, v | v, Int 1 -> v
    | Int 0, _ -> Int 0
    | Int (-1), v | v, Int (-1) -> neg v
    | _ -> apply_binop Div v1 v2

  and modulo v1 v2 = match v1, v2 with
    | _, Int 0 -> raise (Invalid_argument "Arith.modulo")
    | A a, v -> S.modl a v
    | v, A a -> S.modr v a
    | Int i1, Int i2 -> Int (i1 mod i2)
    | Int 0, _ | _, Int 1 -> Int 0
    | _ -> apply_binop Mod v1 v2

  and max v1 v2 = match v1, v2 with
    | A a, v -> S.maxl a v
    | v, A a -> S.maxr v a
    | Int i1, Int i2 -> Int (Pervasives.max i1 i2)
    | _ -> apply_binop Max v1 v2

  and cond v1 rel v2 v3 v4 = match v1, v2 with
    | Int i1, Int i2 when (int_op_of_rel rel) i1 i2 -> v3
    | Int _, Int _ -> v4
    | _ -> Cond (v1, rel, v2, v3, v4)

  let cmp is_large cmpal cmpar int_cmp v1 v2 = match v1, v2 with
    | Int i1, Int i2 -> int_cmp i1 i2
    | A a, v -> cmpal a v
    | v, A a -> cmpar v a
    | _ when is_large -> v1 = v2
    | _ -> false

  let binop_false _ _ = false

  let lt = cmp false S.ltl S.ltr (<)
  let le = cmp true S.lel S.ler (<=)
  let gt = cmp false S.gtl S.gtr (>)
  let ge = cmp true S.gel S.ger (>=)

  let op_of_unop = function
    | Neg -> neg

  let op_of_binop = function
    | Add -> add
    | Sub -> minus
    | Mul -> mul
    | Div -> div
    | Mod -> modulo
    | Max -> max

  let op_of_relation = function
    | Lt -> lt
    | Le -> le
    | Gt -> gt
    | Ge -> ge

  let bool_of_cond = op_of_relation

  let f_compute v subs_res = match v, subs_res with
    | A a, _ -> A (S.compute a)
    | UnOp (unop, _), v :: _ -> (op_of_unop unop) v
    | BinOp (binop, _, _), v1 :: v2 :: _ -> (op_of_binop binop) v1 v2
    | Cond (_, rel, _, _, _), v1 :: v2 :: v3 :: _
      when (op_of_relation rel) v1 v2 -> v3
    | Cond (_, rel, _, _, _), v1 :: v2 :: _ :: v4 :: _
      when (op_of_relation (opposite rel)) v1 v2 -> v4
    | _ -> fill_subs v subs_res

  (** [compute v] partially simplifies the arithmetic expression [v]. *)

  let compute v = fold f_compute v

  let abs v = cond (Int 0) Le v v (neg v)

  let f_replace_vars replacements v subs_res = match fill_subs v subs_res with
    | Var x when Misc.String.Map.mem x replacements ->
      Misc.String.Map.find x replacements
    | v -> v

  let replace_vars replacements = fold (f_replace_vars replacements)

end