BytecodeAnalyzer.scala

package spark.util

import java.io.{ByteArrayInputStream, ByteArrayOutputStream}

import scala.collection.mutable

import org.objectweb.asm.{ClassReader, MethodVisitor}
import org.objectweb.asm.commons.EmptyVisitor
import org.objectweb.asm.Opcodes._



class Edge1[T](val src: Int, val dst: Int, val data: T)



object BytecodeAnalyzer {

  def main(args: Array[String]) {
    println("!!!")
    val test = {e: Edge1[Long] =>
      println(e.data)
    }
    val c = {e: Edge1[Long] => println(e.data); test(e); }
    val c1 = {e: Edge1[Long] => println(e.src); test(e); }
    val c2 = {e: Edge1[Long] => println(e.src); }
    println(methodInvoked(c, "spark/util/Edge1", "data"))
    println("---------------------------------------------------------")
    println(methodInvoked(c1, "spark/util/Edge1", "data"))
    println("---------------------------------------------------------")
    println(methodInvoked(c2, "spark/util/Edge1", "data"))
  }

  def methodInvoked(closure: AnyRef, className: String, methodName: String): Boolean = {
    val classReader = getClassReader(closure.getClass)
    val finder = new MethodAccessFinder("apply", className, methodName)
    classReader.accept(finder, 0)
    finder.methodInvoked
  }

  private def getClassReader(cls: Class[_]): ClassReader = {
    // Copy data over, before delegating to ClassReader - else we can run out of open file handles.
    val className = cls.getName.replaceFirst("^.*\\.", "") + ".class"
    val resourceStream = cls.getResourceAsStream(className)
    // todo: Fixme - continuing with earlier behavior ...
    if (resourceStream == null) return new ClassReader(resourceStream)

    val baos = new ByteArrayOutputStream(128)
    spark.Utils.copyStream(resourceStream, baos, true)
    new ClassReader(new ByteArrayInputStream(baos.toByteArray))
  }

  private def skipClass(className: String): Boolean = {
    val c = className
    c.startsWith("java/") || c.startsWith("scala/") || c.startsWith("javax/")
  }

  class MethodAccessFinder(
      methodToExam: String,
      methodClass: String,
      methodAccess: String,
      examedMethods: mutable.HashSet[(String, String)] = new mutable.HashSet())
    extends EmptyVisitor with spark.Logging {

    var methodInvoked = false

    override def visitMethod(access: Int, name: String, desc: String,
                             sig: String, exceptions: Array[String]): MethodVisitor = {
      if (name == methodToExam) {
        logError("examing method " + name + " " + desc + " " + sig)
        new EmptyVisitor {
          override def visitMethodInsn(op: Int, owner: String, name: String, desc: String) {
            if (!methodInvoked && op == INVOKEVIRTUAL || op == INVOKESPECIAL || op == INVOKESTATIC || op == INVOKEINTERFACE) {
              logError("calling method " + owner + " " + name + " " + desc)
              if (owner == methodClass && name == methodAccess) {
                methodInvoked = true
              } else if (!skipClass(owner) && !examedMethods.contains((owner, name))) {
                logError("Going to exam method ")
                val classReader = getClassReader(Class.forName(owner.replace("/", ".")))
                examedMethods.add((owner, name))
                val finder = new MethodAccessFinder(name, methodClass, methodAccess, examedMethods)
                classReader.accept(finder, 0)
                methodInvoked |= finder.methodInvoked
              }
            }
          }
        }
      } else {
        null
      }
    }
  }
}