; solovay-strassen primality test

(define (square x) (* x x))

(define (expm b e m)
  (define (m* x y) (modulo (* x y) m))
  (cond ((zero? e) 1)
        ((even? e) (expm (m* b b) (/ e 2) m))
        (else (m* b (expm (m* b b) (/ (- e 1) 2) m)))))

(define rand #f)
(define randint #f)
(let ((two31 #x80000000) (a (make-vector 56 -1)) (fptr #f))
  (define (mod-diff x y) (modulo (- x y) two31)) ; generic version
  ; (define (mod-diff x y) (logand (- x y) #x7FFFFFFF)) ; fast version
  (define (flip-cycle)
    (do ((ii 1 (+ ii 1)) (jj 32 (+ jj 1))) ((< 55 jj))
      (vector-set! a ii (mod-diff (vector-ref a ii) (vector-ref a jj))))
    (do ((ii 25 (+ ii 1)) (jj 1 (+ jj 1))) ((< 55 ii))
      (vector-set! a ii (mod-diff (vector-ref a ii) (vector-ref a jj))))
    (set! fptr 54) (vector-ref a 55))
  (define (init-rand seed)
    (let* ((seed (mod-diff seed 0)) (prev seed) (next 1))
      (vector-set! a 55 prev)
      (do ((i 21 (modulo (+ i 21) 55))) ((zero? i))
        (vector-set! a i next) (set! next (mod-diff prev next))
        (set! seed (+ (quotient seed 2) (if (odd? seed) #x40000000 0)))
        (set! next (mod-diff next seed)) (set! prev (vector-ref a i)))
      (flip-cycle) (flip-cycle) (flip-cycle) (flip-cycle) (flip-cycle)))
  (define (next-rand)
    (if (negative? (vector-ref a fptr)) (flip-cycle)
      (let ((next (vector-ref a fptr))) (set! fptr (- fptr 1)) next)))
  (define (unif-rand m)
    (let ((t (- two31 (modulo two31 m))))
      (let loop ((r (next-rand)))
        (if (<= t r) (loop (next-rand)) (modulo r m)))))
  (init-rand 19380110) ; happy birthday donald e knuth
  (set! rand (lambda seed
    (cond ((null? seed) (/ (next-rand) two31))
          ((eq? (car seed) 'get) (cons fptr (vector->list a)))
          ((eq? (car seed) 'set) (set! fptr (caadr seed))
                                 (set! a (list->vector (cdadr seed))))
          (else (/ (init-rand (modulo (numerator
                  (inexact->exact (car seed))) two31)) two31)))))
  (set! randint (lambda args
    (cond ((null? (cdr args))
            (if (< (car args) two31) (unif-rand (car args))
              (floor (* (next-rand) (car args)))))
          ((< (car args) (cadr args))
            (let ((span (- (cadr args) (car args))))
              (+ (car args)
                 (if (< span two31) (unif-rand span)
                   (floor (* (next-rand) span))))))
          (else (let ((span (- (car args) (cadr args))))
                  (- (car args)
                     (if (< span two31) (unif-rand span)
                       (floor (* (next-rand) span))))))))))

(define (jacobi m n) ; also legendre, but not kronecker
  (cond ((<= n m) (jacobi (modulo m n) n))
        ((zero? m) 0) ((= m 1) 1)
        ((= m 2) (if (even? n) 0 (if (member (modulo n 8) '(1 7)) 1 -1)))
        ((even? m) (* (jacobi (/ m 2) n) (jacobi 2 n)))
        ((and (= (modulo m 4) 3) (= (modulo n 4) 3)) (- (jacobi n m))) ; can this be shortened?
        (else (jacobi n m))))

(define (primes n)
  (let* ((m (quotient (- n 1) 2)) (sieve (make-vector m #t)))
    (let loop ((i 0) (p 3) (ps (list 2)))
      (cond ((= m i) (reverse ps))
            ((vector-ref sieve i)
              (do ((j (/ (- (* p p) 3) 2) (+ j p))) ((<= m j))
                (vector-set! sieve j #f))
              (loop (+ i 1) (+ p 2) (cons p ps)))
            (else (loop (+ i 1) (+ p 2) ps))))))

(define (ss-prime? n)
  (if (not (integer? n)) (error 'ss-prime? "must be integer")
    (if (< n 2) #f (if (zero? (modulo n 2)) (= n 2) ; n is even
      (let loop ((k 40) (a (randint 1 n)))
        (if (zero? k) #t ; probably prime
          (let ((j (jacobi a n)) (x (expm a (/ (- n 1) 2) n)))
            (if (or (zero? j) (not (= (modulo j n) x))) #f ; composite
              (loop (- k 1) (randint 1 n))))))))))

(define (ss-proof? n)
  (if (not (integer? n)) (error 'ss-proof? "must be integer")
    (if (< n 2) #f (if (zero? (modulo n 2)) (= n 2) ; n is even
      (let loop ((as (primes (min (- n 1) (inexact->exact (floor (* 2 (square (log n)))))))))
        (if (null? as) #t ; prime on the erh
          (let ((j (jacobi (car as) n)) (x (expm (car as) (/ (- n 1) 2) n)))
            (if (or (zero? j) (not (= (modulo j n) x))) #f ; composite
              (loop (cdr as))))))))))

(define-syntax assert
  (syntax-rules ()
    ((assert expr result)
      (if (not (equal? expr result))
          (for-each display `(
            #\newline "failed assertion:" #\newline
            expr #\newline "expected: " ,result
            #\newline "returned: " ,expr #\newline))))))

(let ((ps (primes 10000)))
  (do ((n 2 (+ n 1))) ((= n 10000))
    (assert (if (member n ps) #t #f) (ss-prime? n))
    (assert (if (member n ps) #t #f) (ss-proof? n))))