(***************************************************************)
(*
** ARIBAS Code
** zur Vorlesung
** O. Forster: Algorithmische Zahlentheorie II
** WS 2009/10 LMU Muenchen
**
** Chap. 4  The AKS primality test
**
** Date of last change: 2009-12-03
** <forster AT lmu.math.de>
*)
(***************************************************************)
(* Example:

==> N := next_prime(random(2**40)).

working ,,,,,,,,,, probable prime:
-: 30_56638_53799

==> bound := bit_length(N)**2.
-: 1521

==> r := AKS_findr(N,bound).
-: 1579

==> a := random(1000).
-: 868

==> AKS_check(N,r,a).
-: true

** For the full AKS primality test one would have to call AKS_check(N,r,a) 
** for all integers a = 1,2,3,..., upto floor(sqrt(r)*log(N)/log(2))
*)
(***************************************************************)
(*
** Constructs an array of the prime factors of x,
** where multiple prime factors are listed repeatedly
** according to their multiplicity
** Uses trial division by primes < 2**16
**
** If the last factor is > 2**32, it is not necessarile prime
*)
function factorlist0(x: integer; verbose := 0): array;
var
    st: stack;
    q: integer;
begin
    x := abs(x);
    if x < 2 then
        return {x};
    end;
    q := 2;
    while q := factor16(x,q) do
        stack_push(st,q);
        x := x div q;
        if verbose then
            writeln(q);
        end;
    end;
    stack_push(st,x);
    if verbose then
        writeln(x);
    end;
    if x > 2**32 then
        writeln("last factor not neccessarily prime");
    end;
    return stack2array(st);
end;
(*---------------------------------------------------------*)
(*
** Constructs an array of primefactors of x, where
** every prime factor is listed only once
*)
function primefactors(x: integer): array;
var
    i,p,p0: integer;
    vec: array;
    st: stack;
begin
    vec := factorlist0(x);
    p0 := 0;
    for i := 0 to length(vec)-1 do
        p := vec[i];
        if p /= p0 then
            stack_push(st,p);
            p0 := p;
        end;
    end;
    return stack2array(st);
end;
(***************************************************************)
(*-------------------------------------------------------------*)
(*
** Argument r must be a prime, and x /= 0 mod p.
** The function calculates the order of x
** in the multiplicative group (Z/r)*
*)
function Zrstar_order(r,x: integer): integer;
var
    s, p, k, z: integer;
    pvec: array;
begin
    if x mod r = 0 then
        writeln("error Zrstar_order: arguments must be coprime");
        return 0;
    end;
    s := r-1;
    pvec := primefactors(s);
    for k := 0 to length(pvec)-1 do
        p := pvec[k];
        while s mod p = 0 do
            s := s div p;
        end;
        z := (x ** s) mod r;
        while z /= 1 do
            z := (z ** p) mod r;
            s := s*p;
        end;  
    end;
    return s;
end;
(***************************************************************)
(*---------------------------------------------------------*)
(*
** Searches for a prime r < N which does not divide N and 
** such that the multiplicative order of N mod r is >= bound
** Return value:
**      r, if such a prime is found
**     -1, if N is not prime
**      0  otherwise
*)
function AKS_findr(N,bound: integer): integer;
var
    r: integer;
begin
    r := next_prime(bound);
    while r < N do
        if N mod r = 0 then
            return -1;
        elsif Zrstar_order(r,N) >= bound then
            return r;
        else
            r := next_prime(r+2);
        end;
    end;
    return 0;
end;
(*---------------------------------------------------------*)
(*
** Checks the congruence
**
**      (X + a)**N = X**N + a mod (N, X**r - 1)
*)
function AKS_check(N,r,a: integer): boolean;
var
    k, i, len, s: integer;
    F,G: array;
begin
    F := (a,1);
    if r = 1 then F := {a+1}; end;
    for k := bit_length(N)-2 to 0 by -1 do
        F := ZnXsquare(N,F);
        if length(F) > r then
            G := F[r..];
            F := (F[0..r-1] + G) mod N;
            F := vec_rtrim(F);
        end;
        if bit_test(N,k) then
            F := ZnXmult(N,F,(a,1));
            if length(F) = r+1 then
                F[0] := (F[0] + F[r]) mod N;
                F := vec_rtrim(F[0..r-1]);
            end;
        end;
    end;
    s := N mod r;
    if length(F) = s+1 then
        F[s] := F[s] - 1;
        F[0] := F[0] - a;
        if vec_rtrim(F) = () then
            return true;
        end;
    end;
    return false;
end;
(*---------------------------------------------------------*)
(*
** Same as AKS_check, but using the Karatsuba algorithm
** for polynomial squaring
*)
function AKS_kcheck(N,r,a: integer): boolean;
var
    k, i, len, s: integer;
    F,G: array;
begin
    (* calculate (X + a)**N mod (N, X**r - 1) *)
    F := (a,1);
    if r = 1 then F := {a+1}; end;
    for k := bit_length(N)-2 to 0 by -1 do
        F := ZnXksquare(N,F);
        if length(F) > r then
            G := F[r..];
            F := (F[0..r-1] + G) mod N;
            F := vec_rtrim(F);
        end;
        if bit_test(N,k) then
            F := ZnXmult(N,F,(a,1));
            if length(F) = r+1 then
                F[0] := (F[0] + F[r]) mod N;
                F := vec_rtrim(F[0..r-1]);
            end;
        end;
    end;
    (* compare with (X**N + a) mod (N, X**r - 1) *)
    s := N mod r;
    if length(F) > s then
        F[s] := F[s] - 1;
        F[0] := F[0] - a;
        if vec_rtrim(F) = () then
            return true;
        end;
    end;
    return false;
end;
(****************************************************************************)
(*
** Multiplication of polynomials over Z/n
** A polynomial of degree n is given by an array
**      (a0,a1,....,an), an /= 0.
** This represents the polynomial
**      an*X**n + ... + a1*X + a0
** The zero polynomial is given by the empty array ()
*)
(*------------------------------------------------------------*)
(*
** Deletes leading zeroes from the array F.
*)
function vec_rtrim(F: array): array;
var
    d, i: integer;
begin
    d := length(F)-1;
	if d < 0 or F[d] /= 0 then
		return F;
	end;
	dec(d);
    while d >= 0 and F[d] = 0 do
        dec(d);
    end;
    return F[0..d];
end;
(*-------------------------------------------------------------*)
(*
** Product of two polynomials in (Z/N)[X]
*)
function ZnXmult(N: integer; F,G: array): array;
var
    n,m,i,i0,i1,k: integer;
    x: integer;
    R: array;
begin
    n := length(F)-1;
    m := length(G)-1;
    if n<0 or m<0 then
        return ();
    end;
    R := alloc(array,n+m+1);
    for k := 0 to n+m do
        x := 0;
        i0 := max(0,k-m);
        i1 := min(n,k);
        for i := i0 to i1 do
            x := x + F[i]*G[k-i];
        end;
        R[k] := x mod N;
    end;
    return vec_rtrim(R);
end;
(*------------------------------------------------------------*)
(*
** Square of a polynomial in (Z/N)[X]
*)
function ZnXsquare(N: integer; F: array): array;
var
    n,m,i,i0,k: integer;
    x: integer;
    R: array;
begin
    n := length(F)-1;
    if n < 0 then
        return ();
    end;
    R := alloc(array,2*n+1);
    for k := 0 to 2*n do
        i0 := max(0,k-n);
        m := (k+1) div 2;
        x := 0;
        for i := i0 to m-1 do
            x := x + F[i]*F[k-i];
        end;
        x := 2*x;
        if even(k) then
            x := x + F[m]**2;
        end;
        R[k] := x mod N;
    end;
    return vec_rtrim(R);
end;
(****************************************************************************)
(*
** Square of a polynomial in (Z/n)[X], using Karatsuba
*)
function ZnXksquare(N: integer; P: array): array;
var
    KK: array;
    len,alen,rlen,n: integer;
    F,A,R: integer;
    P2: array;
begin
    len := length(P);
    alen := 2**bit_length(len);
    rlen := 2*len - 1;
    KK := realloc(P,len+alen+rlen,0);
    F := 0; A := len; R := len+alen;
    n := karatsuba(KK,F,len,A,alen,R,rlen);
    P2 := vec_rtrim(KK[R..R+n-1] mod N);
    return P2;
end;
(*----------------------------------------------------------------------*)
(*
** Auxiliary function for ZnXksquare
** alen >= len must be a power of 2,
** rlen >= 2*len - 1
*)
function karatsuba(var KK: array; F,len, A,alen, R,rlen: integer): integer;
var
    F1,F2,A1,A2,R1,R2: integer;
    len1,len2,alen1,alen2,rlen1,rlen2: integer;
    n1,n2,n3,r,i: integer;
begin
    if len <= 2 then
        if len = 0 then
            r := 0;
        elsif len = 1 then
            KK[R] := KK[F]**2;
            r := 1;
        else    (* len = 2 *)
            KK[R] := KK[F]**2;
            KK[R+1] := 2*KK[F]*KK[F+1];
            KK[R+2] := KK[F+1]**2;
            r := 3;
        end;
        for i := r to rlen-1 do
            KK[R+i] := 0;
        end;
        return r;
    end;
    len1 := (len+1) div 2;
    len2 := len - len1;
    F1 := F; F2 := F + len1;
    alen1 := alen2 := alen div 2;
    A1 := A; A2 := A + alen1;
    rlen1 := 2*len1;
    rlen2 := rlen - rlen1;
    R1 := R; R2 := R + rlen1;
    for i := 0 to len2-1 do
        KK[A1+i] := KK[F1+i] + KK[F2+i];
    end;
    if odd(len) then
        KK[A1+len2] := KK[F1+len2];
    end;
    n1 := karatsuba(KK,F1,len1,A2,alen2,R1,rlen1);
    n2 := karatsuba(KK,F2,len2,A2,alen2,R2,rlen2);
    n3 := karatsuba(KK,A1,len1,A2,alen2,F,len);
    for i := 0 to n1-1 do
        KK[F+i] := KK[F+i] - KK[R1+i];
    end;
    for i := 0 to n2-1 do
        KK[F+i] := KK[F+i] - KK[R2+i];
    end;
    for i := 0 to n3-1 do
        KK[R+len1+i] := KK[R+len1+i] + KK[F+i];
    end;
    for i := R2+n2 to R+rlen-1 do
        KK[i] := 0;
    end;
    return rlen1 + n2;
end;
(********************************************************************)