\ ----------------------------------------------------------------------
\ Prime number finding with the sieve of Erastothenes
\ ----------------------------------------------------------------------
\ example use: gforth nprimes.fs -e '1000 n-first-primes bye'
\ ----------------------------------------------------------------------
\ optimizations:
\   - sieve stores only odd numbers, 2 is treated as a special case
\   - sieve can be resized without recalculating the whole sieve
\   - the would-be-zero of a non-prime following a following a prime
\     stores the size of the hole as a negative number, making
\     iteration faster when trying to extract primes only
\   - 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 ;

\ prime 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 + ;

: neg>loop-increment
	\ neg->pos, 0->1
	negate 1 max
;

: 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 delta
		else
			neg>loop-increment
		then
	+loop \ addr acc
	nip
;

: bp>k ( b p -- k )
	\ for prime p, and array bound b, find lowest integer k
	\ so that p*k >= bn , where k = 2n+1 (natural n), and bn is the number at b
	\ k >= ceil(bn/p) = floor(bn/p) + ((bn%p)==0)?0:1
	swap offset>number swap   \ bn p
	/mod swap if 1+ then \ k (low bound)
	1 or  \ make sure k is odd (a)
	3 max \ make sure k >= 3   (b) ...(a) & (b) => k = 2n+1 (natural n)
;

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

: prune-multiples ( addr hb lb p -- )
	\ set cells that are natural>1 multiples of p to zero
	\ but only for cells with lb <= index < hb
	\ nomenclature: lb/hb array bounds
	\               lk/hk k calculated from array bounds
	\               li/hi indexes calculated from k
	tuck bp>k              \ addr hb p lk
	over * number>offset   \ addr hb p li
	rot 2 pick bp>k        \ addr p li hk
	2 pick * number>offset \ addr p li hi
	swap ?do  \ addr p
		over i cells + 0 swap !
	dup +loop \ addr p
	2drop
;

: iterate-pruning ( addr hb lb -- )
	\ for each cell: 1) if >0: prune multiples
	\                   in the lb-hb index range of it
	\                   go to next cell
	\                2) if <0: make the number positive and jump that many steps forward
	\                3) if =0: go to the next cell
	over 0 do       \ addr lb hb
		3dup 2 pick \ addr lb hb addr lb hb addr
		i cells + @ \ addr lb hb addr lb hb p|0
		dup 0> if
			prune-multiples
			1
		else
			>r 3drop r>
			neg>loop-increment
		then
	+loop
	2drop
	drop
;

: annotate-hole-sizes ( addr hb lb -- )
	\ for a part of an array, filled with primes and zeroes
	\     for each zero that follows a positive number
	\         calculate the array distance to the next
	\         positive number, and store it as a negative
	\         number instead of zero
	dup 0> if \ roll down lb to last prime if it's >0
		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 value
		0= if \ addr hb
			1 -rot       \ delta addr hb
			dup i 1+ ?do \ delta addr hb
				over i cells + @
				0> if
					rot drop i j - -rot \ new-delta addr hb
					leave
				then
			loop         \ delta addr hb
			rot dup negate     \ addr hb delta -delta
			i cells 4 pick + ! \ addr hb delta
		else
			1                  \ addr hb delta
		then
	+loop \ addr hb
	2drop
;

: resize-sieve ( addr oldsz newsz -- addr )
	rot over cells resize throw \ oldsz newsz addr
	-rot swap                   \ addr newsz oldsz
	3dup fill-sieve
	3dup iterate-pruning
	3dup annotate-hole-sizes
	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 + @
		dup 0> if \ n addr v
			. swap 1- dup 0= \ addr n f
			if 2drop unloop exit then
			swap 1
		else
			neg>loop-increment
		then
	+loop
	2drop
;

: n-first-primes ( n -- )
	\ handle special cases
	dup 1 < if drop    exit then
	2 .
	dup 1 = if drop cr exit then

	\ save wanted number of primes
	\ as just that, and a decrementing counter
	\ on the returnstack, 1- because 2 handled
	1-
	dup >r dup >r          \ r:n r:a n

	\ gauge required sieve size but never make
	\ a sieve smaller than 50 entries
	primen>n number>offset 50 max \ r:n r:a minsize
	dup make-sieve swap 0  \ r:n r:a addr ub lb=0
	begin                  \ r:n r:a addr ub lb
		3dup count-primes  \ r:n r:a addr ub lb count
		r> swap - >r       \ r:n r:a-count addr ub lb
	r@ 0> while
		drop dup 15 10 */     \ r:n r:a addr nlb nub, where nlb=oub, nub=oub*1.5
		3dup resize-sieve     \ r:n r:a addr1 ub lb addr2
		2>r nip 2r> -rot swap \ r:n r:a addr2 ub lb
	repeat
	rdrop drop r>  \ addr ub n
	2 pick >r      \ r:addr addr ub n
	print-n-primes \ r:addr
	r> free throw
	cr
;