nim-lib/sieve.nim

import math

type
  Factor* = object
    prime*: uint64
    power*: uint64

  Factors* = object
    numPrimes*: uint64
    factors*: seq[Factor]

  U64Array* = object
    elems*: seq[uint64]

# =========================
# Utilities
# =========================

proc factorAppend*(f: var Factors, p, k: uint64) =
  f.factors.add Factor(prime: p, power: k)
  inc f.numPrimes

# =========================
# Math helpers
# =========================

proc u64Pow*(a, e: uint64): uint64 =
  var base = a
  var exp = e
  var r: uint64 = 1
  while exp > 0:
    if (exp and 1) == 1:
      r *= base
    base *= base
    exp = exp shr 1
  r

proc u64NthRoot*(x, n: uint64): uint64 =
  uint64(pow(float64(x), 1.0 / float64(n)))

proc u64Binom*(n, k: uint64): uint64 =
  if k > n: return 0
  var r: uint64 = 1
  for i in 1..k:
    r = r * (n - k + i) div i
  r

proc i64Gcd*(a, b: int64): int64 =
  var x = a
  var y = b
  while y != 0:
    let t = x mod y
    x = y
    y = t
  x

proc i64Lcm*(a, b: int64): int64 =
  (a div i64Gcd(a, b)) * b

proc i64Egcd*(a, b: int64, x, y: var int64): int64 =
  if b == 0:
    x = 1
    y = 0
    return a

  var x1, y1: int64
  let g = i64Egcd(b, a mod b, x1, y1)
  x = y1
  y = x1 - y1 * (a div b)
  g

# =========================
# Prime estimate helpers
# =========================

proc maxPrimeDivs*(max: uint64): uint64 =
  let primorials = [
    6'u64,30,210,2310,30030,
    510510,9699690,223092870
  ]

  for i, p in primorials:
    if max < p:
      return uint64(i+1)

  15

proc maxPrimesLe*(max: uint64): uint64 =
  uint64(1.25506 * float64(max) / ln(float64(max)))

# =========================
# Omega sieve
# =========================

proc sieveOmega*(max: uint64): seq[uint8] =
  result = newSeq[uint8](max+1)

  for n in 2..max:
    if result[n] != 0: continue
    var m = n
    while m <= max:
      inc result[m]
      m += n

# =========================
# Largest prime factor sieve
# =========================

proc sieveLargestFactors*(max: uint64): seq[uint64] =
  result = newSeq[uint64](max+1)

  for p in 2..max:
    if result[p] != 0: continue
    var m = p
    while m <= max:
      result[m] = p
      m += p

# =========================
# Smallest prime factor sieve
# =========================

proc sieveSmallestFactors*(max: uint64): seq[uint32] =
  result = newSeq[uint32](max+1)

  let r = uint64(sqrt(float64(max)))

  for p in 2..r:
    if result[p] != 0: continue

    var m = p*p
    while m <= max:
      if result[m] == 0:
        result[m] = uint32(p)
      m += p

  for i in 2..max:
    if result[i] == 0:
      result[i] = uint32(i)

# =========================
# Factorization
# =========================

proc fillFactorsFromSmallest*(res: var Factors, n: uint64, spf: seq[uint32]) =
  var x = n
  res.numPrimes = 0
  res.factors.setLen(0)

  while x > 1:
    let p = uint64(spf[x])
    var k: uint64 = 0

    while spf[x] == uint32(p):
      x = x div p
      inc k

    factorAppend(res, p, k)
# =========================
# Totient sieve
# =========================

proc sievePhi*(max: uint64): seq[uint64] =
  result = newSeq[uint64](max+1)

  for i in 1..max:
    result[i] = i

  for p in 2..max:
    if result[p] != p: continue

    var m = p
    while m <= max:
      result[m] -= result[m] div p
      m += p

# =========================
# Sigma_0 (divisor count)
# =========================

proc sigma0*(max: uint64): seq[uint64] =
  result = newSeq[uint64](max+1)

  for i in 1..max:
    var j = i
    while j <= max:
      inc result[j]
      j += i

# =========================
# Mobius sieve
# =========================

proc mobius*(max: uint64): seq[int8] =
  result = newSeq[int8](max+1)

  for i in 1..max:
    result[i] = 1

  for p in 2..max:
    if result[p] == 1:
      var m = p
      while m <= max:
        result[m] = -result[m]
        m += p

      let p2 = p*p
      m = p2
      while m <= max:
        result[m] = 0
        m += p2

# =========================
# Prime sieve
# =========================

proc sievePrimes*(max: uint64): seq[uint64] =
  var isComp = newSeq[bool](max+1)

  for p in 2..max:
    if isComp[p]: continue

    result.add p

    if p*p <= max:
      var m = p*p
      while m <= max:
        isComp[m] = true
        m += p

# =========================
# Main
# =========================

when isMainModule:

  let max = 10_000_000'u64

  let primes = sievePrimes(max)

  echo "Primes up to ", max, ":"
  for p in primes:
    stdout.write($p & " ")
  echo ""

  let phi = sievePhi(max)
  echo "phi(36) = ", phi[36]

  let mu = mobius(max)
  echo "mu(30) = ", mu[30]