import collection.immutable._
import org.scalacheck._
import Prop._
import Gen._
/*
Properties of a Red & Black Tree:
A node is either red or black.
The root is black. (This rule is used in some definitions and not others. Since the
root can always be changed from red to black but not necessarily vice-versa this
rule has little effect on analysis.)
All leaves are black.
Both children of every red node are black.
Every simple path from a given node to any of its descendant leaves contains the same number of black nodes.
*/
package scala.collection.immutable {
abstract class RedBlackTest extends Properties("RedBlack") {
def minimumSize = 0
def maximumSize = 5
import RedBlack._
def nodeAt[A](tree: Node[String, A], n: Int): Option[(String, A)] = if (n < iterator(tree).size && n >= 0)
Some(iterator(tree).drop(n).next)
else
None
def treeContains[A](tree: Node[String, A], key: String) = iterator(tree).map(_._1) contains key
def mkTree(level: Int, parentIsBlack: Boolean = false, label: String = ""): Gen[Node[String, Int]] =
if (level == 0) {
value(null)
} else {
for {
oddOrEven <- choose(0, 2)
tryRed = oddOrEven.sample.get % 2 == 0 // work around arbitrary[Boolean] bug
isRed = parentIsBlack && tryRed
nextLevel = if (isRed) level else level - 1
left <- mkTree(nextLevel, !isRed, label + "L")
right <- mkTree(nextLevel, !isRed, label + "R")
} yield {
if (isRed)
RedNode(label + "N", 0, left, right)
else
BlackNode(label + "N", 0, left, right)
}
}
def genTree = for {
depth <- choose(minimumSize, maximumSize + 1)
tree <- mkTree(depth)
} yield tree
type ModifyParm
def genParm(tree: Node[String, Int]): Gen[ModifyParm]
def modify(tree: Node[String, Int], parm: ModifyParm): Node[String, Int]
def genInput: Gen[(Node[String, Int], ModifyParm, Node[String, Int])] = for {
tree <- genTree
parm <- genParm(tree)
} yield (tree, parm, modify(tree, parm))
}
trait RedBlackInvariants {
self: RedBlackTest =>
import RedBlack._
def rootIsBlack[A](t: Node[String, A]) = isBlack(t)
def areAllLeavesBlack[A](t: Node[String, A]): Boolean = t match {
case null => isBlack(t)
case ne => List(ne.left, ne.right) forall areAllLeavesBlack
}
def areRedNodeChildrenBlack[A](t: Node[String, A]): Boolean = t match {
case RedNode(_, _, left, right) => List(left, right) forall (t => isBlack(t) && areRedNodeChildrenBlack(t))
case BlackNode(_, _, left, right) => List(left, right) forall areRedNodeChildrenBlack
case null => true
}
def blackNodesToLeaves[A](t: Node[String, A]): List[Int] = t match {
case null => List(1)
case BlackNode(_, _, left, right) => List(left, right) flatMap blackNodesToLeaves map (_ + 1)
case RedNode(_, _, left, right) => List(left, right) flatMap blackNodesToLeaves
}
def areBlackNodesToLeavesEqual[A](t: Node[String, A]): Boolean = t match {
case null => true
case ne =>
(
blackNodesToLeaves(ne).distinct.size == 1
&& areBlackNodesToLeavesEqual(ne.left)
&& areBlackNodesToLeavesEqual(ne.right)
)
}
def orderIsPreserved[A](t: Node[String, A]): Boolean =
iterator(t) zip iterator(t).drop(1) forall { case (x, y) => x._1 < y._1 }
def setup(invariant: Node[String, Int] => Boolean) = forAll(genInput) { case (tree, parm, newTree) =>
invariant(newTree)
}
property("root is black") = setup(rootIsBlack)
property("all leaves are black") = setup(areAllLeavesBlack)
property("children of red nodes are black") = setup(areRedNodeChildrenBlack)
property("black nodes are balanced") = setup(areBlackNodesToLeavesEqual)
property("ordering of keys is preserved") = setup(orderIsPreserved)
}
object TestInsert extends RedBlackTest with RedBlackInvariants {
import RedBlack._
override type ModifyParm = Int
override def genParm(tree: Node[String, Int]): Gen[ModifyParm] = choose(0, iterator(tree).size + 1)
override def modify(tree: Node[String, Int], parm: ModifyParm): Node[String, Int] = update(tree, generateKey(tree, parm), 0)
def generateKey(tree: Node[String, Int], parm: ModifyParm): String = nodeAt(tree, parm) match {
case Some((key, _)) => key.init.mkString + "MN"
case None => nodeAt(tree, parm - 1) match {
case Some((key, _)) => key.init.mkString + "RN"
case None => "N"
}
}
property("update adds elements") = forAll(genInput) { case (tree, parm, newTree) =>
treeContains(newTree, generateKey(tree, parm))
}
}
object TestModify extends RedBlackTest {
import RedBlack._
def newValue = 1
override def minimumSize = 1
override type ModifyParm = Int
override def genParm(tree: Node[String, Int]): Gen[ModifyParm] = choose(0, iterator(tree).size)
override def modify(tree: Node[String, Int], parm: ModifyParm): Node[String, Int] = nodeAt(tree, parm) map {
case (key, _) => update(tree, key, newValue)
} getOrElse tree
property("update modifies values") = forAll(genInput) { case (tree, parm, newTree) =>
nodeAt(tree,parm) forall { case (key, _) =>
iterator(newTree) contains (key, newValue)
}
}
}
object TestDelete extends RedBlackTest with RedBlackInvariants {
import RedBlack._
override def minimumSize = 1
override type ModifyParm = Int
override def genParm(tree: Node[String, Int]): Gen[ModifyParm] = choose(0, iterator(tree).size)
override def modify(tree: Node[String, Int], parm: ModifyParm): Node[String, Int] = nodeAt(tree, parm) map {
case (key, _) => delete(tree, key)
} getOrElse tree
property("delete removes elements") = forAll(genInput) { case (tree, parm, newTree) =>
nodeAt(tree, parm) forall { case (key, _) =>
!treeContains(newTree, key)
}
}
}
object TestRange extends RedBlackTest with RedBlackInvariants {
import RedBlack._
override type ModifyParm = (Option[Int], Option[Int])
override def genParm(tree: Node[String, Int]): Gen[ModifyParm] = for {
from <- choose(0, iterator(tree).size)
to <- choose(0, iterator(tree).size) suchThat (from <=)
optionalFrom <- oneOf(Some(from), None, Some(from)) // Double Some(n) to get around a bug
optionalTo <- oneOf(Some(to), None, Some(to)) // Double Some(n) to get around a bug
} yield (optionalFrom, optionalTo)
override def modify(tree: Node[String, Int], parm: ModifyParm): Node[String, Int] = {
val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
range(tree, from, to)
}
property("range boundaries respected") = forAll(genInput) { case (tree, parm, newTree) =>
val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
("lower boundary" |: (from forall ( key => iterator(newTree).map(_._1) forall (key <=)))) &&
("upper boundary" |: (to forall ( key => iterator(newTree).map(_._1) forall (key >))))
}
property("range returns all elements") = forAll(genInput) { case (tree, parm, newTree) =>
val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
val filteredTree = (iterator(tree)
.map(_._1)
.filter(key => from forall (key >=))
.filter(key => to forall (key <))
.toList)
filteredTree == iterator(newTree).map(_._1).toList
}
}
}
object Test extends Properties("RedBlack") {
include(TestInsert)
include(TestModify)
include(TestDelete)
include(TestRange)
}