open Abstract_syntax_tree
open Util

open Value_domain
open Environment_domain


module NonRelational (V : VALUE_DOMAIN) : ENVIRONMENT_DOMAIN = struct
    type env = V.t VarMap.t

    type t = Env of env | Bot

    let init = Env VarMap.empty
    let bottom = Bot

    let get_var env id =
        try VarMap.find id env
        with Not_found -> V.top
    
    (* utilities *)

    let rec eval env expr =
        begin match fst expr with
        | AST_identifier (id, _) -> get_var env id
        | AST_int_const (s, _) -> V.const (Z.of_string s)
        | AST_int_rand ((s, _), (t, _)) -> V.rand (Z.of_string s) (Z.of_string t)
        | AST_unary (AST_UNARY_PLUS, e) -> eval env e
        | AST_unary (AST_UNARY_MINUS, e) -> V.neg (eval env e)
        | AST_unary (_, e) -> V.bottom
        | AST_binary (op, e1, e2) ->
            let f = match op with
            | AST_PLUS -> V.add
            | AST_MINUS -> V.sub
            | AST_MULTIPLY -> V.mul
            | AST_DIVIDE -> V.div
            | AST_MODULO -> V.rem
            | _ -> fun _ _ -> V.bottom
            in f (eval env e1) (eval env e2)
        | _ -> assert false (* unimplemented extension *)
        end

    let strict env =                    (*      env -> t        *)
        if VarMap.exists (fun _ x -> x = V.bottom) env
            then Bot
            else Env env
    let strict_apply f = function       (* (env -> t) -> t -> t *)
        | Bot -> Bot
        | Env e -> match f e with
            | Bot -> Bot
            | Env e -> strict e

    (* implementation *)

    let is_bot env = match env with
        | Bot -> true
        | Env env -> strict env = Bot

    let addvar x id = strict_apply (fun y -> Env (VarMap.add id V.top y)) x
    let rmvar x id = strict_apply (fun y -> Env (VarMap.remove id y)) x

    let assign x id expr = strict_apply
        (fun env -> Env (VarMap.add id (eval env expr) env))
        x
    let compare f x e1 e2 =
        strict_apply (fun env -> match fst e1, fst e2 with
                | AST_identifier(a, _), AST_identifier(b, _) ->
                    let v1, v2 = get_var env a, get_var env b in
                    let v1', v2' = f v1 v2 in
                    Env (VarMap.add a v1' (VarMap.add b v2' env))
                | AST_identifier(a, _), _ ->
                    let v1, v2 = get_var env a, eval env e2 in
                    let v1', v2' = f v1 v2 in
                    if V.bottom = v2'
                        then Bot
                        else Env (VarMap.add a v1' env)
                | _, AST_identifier(b, _) ->
                    let v1, v2 = eval env e1, get_var env b in
                    let v1', v2' = f v1 v2 in
                    if V.bottom = v1'
                        then Bot
                        else Env (VarMap.add b v2' env)
                | _ ->
                    let v1, v2 = eval env e1, eval env e2 in
                    let v1', v2' = f v1 v2 in
                    if V.bottom = v1' || V.bottom = v2'
                        then Bot
                        else Env env)
            x
    let compare_leq = compare V.leq
    let compare_eq = compare (fun x y -> let r = V.meet x y in r, r)

    let join a b = match a, b with
        | Bot, x | x, Bot -> x
        | Env m, Env n ->
            strict (VarMap.map2z (fun _ a b -> V.join a b) m n)
    
    let meet a b = match a, b with
        | Bot, _ | _, Bot -> Bot
        | Env m, Env n ->
            strict (VarMap.map2z (fun _ a b -> V.meet a b) m n)

    let widen a b = match a, b with
        | Bot, x | x, Bot -> x
        | Env m, Env n ->
            strict (VarMap.map2z (fun _ a b -> V.widen a b) m n)
    
    let subset a b = match a, b with
        | Bot, x -> true
        | Env _, Bot -> false
        | Env m, Env n ->
            VarMap.for_all2o
                (fun _ _ -> true)
                (fun _ v -> v = V.top)
                (fun _ a b -> V.subset a b)
                m n

    let eq a b = match a, b with
        | Bot, Bot -> true
        | Env m, Env n ->
            VarMap.for_all2o
                (fun _ _ -> false)
                (fun _ _ -> false)
                (fun _ a b -> a = b)
                m n
        | _ -> false
        
    (* pretty-printing *)
    let var_str env vars =
        match env with
        | Bot -> "bottom"
        | Env env ->
            let v = List.fold_left
                (fun s id ->
                    (if s = "" then s else s ^ ", ") ^
                        id ^ " in " ^ (V.to_string (get_var env id))
                    )
                ""
                vars
            in
                "[ " ^ v ^ " ]"
    
    let vars env = match env with
        | Bot -> []
        | Env env -> List.map fst (VarMap.bindings env)
end