Common Lisp prime testing

primes.lisp

;;;; Needs the iterate package.

(defmacro maybep (&body body)
  "Give this a series of forms.  It will return the result of the first one that
  yields T or NIL.  Any other result will cause the next form to be evaluated.
  Short-circuits like IF or COND."
  (labels ((maybep-expand (forms)
             (if (endp forms)
                 `(error "All evaluated forms yielded MAYBE.")
                 `(let ((result ,(car forms)))
                    (if (or (eq result t) (eq result nil))
                        result
                        ,(maybep-expand (cdr forms)))))))
    (maybep-expand body)))

(defun primep (num)
  (declare (type integer num))
  (and (plusp num)
       (maybep
         (cached-prime-maybe num)
         (trial-division-maybe num)
         (sprp-maybe num)
         (probabilistic-prime-p num))))

;;; This is basically the Sieve of Eratosthenes.  The implementation is
;;; pretty efficient; it only uses integer addition (no division or even
;;; multiplication, aside from one call to isqrt) and it has the usual
;;; boundary constraints (start marking prime multiples at the square of
;;; the prime, stop marking when the primes exceed the square root of the
;;; array limit).
;;;
;;; I tried special-casing 3, 5, and 7 and using a wheel to skip over
;;; their multiples, but the overhead of maintaining the wheel was greater
;;; than the time savings of skipping those numbers.  (The cost of
;;; checking a composite number is very low: one lookup in an array, a
;;; comparison to zero, and another addition.)
(defun generate-prime-array (size)
  "Creates and returns a bit array where each bit is 1 if its index is prime and
   0 if its index is composite.  SIZE is the number of elements in the array."
  (declare (type (integer 0 #.(1- array-dimension-limit)) size))
  (declare (optimize (speed 3)))
  (let ((result (make-array size :element-type 'bit :initial-element 1)))
    (setf (sbit result 0) 0  ; Special case.
          (sbit result 1) 0) ; Special case.
    (iter (for i index-of-vector result from 4 by 2)  ; Multiples of two.
          (declare (type (integer 0 #.(+ 2 array-dimension-limit)) i))
          (setf (sbit result i) 0))
    (iter (for b in-vector result          ; Multiples of the odd primes.
                 from 3 to (isqrt (1- size)) by 2
                 with-index i)
          (declare (type bit b)
                   (type (integer 0 #.(+ 2 array-dimension-limit)) i))
          (when (not (zerop b))
            (iter (for j index-of-vector result from (* i i) by i)
                  ;; Maximum value of j is ARRAY-DIMENSION-LIMIT plus the
                  ;; maximum value of i, which is the square root of (SIZE
                  ;; minus 1).  SIZE is at most one less than
                  ;; ARRAY-DIMENSION-LIMIT.
                  (declare
                   (type (integer 0 #.(+ array-dimension-limit
                                         (isqrt (- array-dimension-limit 2))))
                         j))
                  (setf (sbit result j) 0))))
    result))

(let ((prime-ary (generate-prime-array (- (expt 2 29) 2))))
  (defun cached-prime-maybe (num)
    (declare (type (integer 0) num))
    (if (< num (length prime-ary))
        (= 1 (sbit prime-ary num))
        :maybe)))

(let* ((primes (iter (declare (iterate:declare-variables))
                     (if-first-time (collecting 2 into result))
                     (for (the (integer 0 255) n) from 3 to 254)
                     (when (iter (for d in result)
                                 (never (zerop (mod n d))))
                       (collecting n into result))
                     (finally (return result))))
       (last-prime (the (integer 0 255) (car (last primes)))))
  (defun trial-division-maybe (num)
    (declare (type (integer 0) num)
             (optimize (speed 3)))
    (cond
      ((<= num last-prime)
       (not (null (position num primes))))
      ((iter (declare (iterate:declare-variables))
             (for (the (integer 0 255) d) in primes)
             (thereis (zerop (mod num d))))
       nil)
      (t :maybe))))

(defun sprp-maybe (num)
  (declare (type (integer 0) num))
  (labels ((all-sprp-p (&rest bases)
             (iter (for base in bases)
                   (always (sprp-p num base)))))
    (cond
      ((< num         9080191)
       (all-sprp-p 31 73))
      ((< num      4759123141)
       (all-sprp-p 2 7 61))
      ((< num    118670087467)
       (and (/= num 3215031751)
            (all-sprp-p 2 3 5 7)))
      ((< num   2152302898747)
       (all-sprp-p 2 3 5 7 11))
      ((< num   3474749660383)
       (all-sprp-p 2 3 5 7 11 13))
      ((< num 341550071728321)
       (all-sprp-p 2 3 5 7 11 13 17))
      (t :maybe))))

(defun sprp-p (num base)
  (declare (type integer num base))
  "Checks to see whether NUM is a strong probable-prime base BASE."
  (destructuring-bind (s d)
      (iter (declare (iterate:declare-variables))
            (for (the integer n) first (1- num) then (ash n -1))
            (while (evenp n))
            (for (the integer s) from 1)
            (finally (return (list s n))))
    (let ((test (modpower base d num)))
      (or (= 1 test)
          (iter (declare (iterate:declare-variables))
                (repeat s)
                (with (the integer nm1) = (1- num))
                (for (the integer ti) first test then (mod (* ti ti) num))
                (thereis (= nm1 ti)))))))

(defun modpower (num power modulus)
  "Raises NUM to the POWER power, modulo MODULUS.  Faster than
   (mod (expt num power) modulus)."
  (declare (type (integer 0) num power modulus)
           (values (integer 0)))
  (labels ((modpower-r (b e a)
             (declare (type (integer 0) b e a))
             (if (zerop e)
                 a
                 (modpower-r (mod (* b b) modulus)
                             (truncate e 2)
                             (if (oddp e)
                                 (mod (* b a) modulus)
                                 a)))))
    (modpower-r (mod num modulus) power 1)))

(defun probabilistic-prime-p (num (trials 8))
  (declare (type (integer 0) num))
  (iter (repeat trials)
        (for a = (random num))
        (always (sprp-p num a))))