// compile with
// atscc -o fact fact.dats -lgmp
(* In this file, we are to implement a verified factorial function
 * to serve as an introductory example to programming with theorem
 * proving in ATS.
 *
 * here we define a relation [fact(x,y)] between two natural numbers,
 * with the intended meaning "y = x!". we define [fact] by the use
 * of inductive judgements, and then transliterate them to the syntax of ATS.
 *
 * ----------(fact-bas)
 * fact(0, 1)
 *
 *   fact(x, y)
 * --------------(fact-ind)
 * fact(x+1, y*(x+1))
 *
 * the (fact-bas) rule defines that the factorial of 0 is 1.
 * the (fact-ind) rule states that if [y] is the factorial of [x],
 * then the factorial of [x+1] is [y*x].
 *
 * here is an example derivation of fact(3, 6):
 *
 * ----------(fact-bas)
 * fact(0, 1)
 * ----------(fact-ind)
 * fact(1, 1)
 * ----------(fact-ind)
 * fact(2, 2)
 * ----------(fact-ind)
 * fact(3, 6)
 *
 * due to technical issues it is not possible to express multiplication
 * of integers directly, so we have to use [MUL], which encodes
 * the multiplication relation.
 *)
staload "libats/SATS/intinf.sats"
dynload "libats/DATS/intinf.dats"

dataprop FACT (int, int) =
  | FACTbas (0, 1)
  | {a,b,ab:int | a > 0} FACTind (a, ab) of (FACT (a-1, b), MUL (a, b, ab))

(* we are now to write the first approximation to the solution.
 * note how we manually associate proofs with programs.
 *)
fun fac0 {a:nat} .<a>. (a: int a): [b:int] (FACT (a, b) | int b) =
  if a > 0 then let
    val (pf_fac | b) = fac0 (a-1)
    val (pf_mul | ab) = a imul2 b
  in (FACTind (pf_fac, pf_mul) | ab) end
  else begin (FACTbas () | 1) end

(* this solution is unsatisfactory, because it involves recursion where
 * iteration is more appropriate, so we rewrite it in a tail-recursive
 * fashion.
 *
 * it is also possible to state and prove two lemmas about relation [FACT]:
 * - for any natural number [n], there exists natural number [r] where [FACT (n, r)]
 * - for any natural numbers [n,r0,r1], if [FACT (n, r0)] and [FACT (n, r1)], then [r0 = r1]
 *)

prfun fact_istot {a:nat} .<a>. (): [b: int] FACT (a, b) =
  sif a > 0 then let
    prval [b:int] pf_fact = fact_istot {a-1} ()
    prval pf_mul = mul_istot {a, b} ()
  in FACTind (pf_fact, pf_mul) end
  else FACTbas

prfun fact_isfun {a,b1,b2:int | a >= 0} .<a>.
  (pf_a: FACT (a, b1), pf_b: FACT (a, b2)): [b1 == b2] void =
  case+ (pf_a, pf_b) of
  | (FACTbas (), FACTbas ()) => ()
  | (FACTind (pf_a, pf_mul_a), FACTind (pf_b, pf_mul_b)) => let
    prval () = fact_isfun (pf_a, pf_b)
    prval () = mul_isfun (pf_mul_a, pf_mul_b)
  in () end

(* to rewrite [fac0] in a tail-recursive style, some lemmas on multiplication
 * are stated (one without proof)
 *)
// we need some lemmas on multiplication
// MULlem00 = mul_istot
// 1 * x = x
prfn MULlem10 {x,y:int} (pf: MUL (1, x, y)): [x==y] void =
  case+ pf of MULind (MULbas ()) => ()

// x * 1 = x
prfn MULlem11 {x,y:int} (pf: MUL (x, 1, y)): [x==y] void = let
  prval pf' = MULlem10 (mul_commute pf)
in () end

// multiplication is associative: (xy)z = x(yz)
extern prval MULlem20 : {x,y,z,xy,yz,xyz:int}
  (MUL (x, y, xy), MUL (y, z, yz), MUL (xy, z, xyz)) -<prf> MUL (x, yz, xyz)

// [fac1] implements the factorial function in a tail-recursive style
// the gist of the proof is to show that loop (n,x) = x*n! from which
// loop (n,1) = n! follows
fn fac1 {n:nat} (n: int n): [r0: int] (FACT (n, r0) | int r0) = let
  // [loop] is tail-recursive
  fun loop {n,x:int | n >= 0} .<n>. (n: int n, x: int x)
    : [r,r0:int] (FACT (n, r), MUL (x, r, r0) | int r0) =
    if n > 0 then let
      val (pf_mul_x_n_xn | xn) = x imul2 n
      val (pf_fac_n1_r1, pf_mul_xn_r1_r0 | res) = loop (n-1, xn)
      prval pf_mul_n_r1_nr1 = mul_istot ()
      prval pf_mul_x_nr1_r0 = MULlem20 (pf_mul_x_n_xn, pf_mul_n_r1_nr1, pf_mul_xn_r1_r0)
      prval pf_fac_n_nr1 = FACTind (pf_fac_n1_r1, pf_mul_n_r1_nr1)
    in
      (pf_fac_n_nr1, pf_mul_x_nr1_r0 | res)
    end else let
      prval pf_mul_x_1_y = mul_istot {x,1} ()  // x * 1 = y
      prval () = MULlem11 (pf_mul_x_1_y)       // x = y
    in
      (FACTbas (), pf_mul_x_1_y | x)
    end
  val (pf_fac, pf_mul | res) = loop (n, 1)
  prval () = MULlem10 (pf_mul)
in
  (pf_fac | res)
end

// fac2 is another implementation, employing the GMP library for variable-precision arithmetic
// it can cope with very big numbers
fun fac2 {n:nat} .<n>. (n: int n): [r:int] (FACT (n, r) | intinfptr_gc r) =
  if n > 0 then let
    val (pf1 | (pf1_gc, pf1_at | p1)) = fac2 (n-1)
    val (pf_mul | r) = n * !p1
    val () = intinf_free (pf1_gc, pf1_at | p1)
  in
    (FACTind (pf1, pf_mul) | r)
  end else begin
    (FACTbas () | intinf_make 1)
  end // end of [if]

fn fac_usage (cmd: string): void =
  prerrf ("Usage: %s [integer]\n", @(cmd))

implement main (argc, argv) =
  if argc >= 2 then let
    val n0 = int1_of argv.[1]
    val () = assert_errmsg
      (n0 >= 0, "The integer argument needs to be nonnegative.\n")
    val (pf | res) = fac1 n0
  in
    print "The factorial of "; print n0; print " = "; print res; print_newline ()
  end else begin
    fac_usage (argv.[0]); exit 1
  end