
(** {1 Integer square root} *)

module Square

  use int.Int

  function sqr (x:int) : int = x * x

  lemma sqr_non_neg: forall x:int. sqr x >= 0

  lemma sqr_increasing:
    forall x y:int. 0 <= x <= y -> sqr x <= sqr y

  lemma sqr_sum :
    forall x y : int. sqr(x+y) = sqr x + 2*x*y + sqr y

  predicate isqrt_spec (x res:int) =
    res >= 0 /\ sqr res <= x < sqr (res + 1)
end

(** {2 Simple algorithm} *)

module Simple

  use int.Int
  use ref.Refint
  use Square

  let isqrt (x:int) : int
    requires { x >= 0 }
    ensures  { isqrt_spec x result }
  = let ref count = 0 in
    let ref sum = 1 in
    while sum <= x do
      invariant { count >= 0 }
      invariant { x >= sqr count }
      invariant { sum = sqr (count+1) }
      variant   { x - count }
      count += 1;
      sum += 2 * count + 1
    done;
    count

  let main ()
    ensures { result = 4 }
  = isqrt 17

end

(** {2 Another algorithm, in the style of Newton-Raphson} *)

module NewtonMethod

  use int.Int
  use mach.int.Int
  use ref.Ref
  use Square

  let sqrt (x : int) : int
    requires { x >= 0 }
    ensures  { isqrt_spec x result }
  = if x = 0 then 0 else
    if x <= 3 then 1 else
    let ref y = x in
    let ref z = (1 + x) / 2 in
    while z < y do
      variant { y }
      invariant { z > 0 }
      invariant { y > 0 }
      invariant { z = div (div x y + y) 2 }
      invariant { x < sqr (y + 1) }
      invariant { x < sqr (z + 1) }
      y <- z;
      z <- (x / z + z) / 2;
      (* A few hints to prove preservation of the last invariant *)
      assert { x < sqr (z + 1)
        by let a = div x y in
              x < a * y + y
           so a + y <= 2 * z + 1
           so sqr (a + y + 1) <= sqr (2 * z + 2)
           so 4 * (sqr (z + 1) - x)
             = sqr (2 * z + 2) - 4 * x
             >= sqr (a + y + 1) - 4 * x
             > sqr (a + y + 1) - 4 * (a * y + y)
             = sqr (a + 1 - y)
             >= 0 }
    done;
    assert { y * y <= div x y * y
             by y <= div x y };
    y

end
