;; prime.scm -- prime number utilities
;; Copyright (c) 2004-2014 Alex Shinn.  All rights reserved.
;; BSD-style license: http://synthcode.com/license.txt

;;> Prime and number theoretic utilities.

;; Given \var{n} and a continuation \var{return},
;; returns (\var{return} \var{k2} \var{n2}) where
;; \var{k2} is the power of 2 in the factorization of \var{n}, and
;; \var{n2} is product of all other prime powers dividing \var{n}
(define (factor-twos n return)
  (let ((b (first-set-bit n)))
    (return b (arithmetic-shift n (- b)))))

;;> Returns the multiplicative inverse of \var{a} modulo \var{b}.
(define (modular-inverse a b)
  (let lp ((a1 a) (b1 b) (x 0) (y 1) (last-x 1) (last-y 0))
    (if (zero? b1)
        (if (negative? last-x) (+ last-x b) last-x)
        (let ((q (quotient a1 b1)))
          (lp b1 (remainder a1 b1)
              (- last-x (* q x)) (- last-y (* q y))
              x y)))))

;;> Returns (remainder (expt a e) m).
(define (modular-expt a e m)
  (let lp ((tmp a) (e e) (res 1))
    (if (zero? e)
        res
        (lp (remainder (* tmp tmp) m)
            (arithmetic-shift e -1)
            (if (odd? e) (remainder (* res tmp) m) res)))))

;;> Returns true iff n and m are coprime.
(define (coprime? n m)
  (= 1 (gcd n m)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Exact prime testing.

;; All primes under 1000.
(define prime-table
  '#(  2   3   5   7  11  13  17  19  23  29  31  37  41  43  47  53  59
      61  67  71  73  79  83  89  97 101 103 107 109 113 127 131 137 139
     149 151 157 163 167 173 179 181 191 193 197 199 211 223 227 229 233
     239 241 251 257 263 269 271 277 281 283 293 307 311 313 317 331 337
     347 349 353 359 367 373 379 383 389 397 401 409 419 421 431 433 439
     443 449 457 461 463 467 479 487 491 499 503 509 521 523 541 547 557
     563 569 571 577 587 593 599 601 607 613 617 619 631 641 643 647 653
     659 661 673 677 683 691 701 709 719 727 733 739 743 751 757 761 769
     773 787 797 809 811 821 823 827 829 839 853 857 859 863 877 881 883
     887 907 911 919 929 937 941 947 953 967 971 977 983 991 997 ))

;;> Returns true iff \var{n} is definitely prime.  May take an
;;> impossibly long time for large values.
(define (provable-prime? n)
  (if (or (even? n) (<= n 2))
      (= 2 n)
      (let ((limit (exact (ceiling (sqrt n)))))
        (define (by-twos d)
          (cond ((> d limit) #t)
                ((zero? (remainder n d)) #f)
                (else (by-twos (+ d 2)))))
        (let ((len (vector-length prime-table)))
          (let lp ((i 0))
            (if (>= i len)
                (by-twos (vector-ref prime-table (- len 1)))
                (let ((d (vector-ref prime-table i)))
                  (cond
                   ((> d limit) #t)
                   ((zero? (remainder n d)) #f)
                   (else (lp (+ i 1)))))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Probable primes.

;; Given \var{n}, return a predicate that tests whether
;; its argument \var{a} is a witness for \var{n} not being prime,
;; either (1) because \var{a}^(\var{n}-1)≠1 mod \var{n}
;; \em{or} (2) because \var{a}'s powers include
;; a third square root of 1 beyond {1, -1}
(define (miller-rabin-witnesser n)
  (let ((neg1 (- n 1)))
    (factor-twos neg1
     (lambda (twos odd)
       (lambda (a)
         (let ((b (modular-expt a odd n)))
           (let lp ((i 0) (b b))
             (cond ((= b neg1)
                    ;; found -1 (expected sqrt(1))
                    #f)
                   ((= b 1)
                    ;; !! (previous b)^2=1 and was not 1 or -1
                    (not (zero? i)))
                   ((>= i twos)
                    ;; !! a^(n-1)!=1 mod n
                    )
                   (else
                    (lp (+ i 1) (remainder (* b b) n)))))))))))

;;> Returns true if we can show \var{n} to be composite
;;> using the Miller-Rabin test (i.e., finding a witness \var{a}
;;> where \var{a}^(\var{n}-1)≠1 mod \var{n} or \var{a} reveals
;;> the existence of a 3rd square root of 1 in \b{Z}/(n))
(define (miller-rabin-composite? n)
  (let* ((witness? (miller-rabin-witnesser n))
         ;; Each iteration of Miller Rabin reduces the odds by 1/4, so
         ;; this is a 1 in 2^40 probability of false positive,
         ;; assuming good randomness from SRFI 27 and no bugs, further
         ;; reduced by preliminary sieving.
         (fixed-limit 16)
         (rand-limit (if (< n 341550071728321) fixed-limit 20)))
    (let try ((i 0))
      (and (< i rand-limit)
           (or (witness? (if (< i fixed-limit)
                             (vector-ref prime-table i)
                             (+ (random-integer (- n 3)) 2)))
               (try (+ i 1)))))))

;;> Returns true if \var{n} has a very high probability (enough that
;;> you can assume a false positive will never occur in your lifetime)
;;> of being prime.
(define (probable-prime? n)
  (cond
   ((< n 1) #f)
   (else
    (let ((len (vector-length prime-table)))
      (let lp ((i 0))
        (if (>= i len)
            (not (miller-rabin-composite? n))
            (let ((x (vector-ref prime-table i)))
              (cond
               ((>= x n) (= x n))
               ((zero? (remainder n x)) #f)
               (else (lp (+ i 1)))))))))))

;;> Returns true iff \var{n} is prime.  Uses \scheme{provable-prime?}
;;> for small \var{n}, falling back on \scheme{probable-prime?} for
;;> large values.
(define (prime? n)
  (and (> n 1)
       (if (< n #e1e10)
           (provable-prime? n)
           (probable-prime? n))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Prime iteration and factorization

;;> Returns the nth prime, with 2 being the 0th prime.
(define (nth-prime i)
  (define (by-twos n j)
    (if (prime? n)
        (if (<= j 0) n (by-twos (+ n 2) (- j 1)))
        (by-twos (+ n 2) j)))
  (let ((len (vector-length prime-table)))
    (if (< i len)
        (vector-ref prime-table i)
        (by-twos (+ 2 (vector-ref prime-table (- len 1))) (- i len)))))

;;> Returns the first prime less than or equal to \var{n}, or #f if
;;> there are no such primes.
(define (prime-below n)
  (cond
   ((> n 3)
    (let lp ((n (if (even? n) (- n 1) (- n 2))))
      (if (prime? n) n (lp (- n 2)))))
   ((= n 3)
    2)
   (else
    #f)))

;;> Returns the first prime greater than or equal to \var{n}.  If the
;;> optional \var{limit} is given and not false, returns \scheme{#f}
;;> if no such primes exist below \var{limit}.
(define (prime-above n . o)
  (let ((limit (and (pair? o) (car o))))
    (cond
     ((< n 2)
      2)
     (limit
      (let lp ((n (if (even? n) (+ n 1) (+ n 2))))
        (cond
         ((>= n limit) #f)
         ((prime? n) n)
         (else (lp (+ n 2))))))
     (else
      (let lp ((n (if (even? n) (+ n 1) (+ n 2))))
        (cond
         ((prime? n) n)
         (else (lp (+ n 2)))))))))

;;> Returns the factorization of \var{n} as a monotonically
;;> increasing list of primes.
(define (factor n)
  (when (zero? n)
    (error "cannot factor 0"))
  (factor-twos
   n
   (lambda (b n)
     (let lp ((i 3) (ii 9) (n (abs n))
              (res (append (make-list b 2)
                           (if (negative? n) (list -1) '()))))
       (cond
        ((> ii n)
         (reverse (if (= n 1) res (cons n res))))
        ((zero? (remainder n i))
         (lp i ii (quotient n i) (cons i res)))
        (else
         (lp (+ i 2)
             (+ ii (* (+ i 1) 4))
             n
             res)))))))

;;> Returns the Euler totient function, the number of positive
;;> integers less than \var{n} that are relatively prime to \var{n}.
(define (totient n)
  (let ((limit (exact (ceiling (sqrt n)))))
    (let lp ((i 2) (count 1))
      (cond ((> i limit)
             (if (= count (- i 1))
                 (- n 1)                ; shortcut for prime
                 (let lp ((i i) (count count))
                   (cond ((>= i n) count)
                         ((= 1 (gcd n i)) (lp (+ i 1) (+ count 1)))
                         (else (lp (+ i 1) count))))))
            ((= 1 (gcd n i)) (lp (+ i 1) (+ count 1)))
            (else (lp (+ i 1) count))))))

;;> The aliquot sum s(n), equal to the sum of proper divisors of an
;;> integer n.
(define (aliquot n)
  (let ((limit (+ 1 (quotient n 2))))
    (let lp ((i 2) (sum 1))
      (cond ((> i limit) sum)
            ((zero? (remainder n i)) (lp (+ i 1) (+ sum i)))
            (else (lp (+ i 1) sum))))))

;;> Returns true iff \var{n} is a perfect number, i.e. the sum of its
;;> divisors other than itself equals itself.
(define (perfect? n)
  (and (> n 1) (= n (aliquot n))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Random prime generation

;;> Returns a random prime between \var{lo}, inclusive, and \var{hi},
;;> exclusive.
(define (random-prime lo hi)
  (if (> lo hi)
      (error "bad range: " lo hi))
  (let ((n (bitwise-ior 1 (+ lo (random-integer (- hi lo))))))
    (or (prime-above n hi)
        (prime-above lo n)
        (error "couldn't find prime between: " lo hi))))

;;> Variant of \scheme{random-prime} which ensures the result is
;;> distinct from \var{p}.
(define (random-prime-distinct-from lo hi p)
  (let ((q (random-prime lo hi)))
    (if (= q p)
        (random-prime-distinct-from lo hi p)
        q)))

;;> Returns a random integer less than \var{n} relatively prime to
;;> \var{n}.
(define (random-coprime n)
  (let ((init (+ 2 (random-integer (- n 1)))))
    (let lp ((m init))
      (cond ((>= m n)
             (let lp ((m (- init 1)))
               (cond
                ((<= m 1) (error "couldn't find coprime to " n))
                ((coprime? n m) m)
                (else (lp (- m 1))))))
            ((coprime? n m) m)
            (else (lp (+ m 1)))))))