Decide which lambda classes to save based on LambdaReifier's invoker and not the lambda's class loader

Fixes #21

Author: luontola

README.md

@@ -152,6 +152,9 @@ Version History
 - Added an optional `-Dretrolambda.includedFiles` parameter to support the
   incremental compilers of build tools
   ([Issue #23](https://github.com/orfjackal/retrolambda/pull/23))
+- Decides which lambda classes to save based on the current class being
+  processed, instead of the class loader that loaded the lambda class
+  ([Issue #21](https://github.com/orfjackal/retrolambda/issues/21))
 
 ### Retrolambda 1.3.0 (2014-06-04)
 

retrolambda/src/main/java/net/orfjackal/retrolambda/LambdaReifier.java

@@ -19,6 +19,7 @@ public class LambdaReifier {
     // be an error if these collections contain more than one element.
     private static final BlockingDeque<Handle> currentLambdaImplMethod = new LinkedBlockingDeque<>(1);
     private static final BlockingDeque<Handle> currentLambdaBridgeMethod = new LinkedBlockingDeque<>(1);
+    private static final BlockingDeque<Class<?>> currentInvoker = new LinkedBlockingDeque<>(1);
     private static final BlockingDeque<Type> currentInvokedType = new LinkedBlockingDeque<>(1);
     private static final BlockingDeque<String> currentLambdaClass = new LinkedBlockingDeque<>(1);
 
@@ -27,6 +28,7 @@ public static LambdaFactoryMethod reifyLambdaClass(Handle lambdaImplMethod, Hand
         try {
             setLambdaImplMethod(lambdaImplMethod);
             setLambdaBridgeMethod(lambdaBridgeMethod);
+            setInvoker(invoker);
             setInvokedType(invokedType);
 
             // Causes the lambda class to be loaded. Retrolambda's Java agent
@@ -51,6 +53,10 @@ private static void setLambdaBridgeMethod(Handle lambdaBridgeMethod) {
         currentLambdaBridgeMethod.push(lambdaBridgeMethod);
     }
 
+    private static void setInvoker(Class<?> lambdaInvoker) {
+        currentInvoker.push(lambdaInvoker);
+    }
+
     private static void setInvokedType(Type invokedType) {
         currentInvokedType.push(invokedType);
     }
@@ -59,6 +65,13 @@ public static void setLambdaClass(String lambdaClass) {
         currentLambdaClass.push(lambdaClass);
     }
 
+    public static boolean isLambdaClassToReify(String className) {
+        Class<?> invoker = currentInvoker.peekFirst();
+        return invoker != null
+                && className.startsWith(Type.getInternalName(invoker))
+                && LambdaNaming.LAMBDA_CLASS.matcher(className).matches();
+    }
+
     public static Handle getLambdaImplMethod() {
         return currentLambdaImplMethod.getFirst();
     }
@@ -76,6 +89,7 @@ public static LambdaFactoryMethod getLambdaFactoryMethod() {
     private static void resetGlobals() {
         currentLambdaImplMethod.clear();
         currentLambdaBridgeMethod.clear();
+        currentInvoker.clear();
         currentInvokedType.clear();
         currentLambdaClass.clear();
     }

retrolambda/src/main/java/net/orfjackal/retrolambda/LambdaSavingClassFileTransformer.java

@@ -9,38 +9,31 @@
 import java.lang.instrument.*;
 import java.nio.file.*;
 import java.security.ProtectionDomain;
-import java.util.*;
 
 public class LambdaSavingClassFileTransformer implements ClassFileTransformer {
 
     private final Path outputDir;
     private final int targetVersion;
-    private final List<ClassLoader> ignoredClassLoaders = new ArrayList<>();
 
     public LambdaSavingClassFileTransformer(Path outputDir, int targetVersion) {
         this.outputDir = outputDir;
         this.targetVersion = targetVersion;
-        for (ClassLoader cl = ClassLoader.getSystemClassLoader(); cl != null; cl = cl.getParent()) {
-            ignoredClassLoaders.add(cl);
-        }
-        ignoredClassLoaders.add(null);
     }
 
     @Override
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
-        if (ignoredClassLoaders.contains(loader)) {
-            // Avoid saving any classes from the JDK or Retrolambda itself.
-            // The transformed application classes have their own class loader.
-            return null;
-        }
         if (className == null) {
             // Since JDK 8 build b121 or so, lambda classes have a null class name,
             // but we can read it from the bytecode where the name still exists.
             className = new ClassReader(classfileBuffer).getClassName();
         }
-        if (!isLambdaClass(className)) {
-            return null;
+        if (LambdaReifier.isLambdaClassToReify(className)) {
+            reifyLambdaClass(className, classfileBuffer);
         }
+        return null;
+    }
+
+    private void reifyLambdaClass(String className, byte[] classfileBuffer) {
         try {
             System.out.println("Saving lambda class: " + className);
             byte[] backportedBytecode = LambdaClassBackporter.transform(classfileBuffer, targetVersion);
@@ -53,10 +46,5 @@ public byte[] transform(ClassLoader loader, String className, Class<?> classBein
             System.out.println("ERROR: Failed to backport lambda class: " + className);
             t.printStackTrace(System.out);
         }
-        return null;
-    }
-
-    private static boolean isLambdaClass(String className) {
-        return LambdaNaming.LAMBDA_CLASS.matcher(className).matches();
     }
 }