From 64db5c5bb463b1b222c9ea7317e588e42caf9c68 Mon Sep 17 00:00:00 2001
From: John Benediktsson <mrjbq7@gmail.com>
Date: Tue, 16 Jun 2015 08:26:48 -0700
Subject: [PATCH] math.primes.erato: faster compressed sieve by two
 improvements:

1) fixnum math in unmark-multiples
2) 3 upper sqrt 2 <range> in sieve
---
 basis/math/primes/erato/erato-tests.factor | 12 ++++++++--
 basis/math/primes/erato/erato.factor       | 27 ++++++++++++++--------
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/basis/math/primes/erato/erato-tests.factor b/basis/math/primes/erato/erato-tests.factor
index ff44ec2210..c1efc502db 100644
--- a/basis/math/primes/erato/erato-tests.factor
+++ b/basis/math/primes/erato/erato-tests.factor
@@ -1,5 +1,5 @@
-USING: kernel byte-arrays sequences tools.test ;
-USING: math math.bitwise math.ranges math.primes.erato ;
+USING: fry kernel math math.bitwise math.primes.erato
+math.ranges sequences tools.test ;
 
 [ B{ 255 251 247 126 } ] [ 100 sieve ] unit-test
 [ 1 100 sieve marked-prime? ] [ bounds-error? ] must-fail-with
@@ -14,3 +14,11 @@ USING: math math.bitwise math.ranges math.primes.erato ;
 ! end-point for numbers with all possibilities mod 30. If something
 ! were to go wrong, we'd get a bounds-error.
 [ ] [ 2 100 [a,b] [ dup sieve marked-prime? drop ] each ] unit-test
+
+{ t } [
+    { 2 3 5 7 11 13 } 100 sieve '[ _ marked-prime? ] all?
+] unit-test
+{ t } [
+    { 4 6 8 9 10 12 } 100 sieve '[ _ marked-prime? not ] all?
+] unit-test
+
diff --git a/basis/math/primes/erato/erato.factor b/basis/math/primes/erato/erato.factor
index 60e22db552..6148f5db68 100644
--- a/basis/math/primes/erato/erato.factor
+++ b/basis/math/primes/erato/erato.factor
@@ -1,7 +1,7 @@
 ! Copyright (C) 2009 Samuel Tardieu.
 ! See http://factorcode.org/license.txt for BSD license.
 USING: kernel kernel.private locals math math.bitwise
-math.functions math.order math.ranges sequences
+math.functions math.order math.private math.ranges sequences
 sequences.private ;
 IN: math.primes.erato
 
@@ -15,19 +15,25 @@ CONSTANT: masks
     30 /mod masks nth-unsafe
     { maybe{ fixnum } } declare ; inline
 
-: marked-unsafe? ( n sieve -- ? )
-    [ bit-pos ] dip swap
-    [ [ nth-unsafe ] [ mask zero? not ] bi* ] [ 2drop f ] if* ; inline
+:: marked-unsafe? ( n sieve -- ? )
+    n bit-pos [
+        [ sieve nth-unsafe ] [ mask zero? not ] bi*
+    ] [ drop f ] if* ; inline
 
-: unmark ( n sieve -- )
-    [ bit-pos swap ] dip pick
-    [ [ swap unmask ] change-nth-unsafe ] [ 3drop ] if ; inline
+:: 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 sieve marked-unsafe? [
-        i sq upper i <range> [ sieve unmark ] each
+        i 2 fixnum*fast :> step
+        i i fixnum*fast
+        [ dup upper <= ] [
+            [ sieve unmark ] [ step fixnum+fast ] bi
+        ] while drop
     ] when ; inline
 
 : init-sieve ( n -- sieve )
@@ -38,10 +44,11 @@ PRIVATE>
 :: sieve ( n -- sieve )
     n integer>fixnum-strict init-sieve :> sieve
     sieve upper-bound >fixnum :> upper
-    2 upper sqrt [a,b]
+    3 upper sqrt 2 <range>
     [ upper sieve unmark-multiples ] 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 ;