Lamda context extraction / PoC for zio

lambda-context.scala

/*
libraryDependencies += "org.ow2.asm" % "asm" % "7.1"

Java 8+ only

file:line extraction will work for lambdas only. For anonymous classes it will point to the class definition
*/
import java.lang.invoke.SerializedLambda

import org.objectweb.asm._

import scala.collection.GenTraversableOnce
import scala.collection.generic.CanBuildFrom
import scala.collection.immutable.List
import scala.util.{Failure, Success, Try}

class CV(m: String, ms: Option[String], api: Int) extends ClassVisitor(api) {
  var sourcefile: String = _
  var line: Int = _

  override def visitSource(source: String, debug: String): Unit = {
    super.visitSource(source, debug)
    //println(source, debug)
    sourcefile = source
  }

  override def visitMethod(access: Int, name: String, descriptor: String, signature: String, exceptions: Array[String]): MethodVisitor = {

    if (name == m && ms.forall(_ == descriptor)) {
      new MethodVisitor(api) {
        override def visitLineNumber(line: Int, start: Label): Unit = {
          //println((line, start))
          CV.this.line = line
        }
      }
    } else {
      super.visitMethod(access, name, descriptor, signature, exceptions)
    }
  }
}


object Main {
  var idx = 0
  class X {
    def flatMap(f: X => X): X = {
      run(f)
      f(this)
    }

    def map(f: X => X): X = {
      run(f)
      f(this)
    }
  }


  def run(l: X => X): Option[Unit] = {
    val lcl = l.getClass
    val cl = lcl.getClassLoader
    idx += 1

    //
    import scala.collection.JavaConverters._
    val stackCaller = new RuntimeException().getStackTrace.init.last.getMethodName


    //
    Try(lcl.getDeclaredMethod("writeReplace")) match {
      case Failure(_) =>
        val name = lcl.getName
        val reader = new ClassReader(cl.getResourceAsStream(name.replace('.', '/') + ".class"))
        val cv = new CV("apply", None, Opcodes.ASM7)
        reader.accept(cv, 0)
        println(s"Anon clz $idx invoked from method $stackCaller (${cv.sourcefile}:${cv.line})")

      case Success(replaceMethod) =>
        replaceMethod.setAccessible(true)

        val sl = replaceMethod.invoke(l).asInstanceOf[SerializedLambda]
        val reader = new ClassReader(cl.getResourceAsStream(sl.getImplClass.replace('.', '/') + ".class"))
        val cv = new CV(sl.getImplMethodName, Some(sl.getImplMethodSignature), Opcodes.ASM7)
        reader.accept(cv, 0)
        val dirtyCaller = sl.getImplMethodName.split('$')(2)

        println(s"Lambda $idx invoked from {dirty: ${sl.getImplClass}:$dirtyCaller} ({stack: $stackCaller}) (${cv.sourcefile}:${cv.line})")

        // reflection
        val implclass = cl.loadClass(sl.getImplClass)
        val method = implclass.getMethods.find(_.getName == sl.getImplMethodName)
      //println(implclass, method)

    }

    Some(())
  }

  def x(a: X): X = a

  def main(args: Array[String]): Unit = {
    val f = new Function[X, X] {
      override def apply(v1: X): X = v1
    }

    run(a => a)
    run(x)
    run(f) // points to f defn

    val a = new X()
    for {
      v <- a
      v1 <- v
      v2 <- v1
    } yield {
      v2
    }
  }
}

Output:

    Lambda 1 invoked from {dirty: Main$:main} ({stack: run}) (main.scala:105)
    Lambda 2 invoked from {dirty: Main$:main} ({stack: run}) (main.scala:106)
    Anon clz 3 invoked from method run (main.scala:101)
    Lambda 4 invoked from {dirty: Main$:main} ({stack: run}) (main.scala:112)
    Lambda 5 invoked from {dirty: Main$:main} ({stack: run}) (main.scala:113)
    Lambda 6 invoked from {dirty: Main$:main} ({stack: run}) (main.scala:115)

https://scastie.scala-lang.org/HywGUXj8RXyoVanOMkww5g