/* Author: Jason Eisner <jason@cs.jhu.edu>, 1998.
   Use this code at your own risk; please give appropriate credit.
   Translated from part of the Lisp code by same author in primes.lisp. 
 */

#include <stdlib.h>
#include <assert.h>
#include "prime.h"

/* up to caller to seed the random number generator. */

typedef unsigned long long ullong;   /* we use this during multiplications to prevent overflow */

ulong uliNearbyPrime(ulong uli);
bool fPrimeMillerRabin(ulong n, int iChecks);
ulong uliGcd(ulong uli1, ulong uli2);
ulong uliRandomBase(ulong n);
bool fStrongPseudoprime(ulong n, ulong base, ulong t, int s);
ulong uliExptMod(ulong b, ulong e, ulong m);

ulong uliMaxSquarable = 0;

#define NUM_CHECKS 20   /* gives us confidence of 2^-40, or 1 in a trillion */


/* Pick nearest prime >= uli. */

ulong uliNearbyPrime(ulong uli) {
  if (uli <= 2) uli=2;
  else if (uli % 2 == 0) uli++;
  while (!fPrimeMillerRabin(uli, NUM_CHECKS)) uli+=2;
  assert(uli > 1);  /* no wraparound allowed */
  return uli;
}


/*  The Miller-Rabin test.  Relies on the fact that if n is prime,
    it is an strong pseudoprime to every base in [1, n), whereas if
    it's composite, it is a strong pseudoprime to at most 1/4 of these
    bases.

    We can do enough random spot checks to bring the probability of an
    false positive very low, assuming the random number generator is
    good.  We never have false negatives - if we say it's composite,
    that's because we have a witness.  */

bool fPrimeMillerRabin(ulong n, int iChecks) {   
  ulong t = n-1;
  int s = 0;
  int i;

  assert(n > 1);
  if (uliMaxSquarable == 0) { /* initialize overflow checker */
    ullong ulli1 = 1;
    ullong ulli2 = 1;
    while (ulli2) {
      ulli1 <<= 1;
      ulli2 <<= 2;
    }
    uliMaxSquarable = (ulong)(ulli1-1);
  }

  while (t % 2 == 0) {
    t >>= 1;
    s++;
  }
  /* now n-1 = 2^s * t with t odd */

  for (i=0; i < iChecks; i++) 
    if (!fStrongPseudoprime(n, uliRandomBase(n), t, s))
      return 0;
    
  return 1;
}

/* Pick an integer from [1, n) that is coprime to n. */

ulong uliRandomBase(ulong n) {
  ulong uli;
  do {
    uli = 1 + (lrand48() % (n-1));
  } while (uliGcd(uli, n) > 1);
  return uli;
}

ulong uliRand(void) {   /* lrand48() returns uniform distrib over [0, 2^31) */
  ulong uli = 1 << 31;  /* just bigger than lrand48() can return */
  uli <<= 31;
  uli <<= 31;
  assert (uli==0);  /* make sure long isn't too wide for the following to work (should work up to 64 bits) */
  uli = lrand48();
  uli <<= 31;
  uli |= lrand48();
  uli <<= 31;
  uli |= lrand48();
  return uli;
}


/* Euclid's GCD algorithm. */

ulong uliGcd(ulong uli1, ulong uli2) {
  assert(uli1 > 0);
  assert(uli2 > 0);
  while (1) {
    if ((uli1%=uli2)==0)
      return uli2;
    if ((uli2%=uli1)==0)
      return uli1;
  }
}

bool fStrongPseudoprime(ulong n, ulong base, ulong t, int s) {   /* where n-1 = t*(2^s) */
  ulong uli = uliExptMod(base, t, n);
  int r;

  if (uli == 1)
    return 1;
  for (r=0; r<s; r++) {
    if (uli == n-1) return 1;
    assert (uli <= uliMaxSquarable);
    uli = ((ullong)uli*uli) % n;
  }

  return 0;
}
  
/* Quickly finds b^e (mod m).*/

ulong uliExptMod(ulong b, ulong e, ulong m) {
  ulong prod = 1;
  ulong uli = 1;   /* single bit */
  while (uli <= e)
    uli <<= 1;
  while (uli >>= 1) {
    assert (prod <= uliMaxSquarable);
    prod = ((ullong)prod * prod) % m;
    if (uli & e) {
      assert (prod <= uliMaxSquarable);
      prod = ((ullong)b * prod) % m;
    }
  }
  return prod;
}