improve efficiency of no-op desugarings

- skip lambda desugaring when it won't do anything
- skip ASM class writing when no desugarings apply to an input class
also minor improvements to prefix remapping
RELNOTES: none

PiperOrigin-RevId: 164492293

4 files changed, 80 insertions, 39 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryRewriter.java b/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryRewriter.java
index 85c39ac98f..6a6b06f7d6 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryRewriter.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryRewriter.java
@@ -38,10 +38,10 @@ class CoreLibraryRewriter {
    * with a ClassRemapper that prefixes class names of core library classes if prefix is not empty.
    */
   public ClassReader reader(InputStream content) throws IOException {
-    if (prefix.length() != 0) {
-      return new PrefixingClassReader(content);
-    } else {
+    if (prefix.isEmpty()) {
       return new ClassReader(content);
+    } else {
+      return new PrefixingClassReader(content, prefix);
     }
   }
 
@@ -53,7 +53,7 @@ class CoreLibraryRewriter {
     return new UnprefixingClassWriter(flags);
   }
 
-  private static boolean shouldPrefix(String typeName) {
+  static boolean shouldPrefix(String typeName) {
     return (typeName.startsWith("java/") || typeName.startsWith("sun/")) && !except(typeName);
   }
 
@@ -87,26 +87,21 @@ class CoreLibraryRewriter {
     return false;
   }
 
-  /** Prefixes core library class names with prefix */
-  public String prefix(String typeName) {
-    if (prefix.length() > 0 && shouldPrefix(typeName)) {
-      return prefix + typeName;
-    }
-    return typeName;
-  }
-
   /** Removes prefix from class names */
   public String unprefix(String typeName) {
-    if (prefix.length() == 0 || !typeName.startsWith(prefix)) {
+    if (prefix.isEmpty() || !typeName.startsWith(prefix)) {
       return typeName;
     }
     return typeName.substring(prefix.length());
   }
 
   /** ClassReader that prefixes core library class names as they are read */
-  private class PrefixingClassReader extends ClassReader {
-    PrefixingClassReader(InputStream content) throws IOException {
+  private static class PrefixingClassReader extends ClassReader {
+    private final String prefix;
+
+    PrefixingClassReader(InputStream content, String prefix) throws IOException {
       super(content);
+      this.prefix = prefix;
     }
 
     @Override
@@ -122,6 +117,14 @@ class CoreLibraryRewriter {
               });
       super.accept(cv, attrs, flags);
     }
+
+    /** Prefixes core library class names with prefix. */
+    private String prefix(String typeName) {
+      if (shouldPrefix(typeName)) {
+        return prefix + typeName;
+      }
+      return typeName;
+    }
   }
 
   /**
@@ -135,7 +138,7 @@ class CoreLibraryRewriter {
       super(Opcodes.ASM5);
       this.writer = new ClassWriter(flags);
       this.cv = this.writer;
-      if (prefix.length() != 0) {
+      if (!prefix.isEmpty()) {
         this.cv =
             new ClassRemapperWithBugFix(
                 this.cv,
diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/Desugar.java b/src/tools/android/java/com/google/devtools/build/android/desugar/Desugar.java
index 29d820650e..5ba76999ba 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/Desugar.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/Desugar.java
@@ -387,9 +387,6 @@ class Desugar {
         // any danger of accidentally uncompressed resources ending up in an .apk.
         if (filename.endsWith(".class")) {
           ClassReader reader = rewriter.reader(content);
-          InvokeDynamicLambdaMethodCollector lambdaMethodCollector =
-              new InvokeDynamicLambdaMethodCollector();
-          reader.accept(lambdaMethodCollector, ClassReader.SKIP_DEBUG);
           UnprefixingClassWriter writer = rewriter.writer(ClassWriter.COMPUTE_MAXS);
           ClassVisitor visitor =
               createClassVisitorsForClassesInInputs(
@@ -398,10 +395,14 @@ class Desugar {
                   bootclasspathReader,
                   interfaceLambdaMethodCollector,
                   writer,
-                  lambdaMethodCollector.getLambdaMethodsUsedInInvokeDyanmic());
-          reader.accept(visitor, 0);
-
-          outputFileProvider.write(filename, writer.toByteArray());
+                  reader);
+          if (writer == visitor) {
+            // Just copy the input if there are no rewritings
+            outputFileProvider.write(filename, reader.b);
+          } else {
+            reader.accept(visitor, 0);
+            outputFileProvider.write(filename, writer.toByteArray());
+          }
         } else {
           outputFileProvider.copyFrom(filename, inputFiles);
         }
@@ -435,6 +436,11 @@ class Desugar {
     for (Map.Entry<Path, LambdaInfo> lambdaClass : lambdaClasses.entrySet()) {
       try (InputStream bytecode = Files.newInputStream(lambdaClass.getKey())) {
         ClassReader reader = rewriter.reader(bytecode);
+        InvokeDynamicLambdaMethodCollector collector = new InvokeDynamicLambdaMethodCollector();
+        reader.accept(collector, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
+        ImmutableSet<MethodInfo> lambdaMethods = collector.getLambdaMethodsUsedInInvokeDynamics();
+        checkState(lambdaMethods.isEmpty(),
+            "Didn't expect to find lambda methods but found %s", lambdaMethods);
         UnprefixingClassWriter writer =
             rewriter.writer(ClassWriter.COMPUTE_MAXS /*for invoking bridges*/);
         ClassVisitor visitor =
@@ -536,7 +542,7 @@ class Desugar {
       ClassReaderFactory bootclasspathReader,
       Builder<String> interfaceLambdaMethodCollector,
       UnprefixingClassWriter writer,
-      ImmutableSet<MethodInfo> lambdaMethodsUsedInInvokeDynamic) {
+      ClassReader input) {
     ClassVisitor visitor = checkNotNull(writer);
     if (!allowTryWithResources) {
       visitor =
@@ -558,14 +564,23 @@ class Desugar {
           visitor = new InterfaceDesugaring(visitor, bootclasspathReader, store);
         }
       }
-      visitor =
-          new LambdaDesugaring(
-              visitor,
-              loader,
-              lambdas,
-              interfaceLambdaMethodCollector,
-              lambdaMethodsUsedInInvokeDynamic,
-              allowDefaultMethods);
+      // LambdaDesugaring is relatively expensive, so check first whether we need it.  Additionally,
+      // we need to collect lambda methods referenced by invokedynamic instructions up-front anyway.
+      // TODO(kmb): Scan constant pool instead of visiting the class to find bootstrap methods etc.
+      InvokeDynamicLambdaMethodCollector collector = new InvokeDynamicLambdaMethodCollector();
+      input.accept(collector, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
+      ImmutableSet<MethodInfo> methodsUsedInInvokeDynamics =
+          collector.getLambdaMethodsUsedInInvokeDynamics();
+      if (!methodsUsedInInvokeDynamics.isEmpty() || collector.needOuterClassRewrite()) {
+        visitor =
+            new LambdaDesugaring(
+                visitor,
+                loader,
+                lambdas,
+                interfaceLambdaMethodCollector,
+                methodsUsedInInvokeDynamics,
+                allowDefaultMethods);
+      }
     }
     return visitor;
   }
diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/InvokeDynamicLambdaMethodCollector.java b/src/tools/android/java/com/google/devtools/build/android/desugar/InvokeDynamicLambdaMethodCollector.java
index ed650801a1..5d26a0a049 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/InvokeDynamicLambdaMethodCollector.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/InvokeDynamicLambdaMethodCollector.java
@@ -25,15 +25,29 @@ import org.objectweb.asm.Opcodes;
  * <p>Note that this class only collects lambda methods. If the invokedynamic is used for other
  * purposes, the methods used in the instruction are NOT collected.
  */
-public class InvokeDynamicLambdaMethodCollector extends ClassVisitor {
+class InvokeDynamicLambdaMethodCollector extends ClassVisitor {
 
-  private final ImmutableSet.Builder<MethodInfo> lambdaMethodsUsedInInvokeDyanmic =
+  private final ImmutableSet.Builder<MethodInfo> lambdaMethodsUsedInInvokeDynamic =
       ImmutableSet.builder();
+  private boolean needOuterClassRewrite = false;
 
   public InvokeDynamicLambdaMethodCollector() {
     super(Opcodes.ASM5);
   }
 
+  /**
+   * Returns whether the visited class is declared in the scope of a lambda.  In that case
+   * {@link LambdaDesugaring} will want to rewrite the EnclosingMethod attribute of the class.
+   */
+  public boolean needOuterClassRewrite() {
+    return needOuterClassRewrite;
+  }
+
+  /** Returns methods referenced in invokedynamic instructions that use LambdaMetafactory. */
+  public ImmutableSet<MethodInfo> getLambdaMethodsUsedInInvokeDynamics() {
+    return lambdaMethodsUsedInInvokeDynamic.build();
+  }
+
   @Override
   public MethodVisitor visitMethod(
       int access, String name, String desc, String signature, String[] exceptions) {
@@ -41,6 +55,12 @@ public class InvokeDynamicLambdaMethodCollector extends ClassVisitor {
     return new LambdaMethodCollector(mv);
   }
 
+  @Override
+  public void visitOuterClass(String owner, String name, String desc) {
+    needOuterClassRewrite = needOuterClassRewrite || (name != null && name.startsWith("lambda$"));
+    super.visitOuterClass(owner, name, desc);
+  }
+
   private class LambdaMethodCollector extends MethodVisitor {
 
     public LambdaMethodCollector(MethodVisitor dest) {
@@ -54,13 +74,9 @@ public class InvokeDynamicLambdaMethodCollector extends ClassVisitor {
         return;
       }
       Handle handle = (Handle) bsmArgs[1];
-      lambdaMethodsUsedInInvokeDyanmic.add(
+      lambdaMethodsUsedInInvokeDynamic.add(
           MethodInfo.create(handle.getOwner(), handle.getName(), handle.getDesc()));
       super.visitInvokeDynamicInsn(name, desc, bsm, bsmArgs);
     }
   }
-
-  public ImmutableSet<MethodInfo> getLambdaMethodsUsedInInvokeDyanmic() {
-    return lambdaMethodsUsedInInvokeDyanmic.build();
-  }
 }
diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/LambdaDesugaring.java b/src/tools/android/java/com/google/devtools/build/android/desugar/LambdaDesugaring.java
index a360749dc9..b1eefa80bb 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/LambdaDesugaring.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/LambdaDesugaring.java
@@ -49,6 +49,9 @@ import org.objectweb.asm.tree.TypeInsnNode;
  * Visitor that desugars classes with uses of lambdas into Java 7-looking code. This includes
  * rewriting lambda-related invokedynamic instructions as well as fixing accessibility of methods
  * that javac emits for lambda bodies.
+ *
+ * <p>Implementation note: {@link InvokeDynamicLambdaMethodCollector} needs to detect any class
+ * that this visitor may rewrite, as we conditionally apply this visitor based on it.
  */
 class LambdaDesugaring extends ClassVisitor {
 
@@ -150,6 +153,7 @@ class LambdaDesugaring extends ClassVisitor {
     super.visitEnd();
   }
 
+  // If this method changes then InvokeDynamicLambdaMethodCollector may need changes well
   @Override
   public MethodVisitor visitMethod(
       int access, String name, String desc, String signature, String[] exceptions) {
@@ -192,6 +196,7 @@ class LambdaDesugaring extends ClassVisitor {
         : null;
   }
 
+  // If this method changes then InvokeDynamicLambdaMethodCollector may need changes well
   @Override
   public void visitOuterClass(String owner, String name, String desc) {
     if (name != null && name.startsWith("lambda$")) {
@@ -201,6 +206,8 @@ class LambdaDesugaring extends ClassVisitor {
     super.visitOuterClass(owner, name, desc);
   }
 
+  // When adding visitXxx methods here then InvokeDynamicLambdaMethodCollector may need changes well
+
   static String uniqueInPackage(String owner, String name) {
     String suffix = "$" + owner.substring(owner.lastIndexOf('/') + 1);
     // For idempotency, we only attach the package-unique suffix if it isn't there already.  This