open Naked
open Abstract_syntax_tree

module Congruence (*: NAKED_VALUE_DOMAIN*) = struct
  type t = { multiple : Z.t; offset : Z.t }

  let all_numbers = { multiple = Z.one; offset = Z.zero }
  let const a = { multiple = Z.zero; offset = a }
  let rand a b = if Z.equal a b then const a else all_numbers
  let minus z = { multiple = z.multiple; offset = Z.neg z.offset }
  let is_only_zero z = Z.equal z.multiple Z.zero && Z.equal z.offset Z.zero
  let has_zero z = Z.divisible z.offset z.multiple

  let subset z1 z2 =
    Z.equal z2.multiple Z.zero && Z.equal z1.multiple Z.zero
    && Z.equal z1.offset z2.offset
    || (not (Z.equal z2.multiple Z.zero))
       && (Z.equal z1.multiple Z.zero || Z.divisible z2.multiple z1.multiple)
       && Z.equal (Z.rem z1.offset z2.multiple) (Z.rem z2.offset z2.multiple)

  let join z1 z2 =
    {
      multiple =
        Z.gcd
          (Z.gcd z1.multiple z2.multiple)
          (Z.abs (Z.sub z1.offset z2.offset));
      offset = z1.offset;
    }

  let meet z1 z2 =
    if Z.equal z1.multiple Z.zero then if subset z1 z2 then z1 else raise Absurd
    else if Z.equal z2.multiple Z.zero then
      if subset z2 z1 then z2 else raise Absurd
    else
      let lcm = Z.lcm z1.multiple z2.multiple in
      let gcd = Z.gcd z1.multiple z2.multiple in
      if Z.equal (Z.rem z1.offset gcd) (Z.rem z2.offset gcd) then
        { multiple = lcm; offset = Z.rem z1.offset gcd }
      else raise Absurd

  let multiples_of z =
    if is_only_zero z then z
    else if
      (not (Z.equal z.multiple Z.zero))
      && Z.equal (Z.rem z.offset z.multiple) Z.zero
    then z
    else all_numbers

  let divisors_of z = z
  let remainders z = z
  let convex_sym z = z

  let binary z1 z2 = function
    | AST_PLUS ->
        {
          multiple = Z.gcd z1.multiple z2.multiple;
          offset = Z.add z1.offset z2.offset;
        }
    | AST_MINUS ->
        {
          multiple = Z.gcd z1.multiple z2.multiple;
          offset = Z.sub z1.offset z2.offset;
        }
    | AST_MULTIPLY ->
        {
          multiple =
            Z.gcd
              (Z.gcd
                 (Z.mul z1.multiple z2.multiple)
                 (Z.mul z1.multiple z2.offset))
              (Z.mul z1.offset z2.multiple);
          offset = Z.mul z1.offset z2.offset;
        }
    | AST_DIVIDE ->
        if is_only_zero z2 then raise Absurd
        else if
          Z.equal z2.multiple Z.zero
          && Z.divisible z1.offset z2.offset
          && Z.divisible z1.multiple z2.offset
        then
          {
            multiple = Z.div z1.multiple (Z.abs z2.offset);
            offset = Z.div z1.offset z2.offset;
          }
        else all_numbers
    | AST_MODULO -> if (Z.equal z2.multiple Z.zero) && (Z.divisible z1.multiple z2.offset) then
				const (Z.rem z1.offset z2.offset)
		    else all_numbers

  let compatible z = function AST_EQUAL -> z | _ -> raise NeedTop

  let rec compare z1 z2 = function
    | AST_EQUAL ->
        let r = meet z1 z2 in
        (r, r)
    | AST_NOT_EQUAL ->
        if
          Z.equal z1.multiple Z.zero && Z.equal z2.multiple Z.zero
          && Z.equal z1.offset z2.offset
        then raise Absurd
        else (z1, z2)
    | AST_LESS ->
        if
          Z.equal z1.multiple Z.zero && Z.equal z2.multiple Z.zero
          && Z.geq z1.offset z2.offset
        then raise Absurd
        else (z1, z2)
    | AST_LESS_EQUAL ->
        if
          Z.equal z1.multiple Z.zero && Z.equal z2.multiple Z.zero
          && Z.gt z1.offset z2.offset
        then raise Absurd
        else (z1, z2)
    | AST_GREATER ->
        let r1, r2 = compare z2 z1 AST_LESS in
        (r2, r1)
    | AST_GREATER_EQUAL ->
        let r1, r2 = compare z2 z1 AST_LESS_EQUAL in
        (r2, r1)

  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 ->
        if is_only_zero z1 || is_only_zero z2 then
          if has_zero r then (z1, z2) else raise Absurd
        else
          let z2' = meet z2 (binary r z1 AST_DIVIDE) in
          let z1' = meet z1 (binary r z2 AST_DIVIDE) in
          if has_zero r then
            ( (if has_zero z1' then join z1' (const Z.zero) else z1'),
              if has_zero z2' then join z2' (const Z.zero) else z2' )
          else (z1', z2')
    | AST_DIVIDE ->
        (meet z1 (binary r z2 AST_MULTIPLY), meet z2 (binary z1 r AST_DIVIDE))
    | AST_MODULO -> (z1, z2) (* If z2 is a constant this is somewhat doable with the Chinese Remainder Theorem *)

  let widen = join
  let narrow = meet

  let print fmt z =
    Z.pp_print fmt z.multiple;
    Format.pp_print_string fmt "Z + ";
    Z.pp_print fmt z.offset
end