basis/math/primes/erato/erato.factor

   1 ! Copyright (C) 2009 Samuel Tardieu.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: kernel kernel.private math math.bitwise
   4 math.functions math.order math.private ranges sequences
   5 sequences.private ;
   6 IN: math.primes.erato
   7
   8 <PRIVATE
   9
  10 ! This is a compressed Sieve of Eratosthenes that uses the
  11 ! 2-3-5 wheel to check groups of 8 candidates starting with
  12 ! { 1 7 11 13 17 19 23 29 } allowing us to use a byte-array
  13 ! to store each group of booleans in a byte.
  14
  15 CONSTANT: masks
  16 { 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 }
  17
  18 : bit-pos ( n -- byte mask/f )
  19     { fixnum } declare
  20     30 /mod masks nth-unsafe
  21     { maybe{ fixnum } } declare ; inline
  22
  23 :: marked-unsafe? ( n sieve -- ? )
  24     n bit-pos [
  25         [ sieve nth-unsafe ] [ mask zero? not ] bi*
  26     ] [ drop f ] if* ; inline
  27
  28 :: unmark ( n sieve -- )
  29     n bit-pos [
  30         swap sieve [ swap unmask ] change-nth-unsafe
  31     ] [ drop ] if* ; inline
  32
  33 : upper-bound ( sieve -- n ) length 30 * 1 - ; inline
  34
  35 :: unmark-multiples ( i upper sieve -- )
  36     i 2 fixnum*fast :> step
  37     i i fixnum*fast
  38     [ dup upper <= ] [
  39         [ sieve unmark ] [ step fixnum+fast ] bi
  40     ] while drop ; inline
  41
  42 : init-sieve ( n -- sieve )
  43     30 /i 1 + [ 255 ] B{ } replicate-as ; inline
  44
  45 PRIVATE>
  46
  47 :: sieve ( n -- sieve )
  48     n integer>fixnum-strict init-sieve :> sieve
  49     sieve upper-bound >fixnum :> upper
  50     3 upper sqrt 2 <range> [| i |
  51         i sieve marked-unsafe? [
  52             i upper sieve unmark-multiples
  53         ] when
  54     ] each sieve ;
  55
  56 : marked-prime? ( n sieve -- ? )
  57     [ integer>fixnum-strict ] dip
  58     2dup upper-bound 2 swap between? [ bounds-error ] unless
  59     over { 2 3 5 } member? [ 2drop t ] [ marked-unsafe?  ] if ;