open Ast
open Util
open Ast_util

(* AST for logical formulas *)

(*
  Two representations are used for the formulas :
  - Formula tree, representing the boolean formula directly
    as would be expected
  - "conslist", where all the numerical and enumerate constraints
    considered in a tree of ANDs are regrouped in a big list.
  The second form is used when applying a formula to an abstract value,
  whereas the first form is the one generated by the program transformation
  procedure and is the one we prefer to show the user because it is more
  easily understandable.
*)

(* ------------------------------------ *)
(*         First representation         *)
(* ------------------------------------ *)

(* Numerical part *)

(* bool on numerical operators : is it float ? *)

type num_expr =
  (* constants *)
  | NIntConst of int
  | NRealConst of float
  (* operators *)
  | NBinary of binary_op * num_expr * num_expr * bool
  | NUnary of unary_op * num_expr * bool
  (* identifier *)
  | NIdent of id

type num_cons_op =
  | CONS_EQ | CONS_NE
  | CONS_GT | CONS_GE
type num_cons = num_expr * num_cons_op  (* always imply right member = 0 *)


(* Logical part *)

(* Enumerated types *)
type enum_expr =
  | EIdent of id
  | EItem of string

type enum_op =
  | E_EQ | E_NE

type enum_cons = enum_op * id * enum_expr

type bool_expr =
  (* constants *)
  | BConst of bool
  (* operators from numeric values to boolean values *)
  | BRel of binary_rel_op * num_expr * num_expr * bool
  (* operators on enumerated types *)
  | BEnumCons of enum_cons
  (* boolean operators *)
  | BAnd of bool_expr * bool_expr
  | BOr of bool_expr * bool_expr
  | BTernary of bool_expr * bool_expr * bool_expr
      (* (A and B) or ((not A) and C) *)
  | BNot of bool_expr

(*
  Helper functions to construct formula trees and automatically
  simplifying some tautologies/contradictions
*)

let is_false = function
  | BConst false | BNot(BConst true) -> true
  | _ -> false
let is_true = function
  | BConst true | BNot(BConst false) -> true
  | _ -> false

let f_and a b =
  if is_false a || is_false b then BConst false
  else if is_true a then b
  else if is_true b then a
  else BAnd(a, b)

let f_and_list = List.fold_left f_and (BConst true)

let f_or a b =
  if is_true a || is_true b then BConst true
  else if is_false a then b
  else if is_false b then a
  else BOr(a, b)

let f_ternary c a b =
  if is_true c then a
  else if is_false c then b
  else if is_true a && is_false b then c
  else if is_true b && is_false a then BNot c
  else BTernary(c, a, b)

let f_e_op op a b = match a, b with
  | EItem i, EItem j -> BConst (if op = E_EQ then i = j else i <> j)
  | EIdent x, v | v, EIdent x -> BEnumCons(op, x, v)
let f_e_eq = f_e_op E_EQ

(*
  Transformation functions so that all the logical formula are written
  without using the NOT operator
*)

let rec eliminate_not = function
  | BNot e -> eliminate_not_negate e
  | BAnd(a, b) -> BAnd(eliminate_not a, eliminate_not b)
  | BOr(a, b) -> BOr(eliminate_not a, eliminate_not b)
  | BTernary(c, a, b) ->
    BTernary(eliminate_not c, eliminate_not a, eliminate_not b)
  | x -> x
and eliminate_not_negate = function
  | BConst x -> BConst(not x)
  | BEnumCons(op, a, b) -> BEnumCons((if op = E_EQ then E_NE else E_EQ), a, b)
  | BNot e -> eliminate_not e
  | BRel(r, a, b, is_real) ->
    if r = AST_EQ then
      BOr(BRel(AST_LT, a, b, is_real), BRel(AST_GT, a, b, is_real))
    else
      let r' = match r with
      | AST_EQ -> AST_NE
      | AST_NE -> AST_EQ
      | AST_LT -> AST_GE
      | AST_LE -> AST_GT
      | AST_GT -> AST_LE
      | AST_GE -> AST_LT
      in
      BRel(r', a, b, is_real)
  | BTernary (c, a, b) -> 
    eliminate_not_negate(BOr(BAnd(c, a), BAnd(BNot c, b)))
  | BAnd(a, b) ->
    BOr(eliminate_not_negate a, eliminate_not_negate b)
  | BOr(a, b) ->
    BAnd(eliminate_not_negate a, eliminate_not_negate b)
 


(* ------------------------------------ *)
(*    Simplification functions on the   *)
(*         first representation         *)
(* ------------------------------------ *)

(*
  Extract enumerated and numeric constants that
  are always true in the formula (ie do not appear
  under an OR or a ternary).
  (not mathematically exact : only a subset is returned)
*)

let rec get_root_true = function
  | BAnd(a, b) ->
    get_root_true a @ get_root_true b
  | BEnumCons e -> [BEnumCons e]
  | BRel (a, b, c, d) -> [BRel (a, b, c, d)]
  | _ -> []

(*
  Simplify formula considering something is true
  (only does some clearly visible simplifications ;
  later simplifications are implicitly done when
  abstract-interpreting the formula)
*)

let rec simplify true_eqs e =
  if List.exists
      (fun f ->
        try eliminate_not e = eliminate_not f with _ -> false)
      true_eqs
    then BConst true
  else if List.exists
      (fun f ->
        match e, f with
        | BEnumCons(E_EQ, a, EItem v), BEnumCons(E_EQ, b, EItem w)
          when a = b && v <> w -> true
        | _ -> try eliminate_not e = eliminate_not_negate f with _ -> false)
      true_eqs
    then BConst false
  else
    match e with
    | BAnd(a, b) ->
      f_and
        (simplify true_eqs a)
        (simplify true_eqs b)
    | BOr(a, b) ->
      f_or
        (simplify true_eqs a)
        (simplify true_eqs b)
    | BTernary(c, a, b) ->
      f_ternary
        (simplify true_eqs c)
        (simplify true_eqs a)
        (simplify true_eqs b)
    | BNot(n) ->
      begin match simplify true_eqs n with
      | BConst e -> BConst (not e)
      | x -> BNot x
      end
    | v -> v
    
let rec simplify_k true_eqs e =
  List.fold_left f_and (simplify true_eqs e) true_eqs
  

(*
  Simplify a formula replacing a variable with another
*)

let rec formula_replace_evars repls e =
  let new_name x =
    try List.assoc x repls
    with Not_found -> x
  in
  match e with
  | BOr(a, b) ->
    f_or (formula_replace_evars repls a) (formula_replace_evars repls b)
  | BAnd(a, b) ->
    f_and (formula_replace_evars repls a) (formula_replace_evars repls b)
  | BTernary(c, a, b) ->
    f_ternary
      (formula_replace_evars repls c)
      (formula_replace_evars repls a)
      (formula_replace_evars repls b)
  | BNot(n) ->
    begin match formula_replace_evars repls n with
    | BConst e -> BConst (not e)
    | x -> BNot x
    end
  | BEnumCons (op, a, EIdent b) ->
    let a', b' = new_name a, new_name b in
    if a' = b' then
      match op with | E_EQ -> BConst true | E_NE -> BConst false
    else
      BEnumCons(op, a', EIdent b')
  | BEnumCons (op, a, EItem x) ->
    BEnumCons (op, new_name a, EItem x)
  | x -> x



(*
  Extract names of enumerated variables referenced in formula
*)

let rec refd_evars_of_f = function
  | BAnd (a, b) | BOr(a, b) -> refd_evars_of_f a @ refd_evars_of_f b
  | BNot a -> refd_evars_of_f a
  | BTernary(a, b, c) -> refd_evars_of_f a @ refd_evars_of_f b @ refd_evars_of_f c
  | BEnumCons(_, e, EItem _) -> [e]
  | BEnumCons(_, e, EIdent f) -> [e; f]
  | _ -> []



(* ------------------------------------ *)
(*         Second representation        *)
(* ------------------------------------ *)

(*
  In big ANDs, try to separate levels of /\ and levels of \/
  We also use this step to simplify trues and falses that may be present.
*)

type conslist = enum_cons list * num_cons list * conslist_bool_expr
and conslist_bool_expr =
  | CLTrue
  | CLFalse
  | CLAnd of conslist_bool_expr * conslist_bool_expr
  | CLOr of conslist * conslist

let rec conslist_of_f = function
  | BNot e -> conslist_of_f (eliminate_not_negate e)
  | BRel (op, a, b, is_real) ->
    let x, y, op = match op with
      | AST_EQ -> a, b, CONS_EQ
      | AST_NE -> a, b, CONS_NE
      | AST_GT -> a, b, CONS_GT
      | AST_GE -> a, b, CONS_GE
      | AST_LT -> b, a, CONS_GT
      | AST_LE -> b, a, CONS_GE
    in
    let cons = if y = NIntConst 0 || y = NRealConst 0.
      then (x, op)
      else (NBinary(AST_MINUS, x, y, is_real), op)
    in [], [cons], CLTrue
  | BConst x ->
    [], [], if x then CLTrue else CLFalse
  | BEnumCons e ->
    [e], [], CLTrue
  | BOr(a, b) ->
    let eca, ca, ra = conslist_of_f a in
    let ecb, cb, rb = conslist_of_f b in
    begin match eca, ca, ra, ecb, cb, rb with
      | _, _, CLFalse, _, _, _ -> ecb, cb, rb
      | _, _, _, _, _, CLFalse -> eca, ca, ra
      | [], [], CLTrue, _, _, _ -> [], [], CLTrue
      | _, _, _, [], [], CLTrue -> [], [], CLTrue
      | _ -> [], [], CLOr((eca, ca, ra), (ecb, cb, rb))
    end
  | BAnd(a, b) ->
    let eca, ca, ra = conslist_of_f a in
    let ecb, cb, rb = conslist_of_f b in
    let cons = ca @ cb in
    let econs = eca @ ecb in
    begin match ra, rb with
      | CLFalse, _ | _, CLFalse -> [], [], CLFalse
      | CLTrue, _ -> econs, cons, rb
      | _, CLTrue -> econs, cons, ra
      | _, _ -> econs, cons, CLAnd(ra, rb)
    end
  | BTernary(c, a, b) ->
    conslist_of_f (BOr (BAnd(c, a), BAnd(BNot(c), b)))