Black height invariant for red-black trees

pldi2011-testcases/RedBlackTree.scala

 import scala.collection.immutable.Set
 import funcheck.Annotations._
+import funcheck.Utils._
 
 object RedBlackTree { 
   sealed abstract class Color
@@ -38,9 +39,20 @@ object RedBlackTree {
     case Node(_,l,_,r) => redNodesHaveBlackChildren(l) && redNodesHaveBlackChildren(r)
   }
 
+  def blackBalanced(t : Tree) : Boolean = t match {
+    case Node(_,l,_,r) => blackBalanced(l) && blackBalanced(r) && blackHeight(l) == blackHeight(r)
+    case Empty() => true
+  }
+
+  def blackHeight(t : Tree) : Int = t match {
+    case Empty() => 1
+    case Node(Black(), l, _, _) => blackHeight(l) + 1
+    case Node(Red(), l, _, _) => blackHeight(l)
+  }
+
   // <<insert element x into the tree t>>
   def ins(x: Int, t: Tree): Tree = {
-    require(redNodesHaveBlackChildren(t))
+    require(redNodesHaveBlackChildren(t) && blackBalanced(t))
     t match {
       case Empty() => Node(Red(),Empty(),x,Empty())
       case Node(c,a,y,b) =>
@@ -51,53 +63,54 @@ object RedBlackTree {
   } ensuring (res => content(res) == content(t) ++ Set(x) 
                    && size(t) <= size(res) && size(res) <= size(t) + 1
                    && redDescHaveBlackChildren(res)
+                   && blackBalanced(res)
                    )
 
   def makeBlack(n: Tree): Tree = {
-    require(redDescHaveBlackChildren(n))
+    require(redDescHaveBlackChildren(n) && blackBalanced(n))
     n match {
       case Node(Red(),l,v,r) => Node(Black(),l,v,r)
       case _ => n
     }
-  } ensuring(redNodesHaveBlackChildren(_))
+  } ensuring(res => redNodesHaveBlackChildren(res) && blackBalanced(res))
 
   def add(x: Int, t: Tree): Tree = {
-    require(redNodesHaveBlackChildren(t))
+    require(redNodesHaveBlackChildren(t) && blackBalanced(t))
     makeBlack(ins(x, t))
-  } ensuring (res => content(res) == content(t) ++ Set(x) && redNodesHaveBlackChildren(res))
+  } ensuring (res => content(res) == content(t) ++ Set(x) && redNodesHaveBlackChildren(res) && blackBalanced(res))
   
-  def buggyAdd(x: Int, t: Tree): Tree = {
-    require(redNodesHaveBlackChildren(t))
-    // makeBlack(ins(x, t))
-    ins(x, t)
-  } ensuring (res => content(res) == content(t) ++ Set(x) && redNodesHaveBlackChildren(res))
+//  def buggyAdd(x: Int, t: Tree): Tree = {
+//    require(redNodesHaveBlackChildren(t))
+//    // makeBlack(ins(x, t))
+//    ins(x, t)
+//  } ensuring (res => content(res) == content(t) ++ Set(x) && redNodesHaveBlackChildren(res))
   
   def balance(c: Color, a: Tree, x: Int, b: Tree): Tree = {
-    /*require(
-      Node(c,a,x,b) match {
-        case Node(Black(),Node(Red(),Node(Red(),a,_,b),_,c),_,d) =>
-          redNodesHaveBlackChildren(a) &&
-          redNodesHaveBlackChildren(b) &&
-          redNodesHaveBlackChildren(c) &&
-          redNodesHaveBlackChildren(d)
-        case Node(Black(),Node(Red(),a,_,Node(Red(),b,_,c)),_,d) =>
-          redNodesHaveBlackChildren(a) &&
-          redNodesHaveBlackChildren(b) &&
-          redNodesHaveBlackChildren(c) &&
-          redNodesHaveBlackChildren(d)
-        case Node(Black(),a,_,Node(Red(),Node(Red(),b,_,c),_,d)) => 
-          redNodesHaveBlackChildren(a) &&
-          redNodesHaveBlackChildren(b) &&
-          redNodesHaveBlackChildren(c) &&
-          redNodesHaveBlackChildren(d)
-        case Node(Black(),a,_,Node(Red(),b,_,Node(Red(),c,_,d))) => 
-          redNodesHaveBlackChildren(a) &&
-          redNodesHaveBlackChildren(b) &&
-          redNodesHaveBlackChildren(c) &&
-          redNodesHaveBlackChildren(d)
-        case t => redDescHaveBlackChildren(t)
-      }
-    )*/
+    // require(
+    //   Node(c,a,x,b) match {
+    //     case Node(Black(),Node(Red(),Node(Red(),a,_,b),_,c),_,d) =>
+    //       redNodesHaveBlackChildren(a) &&
+    //       redNodesHaveBlackChildren(b) &&
+    //       redNodesHaveBlackChildren(c) &&
+    //       redNodesHaveBlackChildren(d)
+    //     case Node(Black(),Node(Red(),a,_,Node(Red(),b,_,c)),_,d) =>
+    //       redNodesHaveBlackChildren(a) &&
+    //       redNodesHaveBlackChildren(b) &&
+    //       redNodesHaveBlackChildren(c) &&
+    //       redNodesHaveBlackChildren(d)
+    //     case Node(Black(),a,_,Node(Red(),Node(Red(),b,_,c),_,d)) => 
+    //       redNodesHaveBlackChildren(a) &&
+    //       redNodesHaveBlackChildren(b) &&
+    //       redNodesHaveBlackChildren(c) &&
+    //       redNodesHaveBlackChildren(d)
+    //     case Node(Black(),a,_,Node(Red(),b,_,Node(Red(),c,_,d))) => 
+    //       redNodesHaveBlackChildren(a) &&
+    //       redNodesHaveBlackChildren(b) &&
+    //       redNodesHaveBlackChildren(c) &&
+    //       redNodesHaveBlackChildren(d)
+    //     case t => redDescHaveBlackChildren(t)
+    //   }
+    // )
     Node(c,a,x,b) match {
       case Node(Black(),Node(Red(),Node(Red(),a,xV,b),yV,c),zV,d) => 
         Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
@@ -109,19 +122,20 @@ object RedBlackTree {
         Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
       case Node(c,a,xV,b) => Node(c,a,xV,b)
     }
-  } ensuring (res => content(res) == content(Node(c,a,x,b)) ) // && redDescHaveBlackChildren(res))
-  
-  def buggyBalance(c: Color, a: Tree, x: Int, b: Tree): Tree = {
-    Node(c,a,x,b) match {
-      case Node(Black(),Node(Red(),Node(Red(),a,xV,b),yV,c),zV,d) => 
-        Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
-      case Node(Black(),Node(Red(),a,xV,Node(Red(),b,yV,c)),zV,d) => 
-        Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
-      case Node(Black(),a,xV,Node(Red(),Node(Red(),b,yV,c),zV,d)) => 
-        Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
-      case Node(Black(),a,xV,Node(Red(),b,yV,Node(Red(),c,zV,d))) => 
-        Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
-    }
-  } ensuring (res => content(res) == content(Node(c,a,x,b)) ) // && redDescHaveBlackChildren(res))
+  } ensuring (res => content(res) == content(Node(c,a,x,b)) )// && redDescHaveBlackChildren(res))
+
+  // def buggyBalance(c: Color, a: Tree, x: Int, b: Tree): Tree = {
+  //   Node(c,a,x,b) match {
+  //     case Node(Black(),Node(Red(),Node(Red(),a,xV,b),yV,c),zV,d) => 
+  //       Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
+  //     case Node(Black(),Node(Red(),a,xV,Node(Red(),b,yV,c)),zV,d) => 
+  //       Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
+  //     case Node(Black(),a,xV,Node(Red(),Node(Red(),b,yV,c),zV,d)) => 
+  //       Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
+  //     case Node(Black(),a,xV,Node(Red(),b,yV,Node(Red(),c,zV,d))) => 
+  //       Node(Red(),Node(Black(),a,xV,b),yV,Node(Black(),c,zV,d))
+  //   }
+  // } ensuring (res => content(res) == content(Node(c,a,x,b)) ) // && redDescHaveBlackChildren(res))
 
+  def generateRB(t : Tree) : Boolean = (blackHeight(t) != 4 || !blackBalanced(t)) holds
 }