Smooth.scala

object Smooth {

  type Prime = BigInt
  type Exponent = Int
  type Factorization = Map[Prime, Exponent]

  case class Row(a: BigInt, b: BigInt, f: Factorization)

  def sqrt(a : BigInt): BigInt = {
    def next(n : BigInt, i : BigInt) : BigInt = (n + i/n) >> 1

    val one = BigInt(1)

    var n = one
    var n1 = next(n, a)

    while ((n1 - n).abs > one) {
      n = n1
      n1 = next(n, a)
    }

    while (n1 * n1 > a) {
      n1 -= one
    }

    n1
  }

  def factorBy(p: Prime)(a: BigInt): Option[(Exponent, BigInt)] = {
    var i = 0
    var n = a
    while (n % p == 0) {
      n = n / p
      i += 1
    }

    Option((i, n)).filter{ case (j, _) => j > 0 }
  }

  def factorWrtBase(B: Set[Prime])(a: BigInt): Option[Factorization] =
    Option(B.foldLeft((a, Map.empty[Prime, Exponent])){
      case ((n, f), p) => factorBy(p)(n).map{ case (i, m) => (m, f + ((p, i))) }.getOrElse((n, f))
    }).filter{ case (n, _) => n == 1 }.map(_._2)

  def generateRows(B: Set[Prime])(N: BigInt): Iterator[Row] =
    Iterator.iterate(BigInt(1))(_ + 1)
      .map(sqrt(N) + _).map(a => (a, a.modPow(2, N)))
      .flatMap{ case (a, b) => factorWrtBase(B)(b).map(f => Row(a, b, f)) }

  def main(args: Array[String]): Unit = {
    generateRows(Set(2, 5, 7, 13, 19))(2811929L).take(8).foreach(println)
  }
}