Simple AKS in Perl

taks.pl

#!/usr/bin/env perl
package Acme::Math::AKS;
use warnings;
use strict;
use Math::BigInt   try => 'GMP';
use Math::BigFloat try => 'GMP';
# Technically we can get by without any of these functions
use Math::Prime::Util qw/primes next_prime euler_phi is_prob_prime/;

# A simple standalone Perl implementation of the AKS primality algorithm:
#   http://www.cse.iitk.ac.in/users/manindra/algebra/primality_v6.pdf
# Copyright 2012 by Dana Jacobsen (dana@acm.org)
# This program is free software using the Perl 5 artistic license v2.

my $try_optimistic_r = 0;   # r search: all ints from 2 OR primes from log2n
my $use_sqrtn_test = 1;     # if we've checked all primes under sqrtn, stop.
my $verbose = 0;            # display progress reports?
my $poly_bignum;            # INTERNAL: will poly code use BigInt coefs?

__PACKAGE__->run unless caller;

sub is_aks_prime {
  my $n = shift;
  return wantarray ? (0, "AKS - COMPOSITE - less than two") : 0 if $n < 2;
  $n = Math::BigInt->new("$n") unless ref($n) eq 'Math::BigInt';

  # 1. mpz_perfect_power_p
  return wantarray ? (0, "AKS - COMPOSITE - perfect power") : 0
    if is_perfect_power($n);

  # 2. Find smallest r such that order_r(n) > log^2(n)
  # Notes:
  #  - log = base 2 logarithm (AKS page 3).
  #  - calculate log^2(n) as a bigfloat, then limit = floor as a native int
  #  - we use floor for limit because we check o_r(n) > limit
  #  - one could make the case for log^2n to be:
  #       (a) floor(log(n))^2,
  #       (b) floor(log(n)*log(n))
  #       (c) ceil(log(n))^2
  #    where (a) <= (b) <= (c).  GMP impls use (c) because mpz_sizeinbase(2)
  #    does ceil(log(n)).  Salembier et al. (2005) use (b).  We'll use (b).
  my $sqrtn = int(Math::BigFloat->new($n)->bsqrt->bfloor->bstr);
  my $log2n = Math::BigFloat->new($n)->blog(2);
  my $log2_squared_n = $log2n * $log2n;
  my $limit = int($log2_squared_n->bfloor->bstr);

  print "$n discovering r ($limit)." if $verbose;

  my $r = $try_optimistic_r  ?  2  :  next_prime($limit);
  # Check for small factors up to where we're starting
  foreach my $f (@{primes(0,$r-1)}) {
    return wantarray ? (1, "AKS - PRIME - trial division through $f") : 1
      if $f == $n;
    if ( ! ($n % $f) ) {
      return wantarray ? (0, "AKS - COMPOSITE - factor $f") : 0;
    }
  }

  # AKS Lemma 4.3 says we _will_ find an r <= max(3, ceil(log^5n))
  while ($r < $n) {
    print "." if $verbose && (!$try_optimistic_r || is_prob_prime($r));
    # Check divisibility
    if ( ! ($n % $r) ) {
      print "\n" if $verbose;
      return wantarray ? (0, "AKS - COMPOSITE - factor $r") : 0;
    }
    # We've looked at every prime r up to sqrt(n) and none divide n.
    if ($use_sqrtn_test && $r >= $sqrtn) {
      return wantarray ? (1, "AKS - PRIME - trial division through $r") : 1
    }
    # see if the multiplicative order of n mod r is > limit
    last if order($r, $n, $limit) > $limit;
    # increment r
    $r = $try_optimistic_r  ?  $r+1  :  next_prime($r);
  }
  print "$r\n" if $verbose;

  # 4. n is prime if we've checked all r up to n.
  # (we already would have returned because no primes through sqrtn divided n)
  return wantarray ? (1, "AKS - PRIME - trial division through $r") : 1
    if $r >= $n;

  # 5. Check (X+a)^n = X^n + a (mod X^r-1,n) for a = a to rlimit
  # - if r is prime then euler_phi(r) = r-1
  my $rlimit = int( Math::BigFloat->new(euler_phi($r))
                    ->bsqrt
                    ->bmul($log2n)
                    ->bfloor
                    ->bstr);
  print "testing polynomials from 1 to $rlimit with r = $r\n" if $verbose;

  # If n is less than a half-word, then we can never overflow in native ints.
  # n + (n*n) < word if n < halfword (so 2^16-1 for 32-bit, 2^32-1 for 64-bit)
  $poly_bignum = 1;
  if ( $n < ( (~0 == 4294967295) ? 65535 : 4294967295 ) ) {
    $poly_bignum = 0;
    $n = int($n->bstr) if ref($n) eq 'Math::BigInt';
  }

  for (my $a = 1; $a <= $rlimit; $a++) {
    if ( ! test_anr($a, $n, $r) ) {
      print "COMPOSITE\n" if $verbose;
      return wantarray ? (0, "AKS - COMPOSITE - failed polynomial test 1 .. $rlimit (a = $a, r = $r)") : 0;
    }
    print "." if $verbose;
  }
  print "\n" if $verbose;

  if (!is_prob_prime($n)) {
    die "Whoa -- AKS says $n is prime but BPSW says it isn't.\n";
  }
  return wantarray ? (1, "AKS - PRIME - passed polynomial tests 1 .. $rlimit with r = $r") : 1;
}


sub test_anr {
  my($a, $n, $r) = @_;

  my $pp = poly_mod_pow(poly_new($a, 1), $n, $r, $n);

  # my @want_poly;
  # $want_poly[0] = $a % $n;
  # $want_poly[$n % $r] = 1;
  # print "I want the result to be:\n  "; poly_print(\@want_poly);

  $pp->[$n % $r] = (($pp->[$n % $r] || 0) -  1) % $n;  # subtract X^(n%r)
  $pp->[      0] = (($pp->[      0] || 0) - $a) % $n;  # subtract a

  # return 0 (composite) if any coefficient is non-zero, 1 otherwise.
  return 0 if scalar grep { $_ } @$pp;

  1;
}

# Polynomial functions.

sub poly_new {
  my @poly;
  if ($poly_bignum) {
    foreach my $c (@_) {
      push @poly, (ref $c eq 'Math::BigInt') ? $c->copy : Math::BigInt->new("$c");
    }
  } else {
    push @poly, $_ for (@_);
    push @poly, 0 unless scalar @poly;
  }
  return \@poly;
}

sub poly_print {
  my($poly) = @_;
  warn "poly has null top degree" if $#$poly > 0 && !$poly->[-1];
  foreach my $d (reverse 1 .. $#$poly) {
    my $coef = $poly->[$d];
    print "", ($coef != 1) ? $coef : "", ($d > 1) ? "x^$d" : "x", " + "
      if $coef;
  }
  my $p0 = $poly->[0] || 0;
  print "$p0\n";
}

sub poly_mod_mul {
  my($px, $py, $r, $n) = @_;

  my $px_degree = $#$px;
  my $py_degree = $#$py;
  my @res;

  # convolve(px, py) mod (X^r-1,n)
  my @indices_y = grep { $py->[$_] } (0 .. $py_degree);
  for (my $ix = 0; $ix <= $px_degree; $ix++) {
    my $px_at_ix = $px->[$ix];
    next unless $px_at_ix;
    foreach my $iy (@indices_y) {
      my $py_at_iy = $py->[$iy];
      my $rindex = ($ix + $iy) % $r;  # reduce mod X^r-1
      if (!defined $res[$rindex]) {
        $res[$rindex] = $poly_bignum ? Math::BigInt->bzero : 0
      }
      $res[$rindex] = ($res[$rindex] + ($py_at_iy * $px_at_ix)) % $n;
    }
  }
  # In case we had upper terms go to zero after modulo, reduce the degree.
  pop @res while !$res[-1];
  return \@res;
}

# Square using product scanning, putting off the modulo until the sum is
# complete.  If we could overflow, fall back to poly_mod_mul.
sub poly_mod_sqr {
  my($p, $r, $n) = @_;
  my @res;
  my $degree = $#$p;
  # Sum can be as big as (d+1) * n * n because we're putting off the modulo.
  return poly_mod_mul($p, $p, $r, $n)
    if !$poly_bignum && $n > int((sqrt(~0)/($degree+1)));
  foreach my $d (0 .. 2*$degree) {
    my $sum = $poly_bignum ? Math::BigInt->bzero : 0;
    my $start = ($d <= $degree) ? 0 : $d - $degree;
    my $end   = int($d/2);
    for (my $s = $start; $s <= $end; $s++) {
      my $c1 = $p->[$s];   # my $c2 = $p[$d-$s];
      $sum += ($s*2 == $d)  ?  $c1*$c1  :  2 * $c1 * $p->[$d - $s];
    }
    my $rindex = $d % $r;
    $res[$rindex] = (defined $res[$rindex]) ? ($res[$rindex] + $sum) % $n
                                            : $sum % $n;
  }
  pop @res while !$res[-1];
  return \@res;
}

sub poly_mod_pow {
  my($pn, $power, $r, $mod) = @_;
  my $res = poly_new(1);
  my $p = $power;

  while ($p) {
    $res = poly_mod_mul($res, $pn, $r, $mod) if ($p & 1);
    $p >>= 1;
    $pn  = poly_mod_sqr($pn, $r, $mod) if $p;
  }
  return $res;
}

# Utility functions

sub is_perfect_power {
  my $n = shift;
  my $log2n = _log2($n);
  $n = Math::BigInt->new("$n") unless ref($n) eq 'Math::BigInt';
  for my $e (@{primes($log2n)}) {
    return 1 if $n->copy()->broot($e)->bpow($e) == $n;
  }
  0;
}
sub order {
  my($r, $n, $lim) = @_;
  $lim = $r unless defined $lim;
  $n = Math::BigInt->new("$n") unless ref($n) eq 'Math::BigInt';

  return 1 if $n->copy->bmod($r) == 1;
  for (my $j = 2; $j <= $lim; $j++) {
    return $j if $n->copy->bmodpow($j, $r) == 1;
  }
  return $lim+1;
}

sub _log2 {
  my $n = shift;
  my $log2n = 0;
  $log2n++ while ($n >>= 1);   # floor of n->blog(2);
  $log2n;
}



#######################################################################

sub run {
  require Getopt::Long;
  $|++;

  my $get_certificate = 0;
  my $test_perfect_powers = 0;
  my $test_examples = 0;
  my $test_all_primes = 0;
  my $test_all_composites = 0;
  my $usage = 0;

  Getopt::Long::GetOptions(
     'verbose'             => \$verbose,
     'sqrt!'               => \$use_sqrtn_test,
     'optr|optimr|optimisticr' => \$try_optimistic_r,
     'certificate'         => \$get_certificate,
     'test-perfect-powers' => \$test_perfect_powers,
     'test-suite'          => \$test_examples,
     'test-primes'         => \$test_all_primes,
     'test-composites'     => \$test_all_composites,
     'help|usage|?'        => \$usage,
  ) or show_usage();
  show_usage() if $usage;

  if ($test_perfect_powers) { # Test our is_perfect_power function.
    foreach my $nstr (qw/7 10 11 14 26 39 72355 8812214355 987687645327744 4345465791980381796 543565790912343654309875/) {
      die "is_perfect_power($nstr) is wrong\n"
        if is_perfect_power(Math::BigInt->new($nstr));
    }
    foreach my $b (2 .. 25, 912673, 6436343, 24137569, 69343957, 912673) {
      foreach my $e (2 .. 14) {
        die "is_perfect_power($b ** $e) is wrong\n"
          unless is_perfect_power( Math::BigInt->new($b)->bpow($e) );
      }
    }
    foreach my $n (42875, 225866529, 1564031349) {
      die "is_perfect_power($_) is wrong\n" unless is_perfect_power($n);
    }
  }

  if ($test_examples) { # Run a set of selected test numbers
    $use_sqrtn_test = 0;  # Without this we would need much larger values
    foreach my $n (qw/
                      0 1 4 6 15 27 561 8911 74513 101101
                      2 3 5 7 11 31 41 89 2621 3571 11701 100237
                     /) {
      my ($isp, $cert) = is_aks_prime($n);
      print "$n $cert\n";
      die "INCORRECT RESULT FOR $n\n" unless $isp == !!is_prob_prime($n);
    }
  }

  if ($test_all_primes || $test_all_composites) {
    $use_sqrtn_test = 0;  # Without this we would need much larger values
    my $lastlen = 0;
    my $n = (@ARGV && $ARGV[0] > 2) ? $ARGV[0] : 2;
    $n = Math::BigInt->new("$n") unless $n < 1_000_000_000;
    while (1) {
      if ( ($test_all_composites && !is_prob_prime($n)) ||
           ($test_all_primes && is_prob_prime($n)) ) {
        my($isp, $cert) = is_aks_prime($n);
        my $text = sprintf("%8d %s", $n, $cert);
        print $text;
        die "\nINCORRECT RESULT FOR $n\n" unless $isp == !!is_prob_prime($n);
        print " " x ($lastlen - length($text)) if $lastlen > length($text);
        print "\r";
        $lastlen = length($text);
      }
      $n++;
    }
  }

  foreach my $n (@ARGV) {
    my($is_prime, $cert);
    if ($get_certificate) {
      ($is_prime, $cert) = is_aks_prime($n);
    } else {
      $is_prime = is_aks_prime($n);
    }
    print "$n is ", ($is_prime) ? "PRIME" : "COMPOSITE", "\n";
    print "Cert: $cert\n" if defined $cert;
  }
}

sub show_usage {
  die <<EOU;
Usage: $0 [options] [n ...]

  Options:
     --help            displays this message
     --verbose         display verbose progress output
     --certificate     displays how the result was obtained
     --nosqrt          keep going even if r > sqrt(n)
     --optr            Spend more time calculating r in hopes it will be smaller
     --test-suite      run a simple test suite of numbers
     --test-primes     run a non-stop test of primes starting from n
     --test-composites run a non-stop test of primes starting from n
     --test-perfect-powers  run a simple test of the perfect_power function

  Examples:
     $0 877                         run AKS for 877
     $0 -v 877                      run AKS for 877 showing progress
     $0 -cert 877                   run AKS for 877 and show certificate
     $0 -nosqrt -v 877              run AKS for 877 without stopping early
     $0 -v -cert 100003             run AKS for 100003, showing progress
     $0 --test-primes               run a test of all
EOU
}

1;