Merge pull request #15 from VirtusLab/lambda-unraveler

Lambda unraveler

src/jmh/scala/com/virtuslab/stacktraces/tests/BasicTest.scala

@@ -37,6 +37,30 @@ object BasicTests:
       val z = 1
     } 
 
+  @main
+  def nestedLambdasRec = TestExecutor.executeTest { () =>
+      def rec(int: Int, string: String): IterableOnce[Any] =
+        val y = 1
+        val x = (0 to 10).flatMap { 
+          n => 
+            if int == 1 then
+              throw new RuntimeException("abc")
+            else if int > 0 then
+              rec(int - 1, "ASD")
+            else
+              def rec(n: Int): IterableOnce[Any] = List(n).map { 
+                n => (if n > 5 then List(true) else List(false)).flatMap {
+                  n => (if n then List("0") else List("5")).map { n => 
+                    22
+                  }
+                }
+              } 
+              rec(n)
+        }
+        x
+      rec(3, "ASD")
+    } 
+
   @main 
   def BdoSth = 
     TestExecutor.executeTest(() => B().doSth)

src/main/scala/com/virtuslab/stacktraces/core/LambdaUnraveler.scala

@@ -0,0 +1,22 @@
+package org.virtuslab.stacktraces.core
+
+import scala.quoted.*
+import scala.tasty.inspector.*
+
+class LambdaUnraveler private (val q: Quotes)(
+  private val previousLambdas: List[q.reflect.DefDef],
+  private val previousFrames: Set[String],
+  private val counter: Int = 0
+):
+  def getNextLambdaAndState(using qd: Quotes)(defdefs: List[qd.reflect.DefDef], frameName: String): (Option[qd.reflect.DefDef], LambdaUnraveler) =
+    if defdefs.nonEmpty && !previousFrames.contains(frameName) && previousLambdas == defdefs then
+      (Some(defdefs.reverse(counter)), new LambdaUnraveler(q)(previousLambdas, previousFrames + frameName, counter + 1))
+    else if defdefs.nonEmpty then
+      (Some(defdefs.last), new LambdaUnraveler(qd)(defdefs, Set(frameName), 1))
+    else
+      (None, new LambdaUnraveler(qd)(defdefs, Set(frameName), 0))
+
+
+object LambdaUnraveler:
+  def apply(q: Quotes)(defdefs: List[q.reflect.DefDef]): LambdaUnraveler =
+    new LambdaUnraveler(q)(defdefs, Set.empty, 0)

src/main/scala/com/virtuslab/stacktraces/core/StacktracesInspector.scala

@@ -7,7 +7,6 @@ import org.virtuslab.stacktraces.model.PrettyErrors
 import org.virtuslab.stacktraces.io.TastyFilesLocator
 import org.virtuslab.stacktraces.tasty.TypesSupport
 
-
 import dotty.tools.dotc.util.NameTransformer
 import dotty.tools.dotc.core.Names
 
@@ -54,7 +53,7 @@ class StacktracesInspector private (st: List[StackTraceElement], ctp: Map[String
       PrettyStackTraceElement(ste, ElementType.Method, ste.getMethodName, ste.getFileName, ste.getLineNumber, error = Some(error))
 
     class Traverser(ste: StackTraceElement) extends TreeAccumulator[List[DefDef]]:
-      def traverseTree(defdefs: List[DefDef], tree: Tree)(owner: Symbol): List[DefDef] = 
+      def foldTree(defdefs: List[DefDef], tree: Tree)(owner: Symbol): List[DefDef] =
         val defdef = tree match 
           case d: DefDef => 
             if d.pos.startLine + 1 <= ste.getLineNumber && d.pos.endLine + 1 >= ste.getLineNumber then
@@ -73,50 +72,68 @@ class StacktracesInspector private (st: List[StackTraceElement], ctp: Map[String
             false
 
         if exists && tree.pos.startLine < ste.getLineNumber then 
-          traverseTreeChildren(defdefs ++ defdef, tree)(owner) 
+          foldOverTree(defdefs ++ defdef, tree)(owner)
         else 
           defdefs
-
-      def foldTree(defdefs: List[DefDef], tree: Tree)(owner: Symbol): List[DefDef] = traverseTree(defdefs, tree)(owner)
-
-      protected def traverseTreeChildren(defdefs: List[DefDef], tree: Tree)(owner: Symbol): List[DefDef] = foldOverTree(defdefs, tree)(owner)
-
     end Traverser
 
-    def processDefDefs(defdefs: List[DefDef])(using ste: StackTraceElement): Option[PrettyStackTraceElement] =
-      val decoded = NameTransformer.decode(Names.termName(ste.getMethodName)).toString
+    def processDefDefs(
+      defdefs: List[DefDef],
+      optionalName: Option[String] = None
+    )(
+      using lambdaUnraveler: LambdaUnraveler, 
+      ste: StackTraceElement
+    ): (Option[PrettyStackTraceElement], LambdaUnraveler) =
+      val decoded = NameTransformer.decode(Names.termName(optionalName.getOrElse(ste.getMethodName))).toString
       decoded match
         case d if d.contains("$anonfun$") =>
           val lambdas = defdefs.filter(f => f.name == "$anonfun")
-          lambdas match
-            case head :: Nil =>
-              Some(createPrettyStackTraceElement(head, ste.getLineNumber))
-            case _ =>
-              Some(createErrorWhileBrowsingTastyFiles(PrettyErrors.InlinedLambda))
+          val (defdef, newLambdaUnraveler) = lambdaUnraveler.getNextLambdaAndState(lambdas, decoded)
+          defdef match
+            case Some(head) =>
+              (Some(createPrettyStackTraceElement(head, ste.getLineNumber)), newLambdaUnraveler)
+            case None =>
+              (Some(createErrorWhileBrowsingTastyFiles(PrettyErrors.InlinedLambda)), newLambdaUnraveler)
         case d =>
           defdefs.filter(_.name != "$anonfun") match
             case Nil =>
-              None
+              (None, lambdaUnraveler)
             case head :: Nil =>
-              Some(createPrettyStackTraceElement(head, ste.getLineNumber))
+              (Some(createPrettyStackTraceElement(head, ste.getLineNumber)), lambdaUnraveler)
             case _ => 
               val fun = defdefs.filter(_.name == d)
               fun match // This will probably fail for nested inline functions, though we cannot disambiguate them
                 case head :: Nil =>
-                  Some(createPrettyStackTraceElement(head, ste.getLineNumber))
-                case defdefs =>
-                  Some(createErrorWhileBrowsingTastyFiles(PrettyErrors.Unknown))
-
-    st.foreach { ste =>
+                  (Some(createPrettyStackTraceElement(head, ste.getLineNumber)), lambdaUnraveler)
+                case _ =>
+                  val extraSuffix = """(.*)\$(\d*)""".r
+                  decoded match
+                    case extraSuffix(name, suffix) =>
+                      processDefDefs(defdefs, Some(name))
+                    case _ =>
+                      (Some(createErrorWhileBrowsingTastyFiles(PrettyErrors.Unknown)), lambdaUnraveler)
+
+    def processStackTrace(ste: StackTraceElement)(using LambdaUnraveler): LambdaUnraveler =
       tastys.find(_.path.stripSuffix(".class") endsWith ctp(ste.getClassName)) match 
         case Some(tasty) =>
           given StackTraceElement = ste
           val tree = tasty.ast
           val traverser = Traverser(ste)
-          val defdefs = traverser.traverseTree(List.empty, tree)(tree.symbol)
-          processDefDefs(defdefs) match
+          val defdefs = traverser.foldTree(List.empty, tree)(tree.symbol)
+          val (pse, newLambdaUnraveler) = processDefDefs(defdefs)
+          pse match
             case Some(e) => prettyStackTraceElements += e
             case None => // do nothing
+          newLambdaUnraveler
         case None =>
           prettyStackTraceElements += PrettyStackTraceElement(ste, ElementType.Method, ste.getMethodName, ste.getClassName, ste.getLineNumber, isTasty = false)
-    }
+          summon[LambdaUnraveler]
+
+    def foreachStacktrace(st: List[StackTraceElement])(using LambdaUnraveler): Unit =
+      val ste :: tail = st
+      val newLambdaUnraveler = processStackTrace(ste)
+      tail match
+        case Nil => // do nothing
+        case _ => foreachStacktrace(tail)(using newLambdaUnraveler)
+
+    foreachStacktrace(st)(using LambdaUnraveler(q)(Nil))