Sieve in Scala

Sieve.scala

import scala.collection.mutable.BitSet

def sieve_of_eratosthenes(max: Int) = {
    val sieve = BitSet()

    def index(k: Int): Int  = (k - 3) / 2

    def validIdx(k: Int): Boolean = k >= 3 && (k % 2) == 1

    def is_composite(k: Int): Boolean = {
        assert( validIdx(k) )
        sieve(index(k))
    }

    def set_composite(k: Int): Unit  = {
        assert( validIdx(k) )
        sieve.add(index(k))
    }
    
    for {
        i <- 3 to Math.sqrt(max).toInt by 2
        if !is_composite(i)
        // We increment by 2*i to skip even multiples of i
        multiple_i <- i * i to max by 2 * i
    } set_composite(multiple_i)    

    val primes = for (i <- 3 to max by 2; if !is_composite(i) ) yield i
    
    2 +: primes
}

val max = 99
println(sieve_of_eratosthenes(max))