diff --git a/mytest/WhileTest.scala b/mytest/WhileTest.scala
index 4ec15bf42827d9c626c0aa50ba257abd9d4431a2..c080552e54f5861c1ad368ce4c132e957e9e8554 100644
--- a/mytest/WhileTest.scala
+++ b/mytest/WhileTest.scala
@@ -1,13 +1,19 @@
+import leon.Utils._
+
 object WhileTest {
+//  object InvariantFunction {
+//    def invariant(x: Boolean): Unit = ()
+//  }
+//  implicit def while2Invariant(u: Unit) = InvariantFunction
   def foo(x : Int) : Int = {
     require(x >= 0)
 
     var y : Int = x
 
-    while(y >= 0) {
+    (while (y >= 0) {
       y = y - 1
       // assert(y >= -1)
-    }
+    }) invariant(y >= -1)
 
     y + 1
   } ensuring(_ == 0)
diff --git a/src/main/scala/leon/FunctionClosure.scala b/src/main/scala/leon/FunctionClosure.scala
index 2aea0033a6ed3b816b00b3f6192e79ce28e92c33..dc82a9fba2ec008bfe9796fda3dbbe162fbb8192 100644
--- a/src/main/scala/leon/FunctionClosure.scala
+++ b/src/main/scala/leon/FunctionClosure.scala
@@ -18,7 +18,7 @@ object FunctionClosure extends Pass {
   }
 
   private def functionClosure(expr: Expr, bindedVars: Set[Identifier]): Expr = expr match {
-    case l @ LetDef(FunDef(id, rt, varDecl, Some(funBody), _, _), rest) => {
+    case l @ LetDef(fd@FunDef(id, rt, varDecl, Some(funBody), _, _), rest) => {
       val vars = variablesOf(funBody)
       val capturedVars = vars.intersect(bindedVars).toSeq // this should be the variable used that are in the scope
       val freshVars: Map[Identifier, Identifier] = capturedVars.map(v => (v, FreshIdentifier(v.name).setType(v.getType))).toMap
@@ -27,6 +27,8 @@ object FunctionClosure extends Pass {
       val newVarDecls = varDecl ++ extraVarDecls
       val newFunId = FreshIdentifier(id.name)
       val newFunDef = new FunDef(newFunId, rt, newVarDecls)
+      newFunDef.precondition = fd.precondition.map(expr => replace(freshVars.map(p => (p._1.toVariable, p._2.toVariable)), expr))
+      newFunDef.postcondition = fd.postcondition.map(expr => replace(freshVars.map(p => (p._1.toVariable, p._2.toVariable)), expr))
       def substFunInvocInDef(expr: Expr): Option[Expr] = expr match {
         case FunctionInvocation(fd, args) if fd.id == id => Some(FunctionInvocation(newFunDef, args ++ extraVarDecls.map(_.id.toVariable)))
         case _ => None
diff --git a/src/main/scala/leon/ImperativeCodeElimination.scala b/src/main/scala/leon/ImperativeCodeElimination.scala
index 07d14895e8db9379758ce82cf0e97ff66d3517a2..1ed6208240b067b1283d16c0d16dc3a862b72b95 100644
--- a/src/main/scala/leon/ImperativeCodeElimination.scala
+++ b/src/main/scala/leon/ImperativeCodeElimination.scala
@@ -66,7 +66,7 @@ object ImperativeCodeElimination extends Pass {
 
         (resId.toVariable, scope, modifiedVars.zip(freshIds).toMap)
       }
-      case While(cond, body) => {
+      case wh@While(cond, body) => {
         val (_, bodyScope, bodyFun) = toFunction(body)
         val modifiedVars: Seq[Identifier] = bodyFun.keys.toSeq
 
@@ -84,6 +84,14 @@ object ImperativeCodeElimination extends Pass {
           val whileFunBaseCase = (if(whileFunVars.size == 1) whileFunVars.head.toVariable else Tuple(whileFunVars.map(_.toVariable))).setType(whileFunReturnType)
           val whileFunBody = IfExpr(whileFunCond, whileFunRecursiveCall, whileFunBaseCase).setType(whileFunReturnType)
           whileFunDef.body = Some(whileFunBody)
+          val trivialPostcondition: Option[Expr] = Some(Not(replace(whileFunVars.map(id => (id.toVariable, ResultVariable())).toMap, whileFunCond)))
+          val invariantPrecondition: Option[Expr] = wh.invariant.map(expr => replace(modifiedVars2WhileFunVars, expr))
+          whileFunDef.precondition = invariantPrecondition
+          whileFunDef.postcondition = trivialPostcondition.map(expr => 
+              And(expr, invariantPrecondition match { 
+                case Some(e) => replace(whileFunVars.map(id => (id.toVariable, ResultVariable())).toMap, e)
+                case None => BooleanLiteral(true)
+              }))
 
           val finalVars = modifiedVars.map(id => FreshIdentifier(id.name).setType(id.getType))
           val finalScope = ((body: Expr) => {
@@ -150,8 +158,8 @@ object ImperativeCodeElimination extends Pass {
       case m @ MatchExpr(scrut, cses) => sys.error("not supported: " + expr)
       case _ => sys.error("not supported: " + expr)
     }
-    val codeRepresentation = res._2(Block(res._3.map{ case (id1, id2) => Assignment(id1, id2.toVariable)}.toSeq, res._1))
-    println("res of toFunction on: " + expr + " IS: " + codeRepresentation)
+    //val codeRepresentation = res._2(Block(res._3.map{ case (id1, id2) => Assignment(id1, id2.toVariable)}.toSeq, res._1))
+    //println("res of toFunction on: " + expr + " IS: " + codeRepresentation)
     res.asInstanceOf[(Expr, (Expr) => Expr, Map[Identifier, Identifier])]
   }
 
diff --git a/src/main/scala/leon/UnitElimination.scala b/src/main/scala/leon/UnitElimination.scala
index 6df2aee86c9044ba5059072e52d669d323759558..55b8c6c67bef965e9c6a8a1797a90591e7daff80 100644
--- a/src/main/scala/leon/UnitElimination.scala
+++ b/src/main/scala/leon/UnitElimination.scala
@@ -20,6 +20,8 @@ object UnitElimination extends Pass {
     allFuns.foreach(fd => {
       if(fd.returnType != UnitType && fd.args.exists(vd => vd.tpe == UnitType)) {
         val freshFunDef = new FunDef(FreshIdentifier(fd.id.name), fd.returnType, fd.args.filterNot(vd => vd.tpe == UnitType))
+        freshFunDef.precondition = fd.precondition //TODO: maybe removing unit from the conditions as well..
+        freshFunDef.postcondition = fd.postcondition//TODO: maybe removing unit from the conditions as well..
         fun2FreshFun += (fd -> freshFunDef)
       } else {
         fun2FreshFun += (fd -> fd) //this will make the next step simpler
@@ -86,6 +88,8 @@ object UnitElimination extends Pass {
         else {
           val (newFd, rest) = if(fd.args.exists(vd => vd.tpe == UnitType)) {
             val freshFunDef = new FunDef(FreshIdentifier(fd.id.name), fd.returnType, fd.args.filterNot(vd => vd.tpe == UnitType))
+            freshFunDef.precondition = fd.precondition //TODO: maybe removing unit from the conditions as well..
+            freshFunDef.postcondition = fd.postcondition//TODO: maybe removing unit from the conditions as well..
             fun2FreshFun += (fd -> freshFunDef)
             freshFunDef.body = Some(removeUnit(fd.getBody))
             val restRec = removeUnit(b)