Address Simon (viirya)'s comments

Author: rednaxelafx

core/src/main/scala/org/apache/spark/util/ClosureCleaner.scala

@@ -453,13 +453,13 @@ private[spark] object IndylambdaScalaClosures extends Logging {
     val maybeClosureClass = maybeClosure.getClass
 
     // shortcut the fast check:
-    // indylambda closure classes are generated by Java's LambdaMetafactory, and they're always
-    // synthetic.
-    if (!maybeClosureClass.isSynthetic) return None
+    // 1. indylambda closure classes are generated by Java's LambdaMetafactory, and they're always
+    //    synthetic.
+    // 2. We only care about Serializable closures, so let's check that as well
+    if (!maybeClosureClass.isSynthetic || !maybeClosure.isInstanceOf[Serializable]) return None
 
     val implementedInterfaces = ClassUtils.getAllInterfaces(maybeClosureClass).asScala
-    val isClosureCandidate = implementedInterfaces.exists(_.getName == "scala.Serializable") &&
-      implementedInterfaces.exists(_.getName.startsWith("scala.Function"))
+    val isClosureCandidate = implementedInterfaces.exists(_.getName.startsWith("scala.Function"))
 
     if (isClosureCandidate) {
       try {
@@ -480,7 +480,6 @@ private[spark] object IndylambdaScalaClosures extends Logging {
   def isIndylambdaScalaClosure(lambdaProxy: SerializedLambda): Boolean = {
     lambdaProxy.getImplMethodKind == MethodHandleInfo.REF_invokeStatic &&
       lambdaProxy.getImplMethodName.contains("$anonfun$")
-      // && implements a scala.runtime.java8 functional interface
   }
 
   def inspect(closure: AnyRef): SerializedLambda = {
@@ -514,18 +513,24 @@ private[spark] object IndylambdaScalaClosures extends Logging {
     val implMethodId = MethodIdentifier(
       implClass, lambdaProxy.getImplMethodName, lambdaProxy.getImplMethodSignature)
 
-    val visited = Set[MethodIdentifier[_]](implMethodId)
+    val visited = Set.empty[MethodIdentifier[_]]
     val stack = Stack[MethodIdentifier[_]](implMethodId)
     while (!stack.isEmpty) {
       val currentId = stack.pop
+      visited += currentId
+
       val currentMethodNode = methodsByName(currentId)
       logTrace(s"  scanning $currentId")
       currentMethodNode.accept(new MethodVisitor(ASM7) {
+        // FIXME: record self class name, get a Class[_] for it
+        // val selfClassName: String =
+        // val selfClass: Class[_] =
+
         override def visitFieldInsn(op: Int, owner: String, name: String, desc: String): Unit = {
           if (op == GETFIELD || op == PUTFIELD) {
             val ownerExternalName = owner.replace('/', '.')
             for (cl <- accessedFields.keys if cl.getName == ownerExternalName) {
-              logTrace(s"    found field access $name on $owner")
+              logTrace(s"    found field access $name on $ownerExternalName")
               accessedFields(cl) += name
             }
           }
@@ -534,9 +539,14 @@ private[spark] object IndylambdaScalaClosures extends Logging {
         override def visitMethodInsn(
             op: Int, owner: String, name: String, desc: String, itf: Boolean): Unit = {
           if (owner == implClassInternalName) {
-            logTrace(s"    found intra class call to $owner.$name$desc")
-            stack.push(MethodIdentifier(implClass, name, desc))
+            val ownerExternalName = owner.replace('/', '.')
+            logTrace(s"    found intra class call to $ownerExternalName.$name$desc")
+            val calleeMethodId = MethodIdentifier(implClass, name, desc)
+            if (!visited.contains(calleeMethodId)) {
+              stack.push(calleeMethodId)
+            }
           } else {
+            // FIXME: implement findTransitively
             // keep the same behavior as the original ClosureCleaner
             logTrace(s"    ignoring call to $owner.$name$desc")
           }
@@ -560,7 +570,11 @@ private[spark] object IndylambdaScalaClosures extends Logging {
                 targetHandle.getDesc.startsWith(s"(L$implClassInternalName;")) {
               // this is a lexically nested closure that also captures the enclosing `this`
               logDebug(s"    found inner closure $targetHandle")
-              stack.push(MethodIdentifier(implClass, targetHandle.getName, targetHandle.getDesc))
+              val calleeMethodId =
+                MethodIdentifier(implClass, targetHandle.getName, targetHandle.getDesc)
+              if (!visited.contains(calleeMethodId)) {
+                stack.push(calleeMethodId)
+              }
             }
           }
         }
@@ -577,7 +591,7 @@ private class ReturnStatementFinder(targetMethodName: Option[String] = None)
   override def visitMethod(access: Int, name: String, desc: String,
       sig: String, exceptions: Array[String]): MethodVisitor = {
 
-    // $anonfun$ covers Java 8 lambdas
+    // $anonfun$ covers indylambda closures
     if (name.contains("apply") || name.contains("$anonfun$")) {
       // A method with suffix "$adapted" will be generated in cases like
       // { _:Int => return; Seq()} but not { _:Int => return; true}