Scala TreeSet/TreeMap benchmarks

treemap-bench.scala

import collection.immutable.TreeMap
import util.Random

object TreeMapTest {

  def time(block: => Unit): Double = {
    val start = System.nanoTime()
    block
    val stop = System.nanoTime()
    (stop - start) / 1.0e9
  }

  def bench(name: String, tree: TreeMap[Int, Int], count: Int)(block: TreeMap[Int, Int] => AnyRef): Any = {
    var r: AnyRef = null
    val s1 = time {
      for (_ <- 1 to count) {
        r = block(tree)
      }
    }
    println("%-25s: size: %7d: %7d iterations in %6.3f seconds (%12.2f per second)".format(name, tree.size, count, s1, count / s1))
    r
  }

  def tests(tree: TreeMap[Int, Int]) {
    for (_ <- 1 to 10) {
      val n = Random.nextInt(tree.size)
      assert(tree == tree.take(n) ++ tree.drop(n), "take / drop")
      assert(tree == tree.dropRight(n) ++ tree.takeRight(n), "dropRight / takeRight")
      val (pr, sf) = tree.splitAt(n)
      assert(pr.size == n, "prefix size")
      assert(sf.size == tree.size - n, "suffix size")
      assert(tree == pr ++ sf, "split")
      assert(tree == tree.tail + tree.head, "head / tail")
      assert(tree == tree.init + tree.last, "init / last")
      val rnd = Random.nextInt(tree.size)
      assert(tree == tree.takeWhile(_._1 < rnd) ++ tree.dropWhile(_._1 < rnd), "takeWhile / dropWhile")

      val from = Random.nextInt(tree.size)
      val until = Random.nextInt(tree.size) + 1
      assert(tree == tree.take(from) ++ tree.slice(from, until) ++ tree.takeRight(tree.size - until), "slice")
    }
  }

  def main(args: Array[String]): Unit = {
    val sizes: Seq[Int] = if (args.size < 1) List(1, 10, 100, 1000, 10000, 100000) else args(0).split(",").map(_.toInt).toSeq
    val trees = for (size <- sizes) yield {
      val values = Array.fill(size)({ val r = Random.nextInt; r -> r })
      val tree = TreeMap(values: _*)
      if (values.toSet != tree.iterator.toSet) {
        println("Iterator not correct: expected\n\t<%s> was\n\t<%s>".format(values.toIndexedSeq.sorted, tree.iterator.toIndexedSeq))
      }
      tree
    }
    val counts: Seq[Int] = if (args.size < 2) List(1000) else args(1).split(",").map(_.toInt).toSeq
    trees.foreach(tests)
    for (_ <- 1 to 3; count <- counts; tree <- trees) {
//      bench("headOption", tree, count) { _.headOption }
//      bench("lastOption", tree, count) { _.lastOption }
//      bench("tail", tree, count) { _.tail }
//      bench("init", tree, count) { _.init }
//      bench("drop(5)", tree, count) { _.drop(5) }
//      bench("dropRight(5)", tree, count) { _.dropRight(5) }
//      bench("take(5)", tree, count) { _.take(5) }
//      bench("takeRight(5)", tree, count) { _.takeRight(5) }
//      bench("splitAt(5)", tree, count) { _.splitAt(5) }
//      bench("drop(size / 2)", tree, count) { t => t.drop(t.size / 2) }
//      bench("dropRight(size / 2)", tree, count) { t => t.dropRight(t.size / 2) }
//      bench("take(size / 2)", tree, count) { t => t.take(t.size / 2) }
//      bench("takeRight(size / 2)", tree, count) { t => t.takeRight(t.size / 2) }
//      bench("splitAt(size / 2)", tree, count) { t => t.splitAt(t.size / 2) }
      bench("slice", tree, count) { t => 
        val from = t.size / 3
        val until = 2 * t.size / 3
        t.slice(from, until)
      }
//      bench("dropWhile(_ => false)", tree, count) { _.dropWhile(_ => false) }
//      bench("takeWhile(_ => false)", tree, count) { _.takeWhile(_ => false) }
//      bench("span(_ => false)", tree, count) { _.span(_ => false) }
//      bench("dropWhile(_._1 < 0)", tree, count) { _.dropWhile(_._1 < 0) }
//      bench("takeWhile(_._1 < 0)", tree, count) { _.takeWhile(_._1 < 0) }
      bench("span(_._1 < 0)", tree, count) { _.span(_._1 < 0) }
//      bench("dropWhile(_ => true)", tree, count) { _.dropWhile(_ => true) }
//      bench("takeWhile(_ => true)", tree, count) { _.takeWhile(_ => true) }
//      bench("span(_ => true)", tree, count) { _.span(_ => true) }
//      bench("iterator", tree, count) { _.iterator }
//      bench("iterator.size", tree, count) { _.iterator.size: java.lang.Integer }
//      bench("toStream", tree, count) { _.toStream }
//      bench("toStream.size", tree, count) { _.toStream.size: java.lang.Integer }
//      bench("submap", tree, count) { _.range(Random.nextInt, Random.nextInt) }
    }
  }
}

TreeMapTest.main(args)

treeset-bench.scala

import collection.immutable.TreeSet
import util.Random

object TreeSetTest {

  def time(block: => Unit): Double = {
    val start = System.nanoTime()
    block
    val stop = System.nanoTime()
    (stop - start) / 1.0e9
  }

  def usedMemory = {
    System.gc()
    Runtime.getRuntime.totalMemory - Runtime.getRuntime.freeMemory
  }

  def tests(tree: TreeSet[Int]) {
    for (_ <- 1 to 10) {
      val n = Random.nextInt(tree.size)
      assert(tree == tree.take(n) ++ tree.drop(n), "take / drop")
      assert(tree == tree.dropRight(n) ++ tree.takeRight(n), "dropRight / takeRight")
      val (pr, sf) = tree.splitAt(n)
      assert(pr.size == n, "prefix size was " + pr.size + " expected " + n)
      assert(sf.size == tree.size - n, "suffix size")
      assert(tree == pr ++ sf, "split")
      assert(tree == tree.tail + tree.head, "head / tail")
      assert(tree == tree.init + tree.last, "init / last")
      val rnd = Random.nextInt(tree.size)
      assert(tree == tree.takeWhile(_ < rnd) ++ tree.dropWhile(_ < rnd), "takeWhile / dropWhile")

      val from = Random.nextInt(tree.size)
      val until = Random.nextInt(tree.size) + 1
      assert(tree == tree.take(from) ++ tree.slice(from, until) ++ tree.takeRight(tree.size - until), "slice")
    }
  }

  def bench(name: String, tree: TreeSet[Int], count: Int)(block: TreeSet[Int] => Unit): Unit = {
    val s1 = time {
      for (_ <- 1 to count) {
        block(tree)
      }
    }
    println("%-25s: size: %7d: %7d iterations in %6.3f seconds (%12.2f per second)".format(name, tree.size, count, s1, count / s1))
  }

  def main(args: Array[String]): Unit = {
    val memoryBefore = usedMemory
    println("Used memory before allocating trees: %d KiB".format(memoryBefore / 1024))
    val sizes: Seq[Int] = if (args.size < 1) List(1, 10, 100, 1000, 10000, 100000) else args(0).split(",").map(_.toInt).toSeq
    val trees = for (size <- sizes) yield {
      println("Generating random numbers")
      val values = Array.fill(size)(Random.nextInt)
      println("Building tree")
      val tree = TreeSet(values: _*)
//      if (values.toSet != tree.iterator.toSet) {
//        println("Iterator not correct: expected\n\t<%s> was\n\t<%s>".format(values.toIndexedSeq.sorted, tree.iterator.toIndexedSeq))
//      }
      tree
    }
    val memoryAfter = usedMemory
    println("Used memory after allocating trees: %d KiB (~%d bytes per node including boxed int)".format(memoryAfter / 1024, memoryAfter / trees.map(_.size).sum))
    val counts: Seq[Int] = if (args.size < 2) List(1000) else args(1).split(",").map(_.toInt).toSeq

    trees.foreach(tests)

    for (count <- counts; _ <- 1 to 5; tree <- trees) {
//      val rnd = Random.nextInt
//      bench("lookup", tree, count) { _.contains(rnd) }
//      bench("foreach", tree, count) { _.foreach(_ => ()) }
      bench("head", tree, count) { _.head }
      bench("last", tree, count) { _.last }
//      bench("headOption", tree, count) { _.headOption }
//      bench("lastOption", tree, count) { _.lastOption }
//      bench("tail", tree, count) { _.tail }
//      bench("init", tree, count) { _.init }
      bench("iterator", tree, count) { _.iterator }
      bench("iterator.size", tree, count) { _.iterator.size }
//      bench("toStream", tree, count) { _.toStream }
//      bench("toStream.size", tree, count) { _.toStream.size }
//      bench("drop(5)", tree, count) { _.drop(5) }
//      bench("dropRight(5)", tree, count) { _.dropRight(5) }
//      bench("take(5)", tree, count) { _.take(5) }
//      bench("takeRight(5)", tree, count) { _.takeRight(5) }
//      bench("splitAt(5)", tree, count) { _.splitAt(5) }
//      bench("drop(size / 2)", tree, count) { t => t.drop(t.size / 2) }
//      bench("dropRight(size / 2)", tree, count) { t => t.dropRight(t.size / 2) }
//      bench("take(size / 2)", tree, count) { t => t.take(t.size / 2) }
//      bench("takeRight(size / 2)", tree, count) { t => t.takeRight(t.size / 2) }
//      bench("splitAt(size / 2)", tree, count) { t => t.splitAt(t.size / 2) }
//      bench("slice", tree, count) { t => 
//        val from = t.size / 3
//        val until = 2 * t.size / 3
//        t.slice(from, until)
//      }
//      bench("dropWhile(_ => false)", tree, count) { _.dropWhile(_ => false) }
//      bench("takeWhile(_ => false)", tree, count) { _.takeWhile(_ => false) }
//      bench("span(_ => false)", tree, count) { _.span(_ => false) }
//      bench("dropWhile(_ < 0)", tree, count) { _.dropWhile(_ < 0) }
//      bench("takeWhile(_ < 0)", tree, count) { _.takeWhile(_ < 0) }
//      bench("span(_ < 0)", tree, count) { _.span(_ < 0) }
//      bench("dropWhile(_ => true)", tree, count) { _.dropWhile(_ => true) }
//      bench("takeWhile(_ => true)", tree, count) { _.takeWhile(_ => true) }
//      bench("span(_ => true)", tree, count) { _.span(_ => true) }
//      bench("range", tree, count) { _.range(Random.nextInt, Random.nextInt) }
    }
  }
}

TreeSetTest.main(args)