Avoid zero-ing out dead fields of primitive value class type

- Zero out fields of type Any
- Zero out fields of value class type

3 files changed, 126 insertions, 14 deletions

diff --git a/src/main/scala/scala/async/internal/AsyncTransform.scala b/src/main/scala/scala/async/internal/AsyncTransform.scala
index 18caea4..27d95a4 100644
--- a/src/main/scala/scala/async/internal/AsyncTransform.scala
+++ b/src/main/scala/scala/async/internal/AsyncTransform.scala
@@ -70,12 +70,11 @@ trait AsyncTransform {
     for ((state, flds) <- assignsOf) {
       val assigns = flds.map { fld =>
         val fieldSym = fld.symbol
-        val zero     = gen.mkZero(fieldSym.info)
         Block(
           List(
-            asyncBase.nullOut(global)(Expr[String](Literal(Constant(fieldSym.name.toString))), Expr[Any](zero)).tree
+            asyncBase.nullOut(global)(Expr[String](Literal(Constant(fieldSym.name.toString))), Expr[Any](Ident(fieldSym))).tree
           ),
-          Assign(gen.mkAttributedStableRef(fieldSym.owner.thisType, fieldSym), zero)
+          Assign(gen.mkAttributedStableRef(fieldSym.owner.thisType, fieldSym), gen.mkZero(fieldSym.info))
         )
       }
       val asyncState = asyncBlock.asyncStates.find(_.state == state).get
diff --git a/src/main/scala/scala/async/internal/LiveVariables.scala b/src/main/scala/scala/async/internal/LiveVariables.scala
index 48f03cd..0f95bca 100644
--- a/src/main/scala/scala/async/internal/LiveVariables.scala
+++ b/src/main/scala/scala/async/internal/LiveVariables.scala
@@ -49,10 +49,11 @@ trait LiveVariables {
 
     // determine which fields should be live also at the end (will not be nulled out)
     val noNull: Set[Symbol] = liftedSyms.filter { sym =>
-      liftables.exists { tree =>
+      sym.tpe.typeSymbol.isPrimitiveValueClass || liftables.exists { tree =>
         !liftedSyms.contains(tree.symbol) && tree.exists(_.symbol == sym)
       }
     }
+    AsyncUtils.vprintln(s"fields never zero-ed out: ${noNull.mkString(", ")}")
 
     /**
      *  Traverse statements of an `AsyncState`, collect `Ident`-s refering to lifted fields.
diff --git a/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala b/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
index 2cecffa..be62ed8 100644
--- a/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
+++ b/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
@@ -11,30 +11,142 @@ import org.junit.Test
 import internal.AsyncTestLV
 import AsyncTestLV._
 
+case class Cell[T](v: T)
+
+class Meter(val len: Long) extends AnyVal
+
+case class MCell[T](var v: T)
+
+
 class LiveVariablesSpec {
 
   @Test
-  def liveVars1() {
+  def `zero out fields of reference type`() {
+    val f = async { Cell(1) }
+
+    def m1(x: Cell[Int]): Cell[Int] =
+      async { Cell(x.v + 1) }
+
+    def m2(x: Cell[Int]): String =
+      async { x.v.toString }
+
+    def m3() = async {
+      val a: Cell[Int] = await(f)      // await$1$1
+      // a == Cell(1)
+      val b: Cell[Int] = await(m1(a))  // await$2$1
+      // b == Cell(2)
+      assert(AsyncTestLV.log.exists(_ == ("await$1$1" -> Cell(1))))
+      val res = await(m2(b))           // await$3$1
+      assert(AsyncTestLV.log.exists(_ == ("await$2$1" -> Cell(2))))
+      res
+    }
+
+    assert(m3() == "2")
+  }
+
+  @Test
+  def `zero out fields of type Any`() {
+    val f = async { Cell(1) }
+
+    def m1(x: Cell[Int]): Cell[Int] =
+      async { Cell(x.v + 1) }
+
+    def m2(x: Any): String =
+      async { x.toString }
+
+    def m3() = async {
+      val a: Cell[Int] = await(f)      // await$4$1
+      // a == Cell(1)
+      val b: Any = await(m1(a))        // await$5$1
+      // b == Cell(2)
+      assert(AsyncTestLV.log.exists(_ == ("await$4$1" -> Cell(1))))
+      val res = await(m2(b))           // await$6$1
+      assert(AsyncTestLV.log.exists(_ == ("await$5$1" -> Cell(2))))
+      res
+    }
+
+    assert(m3() == "Cell(2)")
+  }
+
+  @Test
+  def `do not zero out fields of primitive type`() {
     val f = async { 1 }
 
-    def m1(x: Int): Int =
-      async { x + 1 }
+    def m1(x: Int): Cell[Int] =
+      async { Cell(x + 1) }
 
-    def m2(x: Int): String =
+    def m2(x: Any): String =
       async { x.toString }
 
     def m3() = async {
-      val a = await(f)          // await$1$1
+      val a: Int = await(f)            // await$7$1
       // a == 1
-      val b = await(m1(a))      // await$2$1
-      // b == 2
-      assert(AsyncTestLV.log.exists(_ == ("await$1$1" -> 0)))
-      val res = await(m2(b))    // await$3$1
-      assert(AsyncTestLV.log.exists(_ == ("await$2$1" -> 0)))
+      val b: Any = await(m1(a))        // await$8$1
+      // b == Cell(2)
+      assert(!AsyncTestLV.log.exists(p => p._1 == "await$7$1"))
+      val res = await(m2(b))           // await$9$1
+      assert(AsyncTestLV.log.exists(_ == ("await$8$1" -> Cell(2))))
+      res
+    }
+
+    assert(m3() == "Cell(2)")
+  }
+
+  @Test
+  def `zero out fields of value class type`() {
+    val f = async { Cell(1) }
+
+    def m1(x: Cell[Int]): Meter =
+      async { new Meter(x.v + 1) }
+
+    def m2(x: Any): String =
+      async { x.toString }
+
+    def m3() = async {
+      val a: Cell[Int] = await(f)      // await$10$1
+      // a == Cell(1)
+      val b: Meter = await(m1(a))      // await$11$1
+      // b == Meter(2)
+      assert(AsyncTestLV.log.exists(_ == ("await$10$1" -> Cell(1))))
+      val res = await(m2(b.len))       // await$12$1
+      assert(AsyncTestLV.log.exists(entry => entry._1 == "await$11$1" && entry._2.asInstanceOf[Meter].len == 2L))
       res
     }
 
     assert(m3() == "2")
   }
 
+  @Test
+  def `zero out fields after use in loop`() {
+    val f = async { MCell(1) }
+
+    def m1(x: MCell[Int], y: Int): Int =
+      async { x.v + y }
+
+    def m3() = async {
+      // state #1
+      val a: MCell[Int] = await(f)     // await$13$1
+      // state #2
+      var y = MCell(0)
+
+      while (a.v < 10) {
+        // state #4
+        a.v = a.v + 1
+        y = MCell(await(a).v + 1)      // await$14$1
+        // state #7
+      }
+
+      // state #3
+      assert(AsyncTestLV.log.exists(entry => entry._1 == "await$14$1"))
+
+      val b = await(m1(a, y.v))        // await$15$1
+      // state #8
+      assert(AsyncTestLV.log.exists(_ == ("a$1" -> MCell(10))))
+      assert(AsyncTestLV.log.exists(_ == ("y$1" -> MCell(11))))
+      b
+    }
+
+    assert(m3() == 21)
+  }
+
 }