Sieve of Eratosthenes

gistfile1.scala

import scala.testing.Benchmark
import scala.collection.mutable
import scala.collection.immutable
import scala.annotation.tailrec

object PrimeNumbers {
  def main(args: Array[String]) {
    if (args.size != 2) {
      println("usage: scala PrimeNumbers PRIME_MAX RUN_COUNT")
      sys.exit()
    }
    runTest(args(0).toInt, args(1).toInt)
  }

  def runTest(PRIME_MAX: Int, RUN_COUNT: Int) {
    lazy val bitset1 = new Benchmark {
      def run() {
        val bs = mutable.BitSet((3 to PRIME_MAX by 2): _*)
        bs += 2

        val sqrt = math.sqrt(PRIME_MAX).toInt
        var cur = 3

        while(cur <= sqrt) {
          if(bs(cur)) {
            var mul = 2 * cur
            while(mul <= PRIME_MAX) {
              bs -= mul
              mul += cur
            }
          }

          cur += 2
        }

        // println(bs.take(20))
      }
    }.runBenchmark(RUN_COUNT)

    lazy val bitset2 = new Benchmark {
      def run() {
        val bs = mutable.BitSet.empty
        bs += 2

        var cur = 3
        while(cur <= PRIME_MAX) {
          bs += cur
          cur += 2
        }

        val sqrt = math.sqrt(PRIME_MAX).toInt
        cur = 3

        while(cur <= sqrt) {
          if(bs(cur)) {
            var mul = 2 * cur
            while(mul <= PRIME_MAX) {
              bs -= mul
              mul += cur
            }
          }

          cur += 2
        }

        // println(bs.take(20))
      }
    }.runBenchmark(RUN_COUNT)

    lazy val rangeBench = new Benchmark {
      def run() {
        val bs = mutable.BitSet((3 to PRIME_MAX by 2): _*)
      }
    }.runBenchmark(RUN_COUNT)

    val array = new Benchmark {
      def run() {
        val ar = Array.fill(PRIME_MAX + 1)(true)
        ar(0) = false
        ar(1) = false

        var cur = 4
        while(cur <= PRIME_MAX) {
          ar(cur) = false
          cur += 2
        }

        val sqrt = math.sqrt(PRIME_MAX).toInt
        cur = 3
        while(cur <= sqrt) {
          if(ar(cur)) {
            var mul = 2 * cur
            while(mul <= PRIME_MAX) {
              ar(mul) = false
              mul += cur
            }
          }

          cur += 2
        }

        // println(ar.take(72).zipWithIndex.withFilter { _._1 == true } map { _._2 } mkString(", "))
      }
    }.runBenchmark(RUN_COUNT)

    lazy val bitset3 = new Benchmark {
      def run() {
        val bs = Iterator
          .iterate(3) { _ + 2 }
          .takeWhile { _ <= PRIME_MAX }
          .foldLeft(immutable.BitSet.empty + 2) { (bit, num) => bit + num }

        val sqrt = math.sqrt(PRIME_MAX)

        val r = Iterator
          .iterate(3) { _ + 2 }
          .takeWhile { _ <= sqrt }
          .foldLeft(bs) { (bit, num) =>
            if(bit(num))
              Iterator
                .iterate(num * 2) { _ + num }
                .takeWhile { _ <= PRIME_MAX }
                .foldLeft(bit) { (bit2, num2) => bit2 - num2 }
            else bit
          }

        // println(r.take(20))
      }
    }.runBenchmark(RUN_COUNT)

    lazy val stream1 = new Benchmark {
      def run() {
        def func(str: Stream[Int]): Stream[Int] = {
          val x = str.head
          val xs = str.tail
          Stream.cons(x, func(xs.filter { p => p % x != 0 } ))
        }
        val r = func(Stream.from(2)).takeWhile(_ <= PRIME_MAX).toArray
        // println(r.take(20).mkString(", "))
      }
    }.runBenchmark(RUN_COUNT)

    lazy val stream2 = new Benchmark {
      def run() {
        def func(str: Stream[Int]): Stream[Int] = {
          val x = str.head
          val xs = str.tail
          Stream.cons(x, func(xs.filter { p => p < x * x || p % x != 0 } ))
        }
        val r = func(Stream.from(2)).takeWhile(_ <= PRIME_MAX).toArray
        // println(r.take(20).mkString(", "))
      }
    }.runBenchmark(RUN_COUNT)

    lazy val list = new Benchmark {
      def run() {
        @tailrec
        def inner(ls: List[Int], acc: List[Int], end: Int): List[Int] = {
          val x = ls.head
          val xs = ls.tail
          if (x * x <= end) {
            inner(xs.filter(_ % x != 0), x :: acc, end)
          } else acc.reverse ++ ls
        }
        val r = inner((2 to PRIME_MAX).toList, Nil, PRIME_MAX)
        // println(r.take(20))
      }
    }.runBenchmark(RUN_COUNT)

    def statics(res: List[Long]): Long = {
      val sz = res.size
      res.sortWith(_ < _).slice(1, sz - 1).sum / (sz - 2)
    }

    // println("range: " + rangeBench)
    println("bitset1: " + bitset1)
    println(statics(bitset1))
    println("bitset2: " + bitset2)
    println(statics(bitset2))
    println("bitset3: " + bitset3)
    println(statics(bitset3))
    // println("stream1: " + stream1)
    // println("stream2: " + stream2)
    println("array: " + array)
    println(statics(array))
    println("list: " + list)
    println(statics(list))
  }
}