Make small primes faster.

Big improvement is shrinking the bit vector to the right size.
Other improvements include full fixnum arithmetic and less mutation.

pkgs/math-pkgs/math-lib/math/private/number-theory/small-primes.rkt

@@ -10,88 +10,77 @@
          *SMALL-PRIME-LIMIT*)
 
 ; The moduli mod 60 that 2, 3 and 5 do not divide are:
+(: non-235 (Listof Positive-Byte))
 (define non-235 '(1 7 11 13 17 19 23 29 31 37 41 43 47 49 53 59))
-; The differences of these numbers are:
-(define deltas  '( 6 4  2  4  2  4  6  2  6  4  2  4  2  4  6  2))
 ; Note that there are exactly 16 of these moduli, so they fit in a u16.
 ; That is, a single u16 can represent a block of 60 numbers.
 
-(define mod60->bits (make-vector 60 (cast #f (U #f Integer))))
-(for ([x (in-list non-235)]
+(: mod60->bits (Vectorof (U #f Byte)))
+(define mod60->bits (make-vector 60 #f))
+(for ([x : Positive-Byte (in-list non-235)]
       [b (in-naturals)])
-  (vector-set! mod60->bits x b))
+  (vector-set! mod60->bits x (assert b byte?)))
 
-(define-syntax-rule (mod60->bit m) (vector-ref mod60->bits m))
+(: mod60->bit (Byte -> (U #f Byte)))
+(define (mod60->bit m) (vector-ref mod60->bits m))
 
+(: *number-of-groups* Positive-Fixnum)
 (define *number-of-groups* 17000) ; each group contain 16 numbers
-(define *SMALL-PRIME-LIMIT* (- (* 60 *number-of-groups*) 1))
+(: *SMALL-PRIME-LIMIT* Nonnegative-Fixnum)
+(define *SMALL-PRIME-LIMIT* (assert (- (* 60 *number-of-groups*) 1) fixnum?))
 
-; primes holds (* 60 *number-of-groups*) bits each
-; representing a number not congruent to 2, 3, 5
-(define primes (make-bit-vector (* 60 *number-of-groups*) #t))
+; primes holds (* 16 *number-of-groups*) bits
+; each representing a number not congruent to 2, 3, 5
+(define primes (make-bit-vector (* (length non-235) *number-of-groups*) #t))
 
-(define: (set-bit! [x : Integer]) : Void
-  (define-values (q r) (quotient/remainder x 60))
-  (define b (mod60->bit r))
-  (when b (bit-vector-set! primes (+ (* q 16) b) #t)))
-
-(define: (clear-bit! [x : Integer]) : Void
-  (define-values (q r) (quotient/remainder x 60))
-  (define b (mod60->bit r))
+(define: (clear-bit! [x : Nonnegative-Fixnum]) : Void
+  (define q (quotient x 60))
+  (define b (mod60->bit (remainder x 60)))
   (when b (bit-vector-set! primes (+ (* q 16) b) #f)))
 
-(define: (bit [x : Integer]) : Boolean
-  (define-values (q r) (quotient/remainder x 60))
+(define: (inner-bit [q : Nonnegative-Fixnum] [r : Byte]) : Boolean
   (define b (mod60->bit r))
   (if b 
       (bit-vector-ref primes (+ (* q 16) b))
       #f))
 
-(clear-bit! 1) ; 1 is not prime
+(define: (bit [x : Nonnegative-Fixnum]) : Boolean
+  (define q (quotient x 60))
+  (define r (remainder x 60))
+  (inner-bit q r))
 
-(define: (mark-composites [x : Integer]) : Void
-  ; x is prime => mark 2*x, 3*x, 4*x, 5*x, 6*x, 7*x, ... as prime
+
+(define: (mark-composites [x : Nonnegative-Fixnum]) : Void
+  ; x is prime => mark n*x as prime for all n
   ; Well 2*x, 3*x, 4*x, 5*x, 6*x are not in our table,
-  ; so the first number to mark is 7*x .
-  ; Use the deltas to figure out which to mark.  
-  (define y x)
-  (define delta*x 0)
-  (let loop ([ds deltas])
-    ; (for ([delta (in-cycle deltas)] ...
-    (when (empty? ds) (set! ds deltas))
-    (let ([delta (car ds)])
-      (set! delta*x (* delta x))
-      (cond
-        [(> y (- *SMALL-PRIME-LIMIT* delta*x))
-         (void)]
-        [else         
-         (set! y (+ y delta*x))
+  ; so only mark the multiples that are not divisible by 2, 3, or 5.
+  (let/ec: exit : Void
+    (let: loop : Void ([a : Nonnegative-Fixnum 0])
+      (define y-base (* a 60 x))
+      (for: ([d : Positive-Byte (in-list non-235)])
+        (define y (assert (+ y-base (* d x)) fixnum?))
+        (when (not (= y x))
+          (if (<= y *SMALL-PRIME-LIMIT*)
               (clear-bit! y)
-         (loop (cdr ds))]))))
+              (exit (void)))))
+      (loop (assert (add1 a) fixnum?)))))
 
 (define: sieve-done? : Boolean  #f)
 
 (define: (sieve) : Void
-  (define x 1)
-  (let loop ([ds deltas])
-    ; (for ([delta (in-cycle deltas)] ...
-    (when (empty? ds) (set! ds deltas))
-    (let ([delta (car ds)])
-      (cond
-        [(> (* x x) (- *SMALL-PRIME-LIMIT* delta))
-         (void)]
-        [else
-         ; x runs through all numbers incongruent to 2, 3 and 5
-         (set! x (+ x delta))
-         (when (bit x) ; x is prime
-           (mark-composites x))
-         (loop (cdr ds))]))))
+  (clear-bit! 1) ; 1 is not prime
+  (let/ec: exit : Void
+    (let: loop : Void ([q : Nonnegative-Fixnum 0])
+      (for: ([r : Positive-Byte (in-list non-235)])
+        (define x (assert (+ (* q 60) r) fixnum?))
+        (when (> (* x x) *SMALL-PRIME-LIMIT*)
+          (exit (void)))
+        (when (inner-bit q r) ; x is prime
+          (mark-composites x)))
+      (loop (assert (add1 q) fixnum?)))))
 
-(define: (small-prime? [x : Integer]) : Boolean
+(define: (small-prime? [x : Nonnegative-Fixnum]) : Boolean
   (unless sieve-done?
     (sieve)
     (set! sieve-done? #t))
-  
-  (or (= x 2) (= x 3) (= x 5)
-      (and (mod60->bit (modulo x 60))
-           (bit x))))
+  (or (= x 2) (= x 3) (= x 5) (bit x)))