factor: trim using lists

[factor.git] / basis / math / primes / erato / erato.factor

diff --git a/basis/math/primes/erato/erato.factor b/basis/math/primes/erato/erato.factor
index 4df724cfc23bba20f1f0e2229f214d27b82f8ece..d79eb272b4da7c7333e5cb132b5c5dc08649ce2e 100644 (file)
--- a/basis/math/primes/erato/erato.factor
+++ b/basis/math/primes/erato/erato.factor
@@ -1,41 +1,59 @@
 ! 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 math math.bitwise
+math.functions math.order math.private ranges sequences
+sequences.private ;
 IN: math.primes.erato
 
 <PRIVATE
 
-CONSTANT: masks B{ 0 128 0 0 0 0 0 64 0 0 0 32 0 16 0 0 0 8 0 4 0 0 0 2 0 0 0 0 0 1 }
+! This is a compressed Sieve of Eratosthenes that uses the
+! 2-3-5 wheel to check groups of 8 candidates starting with
+! { 1 7 11 13 17 19 23 29 } allowing us to use a byte-array
+! to store each group of booleans in a byte.
 
-: bit-pos ( n -- byte/f mask/f )
-    30 /mod masks nth-unsafe [ drop f f ] when-zero ;
+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 }
 
-: marked-unsafe? ( n arr -- ? )
-    [ bit-pos ] dip swap [ [ nth-unsafe ] [ bitand zero? not ] bi* ] [ 2drop f ] if* ;
+: bit-pos ( n -- byte mask/f )
+    { fixnum } declare
+    30 /mod masks nth-unsafe
+    { maybe{ fixnum } } declare ; inline
 
-: unmark ( n arr -- )
-    [ bit-pos swap ] dip
-    over [ [ swap unmask ] change-nth-unsafe ] [ 3drop ] if ;
+:: marked-unsafe? ( n sieve -- ? )
+    n bit-pos [
+        [ sieve nth-unsafe ] [ mask zero? not ] bi*
+    ] [ drop f ] if* ; inline
 
-: upper-bound ( arr -- n ) length 30 * 1 - ;
+:: unmark ( n sieve -- )
+    n bit-pos [
+        swap sieve [ swap unmask ] change-nth-unsafe
+    ] [ drop ] if* ; inline
 
-: unmark-multiples ( i arr -- )
-    2dup marked-unsafe? [
-        [ [ dup sq ] [ upper-bound ] bi* rot <range> ] keep
-        [ unmark ] curry each
-    ] [
-        2drop
-    ] if ;
+: upper-bound ( sieve -- n ) length 30 * 1 - ; inline
 
-: init-sieve ( n -- arr ) 30 /i 1 + 255 <array> >byte-array ;
+:: unmark-multiples ( i upper sieve -- )
+    i 2 fixnum*fast :> step
+    i i fixnum*fast
+    [ dup upper <= ] [
+        [ sieve unmark ] [ step fixnum+fast ] bi
+    ] while drop ; inline
 
-PRIVATE>
+: init-sieve ( n -- sieve )
+    30 /i 1 + [ 255 ] B{ } replicate-as ; inline
 
-: sieve ( n -- arr )
-    init-sieve [ 2 swap upper-bound sqrt [a,b] ] keep
-    [ [ unmark-multiples ] curry each ] keep ;
+PRIVATE>
 
-: marked-prime? ( n arr -- ? )
+:: sieve ( n -- sieve )
+    n integer>fixnum-strict init-sieve :> sieve
+    sieve upper-bound >fixnum :> upper
+    3 upper sqrt 2 <range> [| i |
+        i sieve marked-unsafe? [
+            i upper sieve unmark-multiples
+        ] when
+    ] each sieve ;
+
+: marked-prime? ( n sieve -- ? )
+    [ integer>fixnum-strict ] dip
     2dup upper-bound 2 swap between? [ bounds-error ] unless
-    over { 2 3 5 } member? [ 2drop t ] [ marked-unsafe? ] if ;
+    over { 2 3 5 } member? [ 2drop t ] [ marked-unsafe?  ] if ;