math.primes: speedup sieve.

basis/math/primes/erato/erato.factor

@@ -1,7 +1,8 @@
 ! Copyright (C) 2009 Samuel Tardieu.
 ! See http://factorcode.org/license.txt for BSD license.
-USING: arrays byte-arrays kernel math math.bitwise math.functions math.order
-math.ranges sequences sequences.private ;
+USING: kernel kernel.private locals math math.bitwise
+math.functions math.order math.ranges sequences
+sequences.private ;
 IN: math.primes.erato
 
 <PRIVATE
@@ -10,34 +11,37 @@ CONSTANT: masks
 { f 128 f f f f f 64 f f f 32 f 16 f f f 8 f 4 f f f 2 f f f f f 1 }
 
 : bit-pos ( n -- byte mask/f )
-    30 /mod masks nth-unsafe ; inline
+    { fixnum } declare
+    30 /mod masks nth-unsafe
+    { maybe{ fixnum } } declare ; inline
 
-: marked-unsafe? ( n arr -- ? )
+: marked-unsafe? ( n sieve -- ? )
     [ bit-pos ] dip swap
-    [ [ nth-unsafe ] [ bitand zero? not ] bi* ] [ 2drop f ] if* ; inline
+    [ [ nth-unsafe ] [ mask zero? not ] bi* ] [ 2drop f ] if* ; inline
 
-: unmark ( n arr -- )
-    [ bit-pos swap ] dip
-    pick [ [ swap unmask ] change-nth-unsafe ] [ 3drop ] if ; inline
+: unmark ( n sieve -- )
+    [ bit-pos swap ] dip pick
+    [ [ swap unmask ] change-nth-unsafe ] [ 3drop ] if ; inline
 
-: upper-bound ( arr -- n ) length 30 * 1 - ; inline
+: upper-bound ( sieve -- n ) length 30 * 1 - ; inline
 
-: unmark-multiples ( i arr -- )
-    2dup marked-unsafe? [
-        [ [ dup sq ] [ upper-bound ] bi* rot <range> ] keep
-        [ unmark ] curry each
-    ] [
-        2drop
-    ] if ; inline
+:: unmark-multiples ( i upper sieve -- )
+    i sieve marked-unsafe? [
+        i sq upper i <range> [ sieve unmark ] each
+    ] when ; inline
 
-: init-sieve ( n -- arr ) 30 /i 1 + 255 <array> >byte-array ; inline
+: init-sieve ( n -- sieve )
+    30 /i 1 + [ 255 ] B{ } replicate-as ; inline
 
 PRIVATE>
 
-: sieve ( n -- arr )
-    init-sieve [ 2 swap upper-bound sqrt [a,b] ] keep
-    [ [ unmark-multiples ] curry each ] keep ;
+:: sieve ( n -- sieve )
+    n integer>fixnum-strict init-sieve :> sieve
+    sieve upper-bound >fixnum :> upper
+    2 upper sqrt [a,b]
+    [ upper sieve unmark-multiples ] each
+    sieve ;
 
-: marked-prime? ( n arr -- ? )
+: marked-prime? ( n sieve -- ? )
     2dup upper-bound 2 swap between? [ bounds-error ] unless
     over { 2 3 5 } member? [ 2drop t ] [ marked-unsafe? ] if ;

basis/math/primes/primes.factor

@@ -8,7 +8,8 @@ IN: math.primes
 
 <PRIVATE
 
-: look-in-bitmap ( n -- ? ) $[ 8999999 sieve ] marked-unsafe? ; inline
+: look-in-bitmap ( n -- ? )
+    $[ 8999999 sieve ] marked-unsafe? ; inline
 
 : (prime?) ( n -- ? )
     dup 8999999 <= [ look-in-bitmap ] [ miller-rabin ] if ;
@@ -26,7 +27,7 @@ IN: math.primes
 : <primes-vector> ( low high -- vector )
     swap [ [ upper-pi ] [ lower-pi ] bi* - >integer
     108 max 10000 min <vector> ] keep
-    3 < [ [ 2 swap push ] keep ] when ;
+    3 < [ 2 suffix! ] when ;
 
 : simple? ( n -- ? ) { [ even? ] [ 3 divisor? ] [ 5 divisor? ] } 1|| ;