Substitute the compacted state ids into the code snipptets in .dot

Author: retronym

src/main/scala/scala/async/internal/ExprBuilder.scala

@@ -397,7 +397,7 @@ trait ExprBuilder {
           dotBuilder.append(s"""${stateLabel(state.state)} [label=""").append("<")
           if (i != length - 1) {
             val CaseDef(_, _, body) = state.mkHandlerCaseForState
-            toHtmlLabel(stateLabel(state.state), showCode(body), dotBuilder)
+            toHtmlLabel(stateLabel(state.state), showCode(compactStateTransform.transform(body)), dotBuilder)
           } else {
             toHtmlLabel(stateLabel(state.state), state.allStats.map(showCode(_)).mkString("\n"), dotBuilder)
           }
@@ -506,18 +506,19 @@ trait ExprBuilder {
             })), EmptyTree)
       }
 
-      private def compactStates(m: Match): Tree = {
-        val stateMemberSymbol = symLookup.stateMachineMember(name.state)
-        val compactStateTransform = new Transformer {
-          override def transform(tree: Tree): Tree = tree match {
-            case as @ Assign(lhs, Literal(Constant(i: Integer))) if lhs.symbol == stateMemberSymbol =>
-              val replacement = switchIds(i)
-              treeCopy.Assign(tree, lhs, Literal(Constant(replacement)))
-            case _: Match | _: CaseDef | _: Block | _: If =>
-              super.transform(tree)
-            case _ => tree
-          }
+      private lazy val stateMemberSymbol = symLookup.stateMachineMember(name.state)
+      private val compactStateTransform = new Transformer {
+        override def transform(tree: Tree): Tree = tree match {
+          case as @ Assign(lhs, Literal(Constant(i: Integer))) if lhs.symbol == stateMemberSymbol =>
+            val replacement = switchIds(i)
+            treeCopy.Assign(tree, lhs, Literal(Constant(replacement)))
+          case _: Match | _: CaseDef | _: Block | _: If =>
+            super.transform(tree)
+          case _ => tree
         }
+      }
+
+      private def compactStates(m: Match): Tree = {
         val cases1 = m.cases.flatMap {
           case cd @ CaseDef(Literal(Constant(i: Integer)), EmptyTree, rhs) =>
             val replacement = switchIds(i)

src/test/scala/scala/async/run/late/LateExpansion.scala

@@ -338,6 +338,30 @@ class LateExpansion {
       """)
   }
 
+  @Test def testDot(): Unit = {
+    val result = run(
+      """
+    import scala.async.run.late.{autoawait,lateasync}
+    case class FixedFoo(foo: Int)
+    class Foobar(val foo: Int, val bar: Double) {
+      def guard: Boolean = true
+      @autoawait @lateasync def getValue = 4.2
+      @autoawait @lateasync def func(f: Any) = {
+        ("": Any) match {
+          case (x1, y1) if guard => x1.toString; y1.toString
+          case (x2, y2) if guard => x2.toString; y2.toString
+          case (x3, y3) if guard => x3.toString; y3.toString
+          case (x4, y4) =>
+             getValue; x4.toString; y4.toString
+        }
+      }
+    }
+    object Test {
+      @lateasync def test() = new Foobar(0, 0).func(4)
+    }
+    """)
+  }
+
   def wrapAndRun(code: String): Any = {
     run(
       s"""