Fix decompilation of while loops

Author: nicolasstucki

library/src/scala/tasty/util/ShowSourceCode.scala

@@ -145,6 +145,39 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
             this
         }
 
+      case While(cond, stats) =>
+        this += "while ("
+        printTree(cond)
+        this += ") "
+        stats match {
+          case stat :: Nil =>
+            printTree(stat)
+          case stats =>
+            this += "{"
+            indented {
+              this += lineBreak()
+              printTrees(stats, lineBreak())
+            }
+            this += lineBreak() += "}"
+        }
+
+      case DoWhile(stats, cond) =>
+        this += "do "
+        stats match {
+          case stat :: Nil =>
+            printTree(stat)
+          case stats =>
+            this += "{"
+            indented {
+              this += lineBreak()
+              printTrees(stats, lineBreak())
+            }
+            this += lineBreak() += "}"
+        }
+        this += " while ("
+        printTree(cond)
+        this += ")"
+
       case ddef@DefDef(name, targs, argss, tpt, rhs) =>
         val flags = ddef.flags
         if (flags.isOverride) sb.append("override ")
@@ -225,7 +258,14 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
         this += " = "
         printTree(rhs)
 
-      case Term.Block(stats, expr) =>
+      case Term.Block(stats0, expr) =>
+        def isLoopEntryPoint(tree: Tree): Boolean = tree match {
+          case Term.Apply(Term.Ident("while$" | "doWhile$"), _) => true
+          case _ => false
+        }
+
+        val stats = stats0.filterNot(isLoopEntryPoint)
+
         expr match {
           case Term.Lambda(_, _) =>
             // Decompile lambda from { def annon$(...) = ...; closure(annon$, ...)}
@@ -235,31 +275,22 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
             this += " => "
             printTree(rhs)
             this += ")"
-
-          case Term.Apply(Term.Ident("while$"), _) =>
-            val DefDef("while$", _, _, _, Some(Term.If(cond, Term.Block(body :: Nil, _), _))) = stats.head
-            this += "while ("
-            printTree(cond)
-            this += ") "
-            printTree(body)
-
-          case Term.Apply(Term.Ident("doWhile$"), _) =>
-            val DefDef("doWhile$", _, _, _, Some(Term.Block(List(body), Term.If(cond, _, _)))) = stats.head
-            this += "do "
-            printTree(body)
-            this += " while ("
-            printTree(cond)
-            this += ")"
-
+          case expr if isLoopEntryPoint(expr) && stats.size == 1 =>
+            // Print { def while$() = ...; while$() }
+            // as while (...) ...
+            // instead of { while (...) ... }
+            printTree(stats.head)
           case _ =>
             this += "{"
             indented {
               if (!stats.isEmpty) {
                 this += lineBreak()
                 printTrees(stats, lineBreak())
               }
-              this += lineBreak()
-              printTree(expr)
+              if (!isLoopEntryPoint(expr)) {
+                this += lineBreak()
+                printTree(expr)
+              }
             }
             this += lineBreak() += "}"
         }
@@ -719,6 +750,20 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
     }
   }
 
+  private object While {
+    def unapply(arg: Tree)(implicit ctx: Context): Option[(Term, List[Statement])] = arg match {
+      case DefDef("while$", _, _, _, Some(Term.If(cond, Term.Block(bodyStats, _), _))) => Some((cond, bodyStats))
+      case _ => None
+    }
+  }
+
+  private object DoWhile {
+    def unapply(arg: Tree)(implicit ctx: Context): Option[(List[Statement], Term)] = arg match {
+      case DefDef("doWhile$", _, _, _, Some(Term.Block(body, Term.If(cond, _, _)))) => Some((body, cond))
+      case _ => None
+    }
+  }
+
   // TODO Provide some of these in scala.tasty.Tasty.scala and implement them using checks on symbols for performance
   private object Types {
 

tests/run/quote-inline-function.check

@@ -1,31 +1,31 @@
 Normal function
 {
-  var i: Int = 0
-  val j: Int = 5
-  <label> def while$(): Unit = 
-    if i.<(j) then 
-      {
-        val x$1: Int = i
-        f.apply(x$1)
-        i = i.+(1)
-        while$()
-      }
-     else ()
-  while$()
+  var i: scala.Int = 0
+  val j: scala.Int = 5
+  while (i.<(j)) {
+    val x$1: scala.Int = i
+    f.apply(x$1)
+    i = i.+(1)
+  }
+  do {
+    val x$2: scala.Int = i
+    f.apply(x$2)
+    i = i.+(1)
+  } while (i.<(j))
 }
 
 By name function
 {
-  var i: Int = 0
-  val j: Int = 5
-  <label> def while$(): Unit = 
-    if i.<(j) then 
-      {
-        val x$1: Int = i
-        scala.Predef.println(x$1)
-        i = i.+(1)
-        while$()
-      }
-     else ()
-  while$()
+  var i: scala.Int = 0
+  val j: scala.Int = 5
+  while (i.<(j)) {
+    val x$1: scala.Int = i
+    scala.Predef.println(x$1)
+    i = i.+(1)
+  }
+  do {
+    val x$2: scala.Int = i
+    scala.Predef.println(x$2)
+    i = i.+(1)
+  } while (i.<(j))
 }

tests/run/quote-inline-function/quoted_1.scala

@@ -15,6 +15,10 @@ object Macros {
         ~f.apply('(i))
         i += 1
       }
+      do {
+        ~f.apply('(i))
+        i += 1
+      } while (i < j)
     }
     res.show.toExpr
   }