\ ----------------------------------------------------------------------
\ Prime number finding with the sieve of Erastothenes
\ ----------------------------------------------------------------------
\ $ gforth-fast primes.fs -e '1000 primes bye'
\ $ gforth-fast primes.fs -e '1000000 nthprime bye'
\ ----------------------------------------------------------------------
\ optimizations:
\   - The sieve stores only odd numbers, 2 is treated as a special case.
\   - The sieve can be resized without recalculating the whole sieve.
\   - At different points in in time a location in the sieve will store
\     either 1) a positive number which is a prime, or a composite that
\     have not yet been culled or 2) a negative number whose negation
\     is used as the step size when iterating over the sieve (outside
\     of culling iteration). The step size is optimized so that the
\     negative number directly following a prime indicates the distance
\     to the next prime.
\   - Iterative culling stops once the square of the culling prime
\     is bigger than the biggest number of the segment, smaller composites
\     will already have been culled, as culling primes are used in
\     ascending order, and bigger composites will also be outside the
\     segment.
\   - Using prime number theorem to guess required size of the sieve.
\ ----------------------------------------------------------------------

: primen>n ( n -- n )
	\ ballparking the value of the nth prime
	\ https://en.wikipedia.org/wiki/Prime_number_theorem
	s>f fdup fln f* f>s
;

: 3dup ( a b c -- a b c a b c ) dup 2over rot ;
: 3drop ( a b c -- )            2drop drop ;

\ stored number: 3  5  7  9 11 13 15
\ sieve offset:  0  1  2  3  4  5  6
: number>offset ( n -- n ) 3 - 2/ ;
: offset>number ( n -- n ) 2 * 3 + ;

: count-primes ( addr ub lb -- n )
	\ counts primes in the lb-ub part of sieve pointed to by addr
	0 -rot
	?do   \ addr acc
		over i cells + @ \ addr acc value
		dup 0> if
			drop 1+ -1  \ addr acc+1 -1
		then
		negate
	+loop \ addr acc
	nip
;

: fill-sieve ( addr stop start -- )
    \ fill some memory space with consecutive odd numbers
    ?do
        i offset>number over i cells + !
    loop
    drop
;

: bp>k ( b p -- k )
	\ calculate smallest k = (2N+1), so that sieve[b] <= p*k
	swap offset>number swap \ sieve[b] p
	/mod swap if 1+ then \ k (low bound)
	1 or 3 max \ make sure k is odd and >= 3 ==> k =(2N+1)
;

: bp>i ( b p -- i )
	\ calculate sieve index for p*k, see bp>k for explanation
	tuck bp>k * number>offset
;

: cull-multiples ( addr hb lb p -- f )
	\ cull cells in the range [lb,hb) that correspond
	\ to natural>1 multiples of p

	\ if p^2 > sieve[hb-1] then there are
	\ no nonprime unculled values left
	dup dup * 3 pick 1- \ addr hb lb p p^2 hb-1
	offset>number > if
		2drop 2drop 1 exit
	then

	     tuck bp>i \ hb p li | convert array bound range
	-rot tuck bp>i \ li p hi | so they align with multiples
	 rot           \ p hi li | of p
	?do       \ addr p
		over i cells + -1 swap !
	dup +loop \ addr p
	2drop
	0
;

: iterate-culling ( addr hb lb -- )
	\ for each cell in the range [0,hb) in the sieve:
	\     if positive, cull multiples>1 of the stored value on indexes [lb,hb)
	over 0 do              \ addr lb hb
		2 pick i cells + @ \ addr lb hb p|neg
		dup 0> if
			>r 3dup r> cull-multiples if leave then
			-1
		then
		negate
	+loop
	3drop
;

: prime-distance ( addr hb lb -- n )
	\ distance to next prime from lb
	dup -rot \ addr lb hb lb
	dup 1+ -rot   \ addr lb lb+1 hb lb
	1+ ?do \ addr lb last-i
		drop
		over i cells + @
		i swap           \ addr lb i p|neg
		0> if leave then
	loop  \ addr lb last-i
	- nip \ -distance
;

: annotate-holes ( addr hb lb -- )
	\ for cells in the range [lb,hb), if the cell is culled
	\ calculate and store the distance to the next unculled cell

	\ roll down lb to last prime if lb>0
	dup 0> if
		0 swap do \ addr hb
			over i cells + @ \ addr hb value
			0> if
				i leave
			then
		-1 +loop
	then \ addr hb lb

	over swap do \ addr hb
		over i cells + @ \ addr hb p|neg
		0< if \ addr hb
			2dup i prime-distance  \ addr hb -distance
			dup 3 pick i cells + ! \ addr hb -distance
		else
			-1
		then
		negate
	+loop \ addr hb
	2drop
;

: print-contents ( addr sz -- )
	0 ?do dup i cells + @ . loop drop
;

: resize-sieve ( addr oldsz newsz -- addr )
	rot over cells resize throw \ oldsz newsz addr
	-rot swap                   \ addr newsz oldsz
	3dup fill-sieve
	3dup iterate-culling
	3dup annotate-holes
	2drop
;

: make-sieve ( sz -- addr )
	0 0 rot resize-sieve
;

: print-n-primes ( addr ub n -- )
	-rot \ n addr ub
	0 do \ n addr
		dup i cells + @ \ n addr p|neg
		dup 0> if \ n addr p
			. swap 1- \ addr n-1
			dup 0= if 2drop unloop exit then
			swap -1
		then
		negate
	+loop
	2drop
;

: n-primes-sieve ( n -- addr sz )
	\ construct a sieve containing primes above 2
	dup 2 < if drop 0 0 exit then
	dup >r
	primen>n number>offset 50 max \ r:c minsize
	dup make-sieve swap 0            \ r:c addr ub lb=0

	\ loop while the counter is > 0
	\     decrement the counter by the number of primes
	\     generated in the last iteration
	\     grow the sieve as needed

	begin                  \ r:c   addr ub lb
		3dup count-primes  \ r:c   addr ub lb n    | count primes in [lb,ub)
		r> swap - >r       \ r:c-n addr ub lb      | decrement counter
	r@ 0> while
		drop dup 15 10 */     \ r:c addr ub ub*1.5 | growth factor 1.5
		2>r 2r@
		resize-sieve 2r> swap \ r:c addr ub lb     | ub, lb = ub*1.5, ub
	repeat                 \ r:c   addr ub lb
	rdrop drop
;

: primes ( n -- )
	\ print n number of primes
	dup 1 < if drop exit then
	2 .
	dup 1 = if drop exit then
	1-
	dup n-primes-sieve   \ n addr sz
	0 ?do
		dup i cells + @
		dup 0> if
			.
			swap 1- swap \ n-1 addr
			over 0= if leave then
			-1
		then
		negate
	+loop
	free throw drop
	cr
;

: nthprime ( n -- )
	\ print the nth prime
	dup 1 < if drop exit then
	dup 1 = if drop 2 . cr exit then
	1-
	dup n-primes-sieve  \ n addr sz
	0 ?do               \ n addr
		dup i cells + @ \ n addr p|neg
		dup 0> if
			rot 1-      \ addr p n-1
			dup 0= if swap . swap leave then
			nip swap
			-1
		then
		negate
	+loop
	free throw drop
	cr
;