-- |-- Module: Math.NumberTheory.Primes.Sieve.Eratosthenes-- Copyright: (c) 2011 Daniel Fischer-- Licence: MIT-- Maintainer: Daniel Fischer <daniel.is.fischer@googlemail.com>-- Stability: Provisional-- Portability: Non-portable (GHC extensions)---- Sieve--{-# LANGUAGE CPP, BangPatterns #-}#if __GLASGOW_HASKELL__ >= 700{-# OPTIONS_GHC -fspec-constr-count=6 #-}#endif{-# OPTIONS_HADDOCK hide #-}moduleMath.NumberTheory.Primes.Sieve.Eratosthenes(primes,sieveFrom,psieveFrom,PrimeSieve(..),psieveList,primeList,primeSieve,nthPrimeCt,countFromTo,countAll,countToNth,sieveBytes,sieveBits,sieveWords,sieveRange,sieveTo)where#include "MachDeps.h"importControl.Monad.STimportData.Array.BaseimportData.Array.SThiding(unsafeFreeze,unsafeThaw,castSTUArray)importControl.Monad(when)importData.BitsimportData.WordimportMath.NumberTheory.Powers.Squares(integerSquareRoot)importMath.NumberTheory.UtilsimportMath.NumberTheory.Primes.Counting.ApproximateimportMath.NumberTheory.Primes.Sieve.Indexing#define IX_MASK 0xFFFFF#define IX_BITS 20#define IX_J_MASK 0x7FFFFF#define IX_J_BITS 23#define J_MASK 7#define J_BITS 3#define SIEVE_KB 128-- Sieve in 128K chunks.-- Large enough to get something done per chunk-- and hopefully small enough to fit in the cache.sieveBytes::IntsieveBytes=SIEVE_KB*1024-- Number of bits per chunk.sieveBits::IntsieveBits=8*sieveBytes-- Last index of chunk.lastIndex::IntlastIndex=sieveBits-1-- Range of a chunk.sieveRange::IntsieveRange=30*sieveBytessieveWords::IntsieveWords=sieveBytes`quot`SIZEOF_HSWORD#if SIZEOF_HSWORD == 8typeCacheWord=Word#define RMASK 63#define WSHFT 6#define TOPB 32#define TOPM 0xFFFFFFFF#elsetypeCacheWord=Word64#define RMASK 31#define WSHFT 5#define TOPB 16#define TOPM 0xFFFF#endif-- | Compact store of primality flags.dataPrimeSieve=PS!Integer{-# UNPACK #-}!(UArrayIntBool)-- | Sieve primes up to (and including) a bound.-- For small enough bounds, this is more efficient than-- using the segmented sieve.---- Since arrays are 'Int'-indexed, overflow occurs when the sieve-- size comes near @'maxBound' :: 'Int'@, that corresponds to an-- upper bound near @15/8*'maxBound'@. On @32@-bit systems, that-- is often within memory limits, so don't give bounds larger than-- @8*10^9@ there.primeSieve::Integer->PrimeSieveprimeSievebound=PS0(runSTUArray$sieveTobound)-- | Generate a list of primes for consumption from a-- 'PrimeSieve'.primeList::PrimeSieve->[Integer]primeList(PS0bs)=2:3:5:[toPrimi|let(lo,hi)=boundsbs,i<-[lo..hi],unsafeAtbsi]primeList(PSvObs)=[vO+toPrimi|let(lo,hi)=boundsbs,i<-[lo..hi],unsafeAtbsi]-- | List of primes.-- Since the sieve uses unboxed arrays, overflow occurs at some point.-- On 64-bit systems, that point is beyond the memory limits, on-- 32-bit systems, it is at about @1.7*10^18@.primes::[Integer]primes=2:3:5:concat[[vO+toPrimi|i<-[0..li],unsafeAtbsi]|PSvObs<-psieveList,let(_,li)=boundsbs]-- | List of primes in the form of a list of 'PrimeSieve's, more compact than-- 'primes', thus it may be better to use @'psieveList' >>= 'primeList'@-- than keeping the list of primes alive during the entire run.psieveList::[PrimeSieve]psieveList=makeSievesplimsqlim00cachewhereplim=4801-- prime #647, 644 of them to usesqlim=plim*plimcache=runSTUArray$dosieve<-sieveTo4801new<-unsafeNewArray_(0,1287)::STs(STUArraysIntCacheWord)letfilljindx|1279<indx=returnnew-- index of 4801 = 159*30 + 31 ~> 159*8+7|otherwise=dop<-unsafeReadsieveindxifpthendolet!i=indx.&.J_MASKk=indx`shiftR`J_BITSstrt1=(k*(30*k+2*rhoi)+bytei)`shiftL`J_BITS+fromIntegral(idxi)!strt=fromIntegral(strt1.&.IX_MASK)!skip=fromIntegral(strt1`shiftR`IX_BITS)!ixes=fromIntegralindx`shiftL`IX_J_BITS+strt`shiftL`J_BITS+fromIntegraliunsafeWritenewjskipunsafeWritenew(j+1)ixesfill(j+2)(indx+1)elsefillj(indx+1)fill00makeSieves::Integer->Integer->Integer->Integer->UArrayIntCacheWord->[PrimeSieve]makeSievesplimsqlimbitOffvalOffcache|valOff'<sqlim=let(nc,bs)=runST$docch<-unsafeThawcache::STs(STUArraysIntCacheWord)bs0<-slicecchfcch<-unsafeFreezecchfbs0<-unsafeFreezebs0return(fcch,fbs0)inPSvalOffbs:makeSievesplimsqlimbitOff'valOff'nc|otherwise=letplim'=plim+4800sqlim'=plim'*plim'(nc,bs)=runST$docch<-growCachebitOffplimcachebs0<-slicecchfcch<-unsafeFreezecchfbs0<-unsafeFreezebs0return(fcch,fbs0)inPSvalOffbs:makeSievesplim'sqlim'bitOff'valOff'ncwherevalOff'=valOff+fromIntegralsieveRangebitOff'=bitOff+fromIntegralsieveBitsslice::STUArraysIntCacheWord->STs(STUArraysIntBool)slicecache=dohi<-snd`fmap`getBoundscachesieve<-newArray(0,lastIndex)Truelettreatpr|hi<pr=returnsieve|otherwise=dow<-unsafeReadcacheprifw/=0thenunsafeWritecachepr(w-1)elsedoixes<-unsafeReadcache(pr+1)let!stj=fromIntegralixes.&.IX_J_MASK-- position of multiple and index of cofactor!ixw=fromIntegral(ixes`shiftR`IX_J_BITS)-- prime data, up to 41 bits!i=ixw.&.J_MASK!k=ixw-i-- On 32-bits, k > 44717396 means overflow is possible in tick!o=i`shiftL`J_BITS!j=stj.&.J_MASK-- index of cofactor!s=stj`shiftR`J_BITS-- index of first multiple to tick off(n,u)<-tickkojslet!skip=fromIntegral(n`shiftR`IX_BITS)!strt=fromIntegral(n.&.IX_MASK)unsafeWritecacheprskipunsafeWritecache(pr+1)((ixes.&.complementIX_J_MASK).|.strt`shiftL`J_BITS.|.fromIntegralu)treat(pr+2)tickstpoffjix|lastIndex<ix=return(ix-sieveBits,j)|otherwise=dop<-unsafeReadsieveixwhenp(unsafeWritesieveixFalse)tickstpoff((j+1).&.J_MASK)(ix+stp*deltaj+tau(off+j))treat0-- | Sieve up to bound in one go.sieveTo::Integer->STs(STUArraysIntBool)sieveTobound=arrwhere(bytes,lidx)=idxPrbound!mxidx=8*bytes+lidxmxval::Integermxval=30*fromIntegralbytes+fromIntegral(rholidx)!mxsve=integerSquareRootmxval(kr,r)=idxPrmxsve!svbd=8*kr+rarr=doar<-newArray(0,mxidx)Trueletstartki=8*(k*(30*k+2*rhoi)+bytei)+idxitickstpoffjix|mxidx<ix=return()|otherwise=dop<-unsafeReadarixwhenp(unsafeWritearixFalse)tickstpoff((j+1).&.J_MASK)(ix+stp*deltaj+tau(off+j))siftix|svbd<ix=returnar|otherwise=dop<-unsafeReadarixwhenp(doleti=ix.&.J_MASKk=ix`shiftR`J_BITS!off=i`shiftL`J_BITS!stp=ix-itickstpoffi(startki))sift(ix+1)sift0growCache::Integer->Integer->UArrayIntCacheWord->STs(STUArraysIntCacheWord)growCacheoffsetplimold=dolet(_,num)=boundsold(bt,ix)=idxPrplim!start=8*bt+ix+1!nlim=plim+4800sieve<-sieveTonlim-- Implement SieveFromTo for this, it's pretty wasteful when nlim isn't(_,hi)<-getBoundssieve-- very small anymoremore<-countFromToWdstarthisievenew<-unsafeNewArray_(0,num+2*more)::STs(STUArraysIntCacheWord)letcopyi|num<i=return()|otherwise=dounsafeWritenewi(old`unsafeAt`i)copy(i+1)copy0letfilljindx|hi<indx=returnnew|otherwise=dop<-unsafeReadsieveindxifpthendolet!i=indx.&.J_MASKk::Integerk=fromIntegral(indx`shiftR`J_BITS)strt0=((k*(30*k+fromIntegral(2*rhoi))+fromIntegral(bytei))`shiftL`J_BITS)+fromIntegral(idxi)strt1=strt0-offset!strt=fromIntegralstrt1.&.IX_MASK!skip=fromIntegral(strt1`shiftR`IX_BITS)!ixes=fromIntegralindx`shiftL`IX_J_BITS.|.strt`shiftL`J_BITS.|.fromIntegraliunsafeWritenewjskipunsafeWritenew(j+1)ixesfill(j+2)(indx+1)elsefillj(indx+1)fill(num+1)start-- Danger: relies on start and end being the first resp. last-- index in a Word-- Do not use except in growCache and psieveFrom{-# INLINE countFromToWd #-}countFromToWd::Int->Int->STUArraysIntBool->STsIntcountFromToWdstartendba=dowa<-(castSTUArray::STUArraysIntBool->STs(STUArraysIntWord))balet!sb=start`shiftR`WSHFT!eb=end`shiftR`WSHFTcount!acci|eb<i=returnacc|otherwise=dow<-unsafeReadwaicount(acc+bitCountWordw)(i+1)count0sb-- count set bits between two indices (inclusive)-- start and end must both be valid indices and start <= endcountFromTo::Int->Int->STUArraysIntBool->STsIntcountFromTostartendba=dowa<-(castSTUArray::STUArraysIntBool->STs(STUArraysIntWord))balet!sb=start`shiftR`WSHFT!si=start.&.RMASK!eb=end`shiftR`WSHFT!ei=end.&.RMASKcount!acci|i==eb=dow<-unsafeReadwaireturn(acc+bitCountWord(w`shiftL`(RMASK-ei)))|otherwise=dow<-unsafeReadwaicount(acc+bitCountWordw)(i+1)ifsb<ebthendow<-unsafeReadwasbcount(bitCountWord(w`shiftR`si))(sb+1)elsedow<-unsafeReadwasblet!w1=w`shiftR`sireturn(bitCountWord(w1`shiftL`(RMASK-ei+si)))-- | @'sieveFrom' n@ creates the list of primes not less than @n@.sieveFrom::Integer->[Integer]sieveFromn=casepsieveFromnofps->dropWhile(<n)(ps>>=primeList)-- | @'psieveFrom' n@ creates the list of 'PrimeSieve's starting roughly-- at @n@. Due to the organisation of the sieve, the list may contain-- a few primes less than @n@.-- This form uses less memory than @['Integer']@, hence it may be preferable-- to use this if it is to be reused.psieveFrom::Integer->[PrimeSieve]psieveFromn=makeSievesplimsqlimbitOffvalOffcachewherek0=max0(n-7)`quot`30valOff=30*k0bitOff=8*k0start=valOff+7ssr=integerSquareRoot(start-1)+1end1=start-6+fromIntegralsieveRangeplim0=integerSquareRootend1plim=plim0+4801-(plim0`rem`4800)sqlim=plim*plimcache=runSTUArray$dosieve<-sieveToplim(lo,hi)<-getBoundssievepct<-countFromToWdlohisievenew<-unsafeNewArray_(0,2*pct-1)::STs(STUArraysIntCacheWord)letfilljindx|hi<indx=returnnew|otherwise=doisPr<-unsafeReadsieveindxifisPrthendolet!i=indx.&.J_MASK!moff=i`shiftL`J_BITSk::Integerk=fromIntegral(indx`shiftR`J_BITS)p=30*k+fromIntegral(rhoi)q0=(start-1)`quot`p(skp0,q1)=q0`quotRem`fromIntegralsieveRange(b0,r0)=idxPr(fromIntegralq1::Int)(b1,r1)|r0==7=(b0+1,0)|otherwise=(b0,r0+1)b2=skp0*fromIntegralsieveBytes+fromIntegralb1strt0=((k*(30*b2+fromIntegral(rhor1))+b2*fromIntegral(rhoi)+fromIntegral(mu(moff+r1)))`shiftL`J_BITS)+fromIntegral(nu(moff+r1))strt1=((k*(30*k+fromIntegral(2*rhoi))+fromIntegral(bytei))`shiftL`J_BITS)+fromIntegral(idxi)(strt2,r2)|p<ssr=(strt0-bitOff,r1)|otherwise=(strt1-bitOff,i)!strt=fromIntegralstrt2.&.IX_MASK!skip=fromIntegral(strt2`shiftR`IX_BITS)!ixes=fromIntegralindx`shiftL`IX_J_BITS.|.strt`shiftL`J_BITS.|.fromIntegralr2unsafeWritenewjskipunsafeWritenew(j+1)ixesfill(j+2)(indx+1)elsefillj(indx+1)fill00-- prime countingnthPrimeCt::Integer->IntegernthPrimeCt1=2nthPrimeCt2=3nthPrimeCt3=5nthPrimeCt4=7nthPrimeCt5=11nthPrimeCt6=13nthPrimeCtn|n<1=error"nthPrimeCt: negative argument"|n<200000=letbd0=nthPrimeApproxnbnd=bd0+bd0`quot`32+37!sv=primeSievebndincountToNth(n-3)[sv]|otherwise=countToNth(n-3)(psieveFrom(fromIntegral$fromIntegern.&.(7::Int)))-- find the n-th set bit in a list of PrimeSieves,-- aka find the (n+3)-rd primecountToNth::Integer->[PrimeSieve]->IntegercountToNth!nps=runST(countDownnps)countDown::Integer->[PrimeSieve]->STsIntegercountDown!n(ps@(PSv0bs):more)|n>278734||(v0/=0&&n>253000)=doct<-countAllpscountDown(n-fromIntegralct)more|otherwise=dostu<-unsafeThawbswa<-(castSTUArray::STUArraysIntBool->STs(STUArraysIntWord))stuletgo!ki|i==sieveWords=countDownkmore|otherwise=dow<-unsafeReadwailet!bc=fromIntegral$bitCountWordwifbc<kthengo(k-bc)(i+1)elselet!j=fromIntegral(bc-k)!px=topwj(fromIntegralbc)inreturn(v0+toPrim(px+(i`shiftL`WSHFT)))gon0countDown_[]=error"Prime stream ended prematurely"-- count all set bits in a chunk, do it wordwise for speed.countAll::PrimeSieve->STsIntcountAll(PS_bs)=dostu<-unsafeThawbswa<-(castSTUArray::STUArraysIntBool->STs(STUArraysIntWord))stuletgo!cti|i==sieveWords=returnct|otherwise=dow<-unsafeReadwaigo(ct+bitCountWordw)(i+1)go00-- Find the j-th highest of bc set bits in the Word w.top::Word->Int->Int->Inttopwjbc=go0TOPBTOPMbnwwhere!bn=bc-jgo!__!_!_0=error"Too few bits set"gobs0__wd=ifwd.&.1==0thenerror"Too few bits, shift 0"elsebsgobsamskixwd=casebitCountWord(wd.&.msk)oflc|lc<ix->go(bs+a)amsk(ix-lc)(wd`uncheckedShiftR`a)|otherwise->let!na=a`shiftR`1ingobsna(msk`uncheckedShiftR`na)ixwd