(* The environment keeping track of possible and impossible values for
   variables *)

open Ast

module IntSet = Set.Make(struct type t = int let compare = compare end)
module StringMap = Map.Make(struct type t = string let compare = compare end)

let list_to_set l =
  List.fold_right IntSet.add l IntSet.empty

type info = Pos of IntSet.t | Neg of IntSet.t

type t = info StringMap.t

(* A variable associated with Pos l can take the values given in l
   and no other.
   A variable associated with Neg l cannot take the values given in l.
   In particular, a variable associated with Neg empty can take any value. *)

let all_variables = ref []
let all_states = ref []

let empty() =
  List.fold_right (fun v env -> StringMap.add v (Neg IntSet.empty) env)
                  !all_variables StringMap.empty

let init() =
  StringMap.add "state" (Pos (list_to_set !all_states)) (empty())

type equality_result = Yes | No | Maybe

let equals env v n =
  match StringMap.find v env with
      Pos l -> if IntSet.mem n l
               then if IntSet.cardinal l = 1 then Yes else Maybe
               else No
    | Neg l -> if IntSet.mem n l then No else Maybe

let matches_arm env v lbls =
    match StringMap.find v env with
      Pos l -> IntSet.subset l (list_to_set lbls)
    | _ -> false

let refine_arm env v lbls =
    match StringMap.find v env with
      Pos l ->
        (* keep only the labels that are in l *)
        List.filter (fun n -> IntSet.mem n l) lbls
    | Neg l ->
        (* remove all labels that are in l *)
        List.filter (fun n -> not(IntSet.mem n l)) lbls

let exclude_arm env v lbls =
    match StringMap.find v env with
      Pos l ->
        (* remove all values that are in lbls *)
        let l' = List.fold_right IntSet.remove lbls l in
        StringMap.add v (Pos l') env 
    | Neg l ->
        (* add all values that are in lbls *)
        let l' = List.fold_right IntSet.add lbls l in
        StringMap.add v (Neg l') env 

let exclude_arms env v armlist =
    List.fold_right (fun (lbls, stmt) env -> exclude_arm env v lbls)
                    armlist env

let set_value env v n =
  StringMap.add v (Pos (IntSet.singleton n)) env

let exclude_value env v n =
    match StringMap.find v env with
      Pos l ->
        let l' = IntSet.remove n l in
        StringMap.add v (Pos l') env
    | Neg l ->
        StringMap.add v (Neg (IntSet.add n l)) env

let set_possible_values env v lbls =
    match StringMap.find v env with
      Pos l ->
        (* take the intersection of l and lbls *)
        let l' = IntSet.inter l (list_to_set lbls) in
        StringMap.add v (Pos l') env
    | Neg l ->
        (* remove from lbls all values in l and set a positive info on v *)
        let l' = IntSet.diff (list_to_set lbls) l in
        StringMap.add v (Pos l') env

let exclude_values env v lbls =
    match StringMap.find v env with
      Pos l ->
        (* remove from l all values from lbls *)
        let l' = IntSet.diff l (list_to_set lbls) in
        StringMap.add v (Pos l') env
    | Neg l ->
        (* add to l all values in lbls *)
        StringMap.add v (Neg (IntSet.union l (list_to_set lbls))) env

let possible_values env v lbls =
    match StringMap.find v env with
      Pos l ->
        (* return the intersection of l and lbls *)
        List.filter (fun n -> IntSet.mem n l) lbls
    | Neg l ->
        (* remove from lbls all values in l *)
        List.filter (fun n -> not(IntSet.mem n l)) lbls

let value_of env v =
    match StringMap.find v env with
      Pos l ->
        begin match IntSet.elements l with [n] -> n | _ -> raise Not_found end
    | Neg l ->
        raise Not_found

let conj_info i1 i2 =
  match (i1, i2) with
    Pos l1, Pos l2 -> Pos(IntSet.inter l1 l2)
  | Neg l1, Neg l2 -> Neg(IntSet.union l1 l2)
  | Pos l1, Neg l2 -> Pos(IntSet.diff l1 l2)
  | Neg l1, Pos l2 -> Pos(IntSet.diff l2 l1)

let disj_info i1 i2 =
  match (i1, i2) with
    Pos l1, Pos l2 -> Pos(IntSet.union l1 l2)
  | Neg l1, Neg l2 -> Neg(IntSet.inter l1 l2)
  | _, _ -> Neg IntSet.empty (* suboptimal, but safe *)

let conj e1 e2 =
  List.fold_right
    (fun v env ->
      StringMap.add v (conj_info (StringMap.find v e1) (StringMap.find v e2))
                    env)
    !all_variables StringMap.empty

let disj e1 e2 =
  List.fold_right
    (fun v env ->
      StringMap.add v (disj_info (StringMap.find v e1) (StringMap.find v e2))
                    env)
    !all_variables StringMap.empty

let rec refine_cond env = function
    Cequals(v, n) ->
      (set_value env v n, exclude_value env v n)
  | Cor [] ->
      (env, env) (* should not happen *)
  | Cor [c1] ->
      refine_cond env c1
  | Cor (c1::cl) ->
      let (iftrue1, iffalse1) = refine_cond env c1 in
      let (iftruel, iffalsel) = refine_cond env (Cor cl) in
      (disj iftrue1 iftruel, conj iffalse1 iffalsel)
  | Cand [] ->
      (env, env) (* should not happen *)
  | Cand [c1] ->
      refine_cond env c1
  | Cand (c1::cl) ->
      let (iftrue1, iffalse1) = refine_cond env c1 in
      let (iftruel, iffalsel) = refine_cond env (Cand cl) in
      (conj iftrue1 iftruel, disj iffalse1 iffalsel)

open Format

let pretty_infos v i =
  print_string v; print_string ": ";
  begin match i with
    Pos l ->
      begin match IntSet.elements l with
        [] -> print_string "impossible"
      | elts -> print_string "one of ";
                List.iter (fun n -> print_int n; print_space()) elts
      end
  | Neg l ->
      begin match IntSet.elements l with
        [] -> print_string "any value"
      | elts -> print_string "none of ";
                List.iter (fun n -> print_int n; print_space()) elts
      end
  end;
  print_newline()

let pretty env =
  StringMap.iter pretty_infos env