open Ast
open Formula
open Util

exception Bot

(*
  Interface and implementation of an abstract domain meant to handle
  constraints on enumerated variables.

  Constraints are basically of two types :
  - var OP value
  - var OP var'
  with OP either == or !=

  (see formula.ml for a definition of the enum_cons type)
*)

module type ENUM_ENVIRONMENT_DOMAIN = sig
    type t

    type item = string

    (* construction *)
    val top           : (id * item list) list -> t
    val bottom        : (id * item list) list -> t
    val vtop          : t -> t  (* top with same vars *)
    val is_bot        : t -> bool
    val is_top        : t -> bool

    (* variable management *)
    val vars          : t -> (id * item list) list

    val forgetvar     : t -> id -> t
    val forgetvars    : t -> id list -> t
    val project       : t -> id -> item list

    (* set-theoretic operations *)
    val join          : t -> t -> t   (* union *)
    val meet          : t -> t -> t   (* intersection, may raise Bot *)

    val subset        : t -> t -> bool
    val eq            : t -> t -> bool

    (* abstract operations *)
    val apply_cons    : t -> enum_cons -> t list (* may disjunct *)
    val apply_cl      : t -> enum_cons list -> t (* join any disjunction ; may raise Bot *)
    val assign        : t -> (id * id) list -> t

    (* pretty-printing *)
    val print         : Format.formatter -> t -> unit
end


(*
    Simple domain : one valuation for each variable, or T
*)

module Valuation : ENUM_ENVIRONMENT_DOMAIN = struct
    type item = string

    type t = {
        vars      : (id * item list) list;
        value     : item VarMap.t;
    }

    let top vars = { vars; value = VarMap.empty }
    let vtop x = { x with value = VarMap.empty }
    let bottom vars = raise Bot  (* we don't represent Bot *)
    let is_bot x = false
    let is_top x = VarMap.is_empty x.value

    let vars x = x.vars

    let forgetvar x id =
        { x with value = VarMap.remove id x.value }
    let forgetvars = List.fold_left forgetvar
    let project x id =
        try [VarMap.find id x.value]
        with Not_found -> List.assoc id x.vars

    let join x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, Some b when a = b -> Some a
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }
    let meet x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, None | None, Some a -> Some a
            | Some a, Some b -> if a = b then Some a else raise Bot
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }

    let subset a b =
        VarMap.for_all
          (fun id v ->
              try v = VarMap.find id b.value
              with Not_found -> true)
          a.value
    let eq a b =
        VarMap.equal (=) a.value b.value

    let cc op = match op with
        | E_EQ -> (=)
        | E_NE -> (<>)

    let apply_cons x (op, id, e) =
        match e with
        | EItem s ->
          begin try let t = VarMap.find id x.value in
              if cc op s t then [x] else []
          with Not_found ->
            List.map (fun v -> { x with value = VarMap.add id v x.value })
              (List.filter (cc op s) (List.assoc id x.vars))
          end
        | EIdent id2 ->
          match
            (try Some (VarMap.find id x.value) with _  -> None),
            (try Some (VarMap.find id2 x.value) with _  -> None)
          with
          | None, None -> [x]
          | None, Some s ->
            List.map (fun v -> { x with value = VarMap.add id v x.value})
              (List.filter (cc op s) (List.assoc id x.vars))
          | Some s, None ->
            List.map (fun v -> { x with value = VarMap.add id2 v x.value})
              (List.filter (cc op s) (List.assoc id2 x.vars))
          | Some s, Some t ->
            if cc op s t then [x] else []

    let apply_cl x l =
      let f x =
        List.fold_left
          (fun k c -> match apply_cons k c with
            | [] -> raise Bot
            | p::q -> List.fold_left join p q )
          x l
      in
        fix eq f x
   
    let assign x idl =
        let v = List.fold_left
            (fun v (id, id2) ->
                try VarMap.add id (VarMap.find id2 x.value) v
                with Not_found -> VarMap.remove id v )
            x.value idl
        in { x with value = v }


    let print fmt x =
        let b = List.map (fun (x, y) -> y, x) (VarMap.bindings x.value) in
        let s = List.sort Pervasives.compare b in
        let rec bl = function
          | [] -> []
          | (v, id)::q ->
            begin match bl q with
            | (vv, ids)::q when vv = v -> (v, id::ids)::q
            | r -> (v, [id])::r
            end
        in
        let sbl = bl s in
        Format.fprintf fmt "@[<v 2>{ ";
        List.iteri
          (fun j (v, ids) ->
            if j > 0 then Format.fprintf fmt "@ ";
            Format.fprintf fmt "@[<hov 4>";
            List.iteri
              (fun i id ->
                if i > 0 then Format.fprintf fmt ",@ ";
                Format.fprintf fmt "%a" Formula_printer.print_id id)
              ids;
            Format.fprintf fmt " ≡ %s@]" v)
          sbl;
        Format.fprintf fmt " }@]"

end


(*
    Less simple domain : a set of possible values for each variable, or T
*)

module MultiValuation : ENUM_ENVIRONMENT_DOMAIN = struct
    type item = string

    type t = {
        vars      : (id * item list) list;
        value     : item list VarMap.t;
    }

    let sort = List.sort Pervasives.compare
    let uniq x = uniq_sorted (sort x)

    let top vars = { vars; value = VarMap.empty }
    let vtop x = { x with value = VarMap.empty }
    let bottom vars = raise Bot  (* we don't represent Bot *)
    let is_bot x = false
    let is_top x =
      VarMap.for_all
        (fun k v -> List.for_all (fun vv -> List.mem vv v)
            (List.assoc k x.vars))
        x.value

    let vars x = x.vars

    let forgetvar x id =
        { x with value = VarMap.remove id x.value }
    let forgetvars = List.fold_left forgetvar
    let project x id =
        try VarMap.find id x.value
        with Not_found -> List.assoc id x.vars

    let join x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, Some b -> Some (uniq (a@b))
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }
    let meet x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, None | None, Some a -> Some a
            | Some a, Some b ->
              begin match List.filter (fun x -> List.mem x b) a with
              | [] -> raise Bot | l -> Some l end
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }

    let subset a b =
        VarMap.for_all
          (fun id v ->
              try let v2 = VarMap.find id b.value in
                  List.for_all (fun x -> List.mem x v2) v
              with Not_found -> true)
          a.value
    let eq a b = subset a b && subset b a

    let apply_cons x (op, id, e) =
        let op = match op with | E_EQ -> (=) | E_NE -> (<>) in

        match e with
        | EItem s ->
          let pv = project x id in
          begin match List.filter (op s) pv with
          | [] -> []
          | vals -> [{ x with value = VarMap.add id vals x.value }]
          end
        | EIdent id2 ->
          let v1 = project x id in
          let v2 = project x id2 in
          List.fold_left
            (fun l v1 ->
                let v2 = List.filter (op v1) v2 in
                let x = { x with value = VarMap.add id [v1] x.value } in
                match v2 with
                | [] -> l
                | _ -> { x with value = VarMap.add id2 v2 x.value }::l)
            [] v1

    let apply_cl x l =
      let f x =
        List.fold_left
          (fun k c -> match apply_cons k c with
            | [] -> raise Bot
            | p::q -> List.fold_left join p q )
          x l
      in
        fix eq f x
    
    let assign x idl =
        let v = List.fold_left
            (fun v (id, id2) ->
                try VarMap.add id (VarMap.find id2 x.value) v
                with Not_found -> VarMap.remove id v )
            x.value idl
        in { x with value = v }


    let print fmt x =
        let b = List.map (fun (x, y) -> y, x) (VarMap.bindings x.value) in
        let s = List.sort Pervasives.compare b in
        let rec bl = function
          | [] -> []
          | (v, id)::q ->
            let v = sort v in
            if v <> sort (List.assoc id x.vars) then
              match bl q with
              | (vv, ids)::q when vv = v -> (v, id::ids)::q
              | r -> (v, [id])::r
            else
              bl q
        in
        let sbl = bl s in
        Format.fprintf fmt "@[<v 2>{ ";
        List.iteri
          (fun j (v, ids) ->
            if j > 0 then Format.fprintf fmt "@ ";
            Format.fprintf fmt "@[<hov 4>";
            List.iteri
              (fun i id ->
                if i > 0 then Format.fprintf fmt ",@ ";
                Format.fprintf fmt "%a" Formula_printer.print_id id)
              ids;
            match v with
            | [v0] -> Format.fprintf fmt " ≡ %s@]" v0
            | l -> Format.fprintf fmt " ∊ @[<hov>{ %a }@]@]" (print_list Format.pp_print_string ", ") l)
          sbl;
        Format.fprintf fmt " }@]"

end