open Naked
open Abstract_syntax_tree

module Interval : NAKED_VALUE_DOMAIN = struct
  type t = { lower : Z.t; upper : Z.t }

  let const a = { lower = a; upper = a }
  let rand a b = { lower = Z.min a b; upper = Z.max a b }
  let infty = Z.of_int64_unsigned Int64.max_int
  let min_infty = Z.neg infty

  let rand3 a b c =
    let mi = List.fold_left Z.min a [ a; b; c ] in
    let ma = List.fold_left Z.max a [ a; b; c ] in
    rand mi ma

  let rand4 a b c d =
    let mi = List.fold_left Z.min a [ a; b; c; d ] in
    let ma = List.fold_left Z.max a [ a; b; c; d ] in
    rand mi ma

  let minus z = rand (Z.neg z.upper) (Z.neg z.lower)
  
  let join z1 z2 = rand4 z1.lower z2.lower z1.upper z2.upper
  let meet z1 z2 =
    if Z.leq z1.upper z2.lower || Z.leq z2.upper z1.lower then raise Absurd
    else rand (Z.max z1.lower z2.lower) (Z.min z1.upper z2.upper)
  let widen z1 z2 =
    rand
      (if Z.leq z1.lower z2.lower then z1.lower else min_infty)
      (if Z.geq z1.upper z2.upper then z1.upper else infty)
  let narrow = meet
  let subset z1 z2 = Z.geq z1.lower z2.lower && Z.leq z1.upper z2.upper


  let binary z1 z2 = function
    | AST_PLUS -> rand (Z.add z1.lower z2.lower) (Z.add z1.upper z2.upper)
    | AST_MINUS -> rand (Z.sub z1.lower z2.upper) (Z.sub z1.upper z2.lower)
    | AST_MULTIPLY ->
        rand4 (Z.mul z1.lower z2.lower) (Z.mul z1.upper z2.lower)
          (Z.mul z1.lower z2.upper) (Z.mul z1.upper z2.upper)
    | AST_DIVIDE ->
        if Z.sign z2.lower <> Z.sign z2.upper || Z.sign z2.lower = 0 then
          raise Absurd
        else
          rand4 (Z.div z1.lower z2.lower) (Z.div z1.upper z2.lower)
            (Z.div z1.lower z2.upper) (Z.div z1.upper z2.upper)
    | AST_MODULO ->
        if Z.sign z2.lower <> Z.sign z2.upper || Z.sign z2.lower = 0 then raise Absurd
        else
          rand4 (Z.rem z1.lower z2.lower) (Z.rem z1.upper z2.lower)
            (Z.rem z1.lower z2.upper) (Z.rem z1.upper z2.upper)

  let is_only_zero z = Z.equal z.lower Z.zero && Z.equal z.upper Z.zero
  let multiples_of z = if is_only_zero z then const Z.zero else raise NeedTop
  let divisors_of z = if is_only_zero z then const Z.zero else raise NeedTop
  let remainders z = if is_only_zero z then const Z.zero else raise NeedTop
  let convex_sym z = if is_only_zero z then const Z.zero else raise NeedTop

  let compatible z = function
    | AST_EQUAL -> z
    | AST_NOT_EQUAL -> raise NeedTop
    | AST_LESS ->
        if Z.equal z.lower min_infty then raise Absurd
        else rand min_infty z.lower
    | AST_LESS_EQUAL -> rand min_infty z.lower
    | AST_GREATER ->
        if Z.equal z.upper infty then raise Absurd else rand z.upper infty
    | AST_GREATER_EQUAL -> rand z.upper infty

  let rec compare z1 z2 = function
    | AST_EQUAL ->
        let z = join z1 z2 in
        (z, z)
    | AST_NOT_EQUAL ->
        if
          Z.equal z1.lower z1.upper && Z.equal z1.upper z2.lower
          && Z.equal z2.lower z2.upper
        then raise Absurd
        else (z1, z2)
    | AST_LESS ->
        if Z.leq z2.upper z1.lower then raise Absurd
        else
          let r1 =
            rand z1.lower
              (if Z.geq z1.upper z2.upper then Z.pred z2.upper else z1.upper)
          in
          let r2 =
            rand
              (if Z.geq z1.lower z2.lower then Z.succ z2.lower else z1.lower)
              z2.upper
          in
          (r1, r2)
    | AST_LESS_EQUAL ->
        if Z.lt z2.upper z1.lower then raise Absurd
        else
          let r1 = rand z1.lower (Z.min z1.upper z2.upper) in
          let r2 = rand (Z.max z1.lower z2.lower) z2.upper in
          (r1, r2)
    | AST_GREATER -> compare z2 z1 AST_LESS
    | AST_GREATER_EQUAL -> compare z2 z1 AST_GREATER_EQUAL

  let bwd_binary z1 z2 op r =
    match op with
    | AST_PLUS ->
        (meet z1 (binary r z2 AST_MINUS), meet z2 (binary r z1 AST_MINUS))
    | AST_MINUS ->
        (meet z1 (binary r z2 AST_PLUS), meet z2 (binary z1 r AST_MINUS))
    | AST_MULTIPLY ->
        (meet z1 (binary r z2 AST_DIVIDE), meet z2 (binary r z1 AST_DIVIDE))
    | AST_DIVIDE ->
        (meet z1 (binary r z2 AST_MULTIPLY), meet z2 (binary z1 r AST_DIVIDE))
    | AST_MODULO -> (z1, z2)


  let print fmt z =
    Format.pp_print_string fmt "[";
    Z.pp_print fmt z.lower;
    Format.pp_print_string fmt "; ";
    Z.pp_print fmt z.upper;
    Format.pp_print_string fmt "]"
end