diff --git a/cp-demo/ChooseCalls.scala b/cp-demo/ChooseCalls.scala
index 2e4480db2f80bb64d9bc8dcf50cc8ea5e7f5f63b..77f37349e648bfe897a8162ad17ea9b7e64212c5 100644
--- a/cp-demo/ChooseCalls.scala
+++ b/cp-demo/ChooseCalls.scala
@@ -38,11 +38,21 @@ object ChooseCalls {
   }
 
   def chooseTree(height : Int) : Tree = {
-    choose((t: Tree) => blackBalanced(t) && size(t) == height)
+    choose((t: Tree) => blackBalanced(t) && redNodesHaveBlackChildren(t) && isBlack(t) && size(t) == height)
   }
 
   def main(args: Array[String]) : Unit = {
-    val height = 5
-    println("The chosen tree (of height " + height + ") is : " + chooseTree(height))
+
+    /** Printing trees */
+    def indent(s: String) = ("  "+s).split('\n').mkString("\n  ")
+
+    def print(tree: Tree): String = tree match {
+      case Node(c,l,v,r) =>
+        indent(print(r)) + "\n" + (if (c == Black()) "B" else "R") + " " + v.toString + "\n" + indent(print(l))
+      case Empty() => "E"
+    }
+
+    val height = if (args.isEmpty) 3 else args(0).toInt
+    println("The chosen tree (of height " + height + ") is : \n" + print(chooseTree(height)))
   }
 }
diff --git a/cp-runner b/cp-runner
index efc3278031e7a9b7195106d6bf90c1d2fc6a55ce..dc279546167a78bcb1e47667fd1ae5b353af2871 100755
--- a/cp-runner
+++ b/cp-runner
@@ -1,4 +1,4 @@
 #!/bin/bash
 
 LD_LIBRARY_PATH=lib-bin \
-scala -classpath bin/purescala/purescala-definitions_2.8.1-1.0.jar:bin/funcheck/funcheck-plugin_2.8.1-1.0.jar:lib/z3.jar:bin/cp/constraint-programming-plugin_2.8.1-1.0.jar:out ${1}
+scala -classpath bin/purescala/purescala-definitions_2.8.1-1.0.jar:bin/funcheck/funcheck-plugin_2.8.1-1.0.jar:lib/z3.jar:bin/cp/constraint-programming-plugin_2.8.1-1.0.jar:out ${@}
diff --git a/src/cp/CPComponent.scala b/src/cp/CPComponent.scala
index 0ce12c08fcdda1383430d9a27e394a54bdbfe900..a9804fc1b7c0e204b66294871acb0a127e933542 100644
--- a/src/cp/CPComponent.scala
+++ b/src/cp/CPComponent.scala
@@ -45,16 +45,6 @@ class CPComponent(val global: Global, val pluginInstance: CPPlugin)
 
       transformCalls(unit, prog, progString, progId)
       println("Finished transformation")
-
-      /*
-      try {
-        val recovered = readProgram(filename)
-        println
-        println("Recovered: " + recovered)
-      } catch {
-        case e => e.printStackTrace()
-      }
-      */
     }
   }
 }
diff --git a/src/cp/CallTransformation.scala b/src/cp/CallTransformation.scala
index 5715bcde8d9862537a4d94c767c6e04c588b72a6..320504b316280f17b9b0623c993715525976166b 100644
--- a/src/cp/CallTransformation.scala
+++ b/src/cp/CallTransformation.scala
@@ -126,6 +126,7 @@ trait CallTransformation
   }
 }
 
+/** A collection of methods that are called on runtime */
 object CallTransformation {
   /* Get list of assignments in the order specified by outputVars list */
   def outputAssignments(outputVars: List[String], model: Map[Identifier, Expr]) : List[Any] = {
diff --git a/src/funcheck/CodeExtraction.scala b/src/funcheck/CodeExtraction.scala
index f139f1582ec169eb9b45ec0cdc6ff0ac34e3e867..f8fc276d68843e5a207dbed761ee0d575b613a63 100644
--- a/src/funcheck/CodeExtraction.scala
+++ b/src/funcheck/CodeExtraction.scala
@@ -19,17 +19,23 @@ trait CodeExtraction extends Extractors {
   private lazy val setTraitSym = definitions.getClass("scala.collection.immutable.Set")
   private lazy val multisetTraitSym = definitions.getClass("scala.collection.immutable.Multiset")
 
-  val varSubsts: scala.collection.mutable.Map[Symbol,Function0[Expr]] =
+  private val varSubsts: scala.collection.mutable.Map[Symbol,Function0[Expr]] =
     scala.collection.mutable.Map.empty[Symbol,Function0[Expr]]
-  val classesToClasses: scala.collection.mutable.Map[Symbol,ClassTypeDef] =
+  private val classesToClasses: scala.collection.mutable.Map[Symbol,ClassTypeDef] =
     scala.collection.mutable.Map.empty[Symbol,ClassTypeDef]
   private val defsToDefs: scala.collection.mutable.Map[Symbol,FunDef] =
     scala.collection.mutable.Map.empty[Symbol,FunDef]
 
-  val reverseVarSubsts: scala.collection.mutable.Map[Expr,Symbol] =
+  private val reverseVarSubsts_ : scala.collection.mutable.Map[Expr,Symbol] =
     scala.collection.mutable.Map.empty[Expr,Symbol]
-  val reverseClassesToClasses: scala.collection.mutable.Map[ClassTypeDef,Symbol] =
+  private val reverseClassesToClasses_ : scala.collection.mutable.Map[ClassTypeDef,Symbol] =
     scala.collection.mutable.Map.empty[ClassTypeDef,Symbol]
+
+  def reverseVarSubsts: scala.collection.immutable.Map[Expr,Symbol] =
+    scala.collection.immutable.Map() ++ reverseVarSubsts_
+
+  def reverseClassesToClasses: scala.collection.immutable.Map[ClassTypeDef,Symbol] =
+    scala.collection.immutable.Map() ++ reverseClassesToClasses_
   
   protected def stopIfErrors: Unit = {
     if(reporter.hasErrors) {
@@ -271,8 +277,8 @@ trait CodeExtraction extends Extractors {
     stopIfErrors
 
     // Reverse map for Scala class symbols
-    reverseClassesToClasses ++= classesToClasses.map{ case (a, b) => (b, a) }
-    reverseVarSubsts ++= varSubsts.map{ case (a, b) => (b(), a) }
+    reverseClassesToClasses_ ++= classesToClasses.map{ case (a, b) => (b, a) }
+    reverseVarSubsts_ ++= varSubsts.map{ case (a, b) => (b(), a) }
 
     val programName: Identifier = unit.body match {
       case PackageDef(name, _) => FreshIdentifier(name.toString)