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 70a9c10ff5367ff1f3ff356a77f705919f2f2a60..d79eb272b4da7c7333e5cb132b5c5dc08649ce2e 100644 (file)
--- a/basis/math/primes/erato/erato.factor
+++ b/basis/math/primes/erato/erato.factor
@@ -1,25 +1,59 @@
 ! Copyright (C) 2009 Samuel Tardieu.
 ! See http://factorcode.org/license.txt for BSD license.
-USING: bit-arrays kernel math math.functions math.ranges sequences ;
+USING: kernel kernel.private math math.bitwise
+math.functions math.order math.private ranges sequences
+sequences.private ;
 IN: math.primes.erato
 
-: >index ( n -- i )
-    3 - 2 /i ; inline
+<PRIVATE
 
-: index> ( i -- n )
-    2 * 3 + ; inline
+! 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.
 
-: mark-multiples ( i arr -- )
-    [ index> [ sq >index ] keep ] dip
-    [ length 1 - swap <range> f swap ] keep
-    [ set-nth ] curry with each ;
+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 }
 
-: maybe-mark-multiples ( i arr -- )
-    2dup nth [ mark-multiples ] [ 2drop ] if ;
+: bit-pos ( n -- byte mask/f )
+    { fixnum } declare
+    30 /mod masks nth-unsafe
+    { maybe{ fixnum } } declare ; inline
 
-: init-sieve ( n -- arr )
-    >index 1 + <bit-array> dup set-bits ;
+:: marked-unsafe? ( n sieve -- ? )
+    n bit-pos [
+        [ sieve nth-unsafe ] [ mask zero? not ] bi*
+    ] [ drop f ] if* ; inline
 
-: sieve ( n -- arr )
-    [ init-sieve ] [ sqrt >index [0,b] ] bi
-    over [ maybe-mark-multiples ] curry each ; foldable
+:: unmark ( n sieve -- )
+    n bit-pos [
+        swap sieve [ swap unmask ] change-nth-unsafe
+    ] [ drop ] if* ; inline
+
+: upper-bound ( sieve -- n ) length 30 * 1 - ; inline
+
+:: 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
+
+: init-sieve ( n -- sieve )
+    30 /i 1 + [ 255 ] B{ } replicate-as ; inline
+
+PRIVATE>
+
+:: 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 ;