Merge "[HostStubGen] General processing speedup" into main

5 files changed, 155 insertions, 187 deletions

diff --git a/ravenwood/tools/hoststubgen/Android.bp b/ravenwood/tools/hoststubgen/Android.bp
index 004834eed983..e605318f4a10 100644
--- a/ravenwood/tools/hoststubgen/Android.bp
+++ b/ravenwood/tools/hoststubgen/Android.bp
@@ -99,6 +99,7 @@ java_library_host {
         "ow2-asm-commons",
         "ow2-asm-tree",
         "ow2-asm-util",
+        "apache-commons-compress",
     ],
 }
 
diff --git a/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/Utils.kt b/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/Utils.kt
index b2af7827f8c5..a62f66dd18e5 100644
--- a/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/Utils.kt
+++ b/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/Utils.kt
@@ -16,6 +16,15 @@
 package com.android.hoststubgen
 
 import java.io.PrintWriter
+import java.util.zip.CRC32
+import org.apache.commons.compress.archivers.zip.ZipArchiveEntry
+import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream
+import org.apache.commons.compress.archivers.zip.ZipFile
+
+/**
+ * Whether to skip compression when adding processed entries back to a zip file.
+ */
+private const val SKIP_COMPRESSION = false
 
 /**
  * Name of this executable. Set it in the main method.
@@ -118,3 +127,29 @@ inline fun runMainWithBoilerplate(realMain: () -> Unit) {
 
     System.exit(if (success) 0 else 1 )
 }
+
+/**
+ * Copy a single ZIP entry to the output.
+ */
+fun copyZipEntry(
+    inZip: ZipFile,
+    entry: ZipArchiveEntry,
+    out: ZipArchiveOutputStream,
+) {
+    inZip.getRawInputStream(entry).use { out.addRawArchiveEntry(entry, it) }
+}
+
+/**
+ * Add a single ZIP entry with data.
+ */
+fun ZipArchiveOutputStream.addBytesEntry(name: String, data: ByteArray) {
+    val newEntry = ZipArchiveEntry(name)
+    if (SKIP_COMPRESSION) {
+        newEntry.method = 0
+        newEntry.size = data.size.toLong()
+        newEntry.crc = CRC32().apply { update(data) }.value
+    }
+    putArchiveEntry(newEntry)
+    write(data)
+    closeArchiveEntry()
+}
diff --git a/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/asm/ClassNodes.kt b/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/asm/ClassNodes.kt
index e2647eb13ed3..40d343ab9658 100644
--- a/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/asm/ClassNodes.kt
+++ b/ravenwood/tools/hoststubgen/lib/com/android/hoststubgen/asm/ClassNodes.kt
@@ -18,21 +18,20 @@ package com.android.hoststubgen.asm
 import com.android.hoststubgen.ClassParseException
 import com.android.hoststubgen.InvalidJarFileException
 import com.android.hoststubgen.log
-import org.objectweb.asm.ClassReader
-import org.objectweb.asm.tree.AnnotationNode
-import org.objectweb.asm.tree.ClassNode
-import org.objectweb.asm.tree.FieldNode
-import org.objectweb.asm.tree.MethodNode
-import org.objectweb.asm.tree.TypeAnnotationNode
-import java.io.BufferedInputStream
 import java.io.PrintWriter
 import java.util.Arrays
 import java.util.concurrent.Executors
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
 import java.util.function.Consumer
-import java.util.zip.ZipEntry
-import java.util.zip.ZipFile
+import org.apache.commons.compress.archivers.zip.ZipArchiveEntry
+import org.apache.commons.compress.archivers.zip.ZipFile
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.tree.AnnotationNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.FieldNode
+import org.objectweb.asm.tree.MethodNode
+import org.objectweb.asm.tree.TypeAnnotationNode
 
 /**
  * Stores all classes loaded from a jar file, in a form of [ClassNode]
@@ -189,7 +188,8 @@ class ClassNodes {
          * Load all the classes, without code.
          */
         fun loadClassStructures(
-            inJar: String,
+            inJar: ZipFile,
+            jarName: String,
             timeCollector: Consumer<Double>? = null,
         ): ClassNodes {
             val allClasses = ClassNodes()
@@ -201,20 +201,19 @@ class ClassNodes {
             val exception = AtomicReference<Throwable>()
 
             // Called on a BG thread. Read a single jar entry and add it to [allClasses].
-            fun parseClass(inZip: ZipFile, entry: ZipEntry) {
+            fun parseClass(inZip: ZipFile, entry: ZipArchiveEntry) {
                 try {
-                    inZip.getInputStream(entry).use { ins ->
-                        val cr = ClassReader(BufferedInputStream(ins))
-                        val cn = ClassNode()
-                        cr.accept(
-                            cn, ClassReader.SKIP_CODE
-                                    or ClassReader.SKIP_DEBUG
-                                    or ClassReader.SKIP_FRAMES
-                        )
-                        synchronized(allClasses) {
-                            if (!allClasses.addClass(cn)) {
-                                log.w("Duplicate class found: ${cn.name}")
-                            }
+                    val classBytes = inZip.getInputStream(entry).use { it.readAllBytes() }
+                    val cr = ClassReader(classBytes)
+                    val cn = ClassNode()
+                    cr.accept(
+                        cn, ClassReader.SKIP_CODE
+                                or ClassReader.SKIP_DEBUG
+                                or ClassReader.SKIP_FRAMES
+                    )
+                    synchronized(allClasses) {
+                        if (!allClasses.addClass(cn)) {
+                            log.w("Duplicate class found: ${cn.name}")
                         }
                     }
                 } catch (e: Throwable) {
@@ -224,35 +223,30 @@ class ClassNodes {
             }
 
             // Actually open the jar and read it on worker threads.
-            val time = log.iTime("Reading class structure from $inJar") {
+            val time = log.iTime("Reading class structure from $jarName") {
                 log.withIndent {
-                    ZipFile(inJar).use { inZip ->
-                        val inEntries = inZip.entries()
-
-                        while (inEntries.hasMoreElements()) {
-                            val entry = inEntries.nextElement()
-
-                            if (entry.name.endsWith(".class")) {
-                                executor.submit {
-                                    parseClass(inZip, entry)
-                                }
-                            } else if (entry.name.endsWith(".dex")) {
-                                // Seems like it's an ART jar file. We can't process it.
-                                // It's a fatal error.
-                                throw InvalidJarFileException(
-                                    "$inJar is not a desktop jar file."
-                                            + " It contains a *.dex file."
-                                )
-                            } else {
-                                // Unknown file type. Skip.
+                    inJar.entries.asSequence().forEach { entry ->
+                        if (entry.name.endsWith(".class")) {
+                            executor.submit {
+                                parseClass(inJar, entry)
                             }
+                        } else if (entry.name.endsWith(".dex")) {
+                            // Seems like it's an ART jar file. We can't process it.
+                            // It's a fatal error.
+                            throw InvalidJarFileException(
+                                "$jarName is not a desktop jar file."
+                                        + " It contains a *.dex file."
+                            )
+                        } else {
+                            // Unknown file type. Skip.
                         }
-                        // Wait for all the work to complete. (must do it before closing the zip)
-                        log.i("Waiting for all loaders to finish...")
-                        executor.shutdown()
-                        executor.awaitTermination(5, TimeUnit.MINUTES)
-                        log.i("All loaders to finished.")
                     }
+
+                    // Wait for all the work to complete. (must do it before closing the zip)
+                    log.i("Waiting for all loaders to finish...")
+                    executor.shutdown()
+                    executor.awaitTermination(5, TimeUnit.MINUTES)
+                    log.i("All loaders to finished.")
                 }
 
                 // If any exception is detected, throw it.
@@ -261,13 +255,13 @@ class ClassNodes {
                 }
 
                 if (allClasses.size == 0) {
-                    log.w("$inJar contains no *.class files.")
+                    log.w("$jarName contains no *.class files.")
                 } else {
-                    log.i("Loaded ${allClasses.size} classes from $inJar.")
+                    log.i("Loaded ${allClasses.size} classes from $jarName.")
                 }
             }
             timeCollector?.accept(time)
             return allClasses
         }
     }
-}
\ No newline at end of file
+}
diff --git a/ravenwood/tools/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt b/ravenwood/tools/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
index 333540573364..2edcb2a6c199 100644
--- a/ravenwood/tools/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
+++ b/ravenwood/tools/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
@@ -19,13 +19,11 @@ import com.android.hoststubgen.asm.ClassNodes
 import com.android.hoststubgen.dumper.ApiDumper
 import com.android.hoststubgen.filters.FilterPolicy
 import com.android.hoststubgen.filters.printAsTextPolicy
-import java.io.BufferedInputStream
-import java.io.BufferedOutputStream
 import java.io.FileOutputStream
 import java.io.PrintWriter
-import java.util.zip.ZipEntry
-import java.util.zip.ZipFile
-import java.util.zip.ZipOutputStream
+import org.apache.commons.compress.archivers.zip.ZipArchiveEntry
+import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream
+import org.apache.commons.compress.archivers.zip.ZipFile
 
 /**
  * Actual main class.
@@ -34,9 +32,10 @@ class HostStubGen(val options: HostStubGenOptions) {
     fun run() {
         val errors = HostStubGenErrors()
         val stats = HostStubGenStats()
+        val inJar = ZipFile(options.inJar.get)
 
         // Load all classes.
-        val allClasses = ClassNodes.loadClassStructures(options.inJar.get)
+        val allClasses = ClassNodes.loadClassStructures(inJar, options.inJar.get)
 
         // Dump the classes, if specified.
         options.inputJarDumpFile.ifSet {
@@ -59,7 +58,7 @@ class HostStubGen(val options: HostStubGenOptions) {
         val processor = HostStubGenClassProcessor(options, allClasses, errors, stats)
 
         // Transform the jar.
-        convert(
+        inJar.convert(
             options.inJar.get,
             options.outJar.get,
             processor,
@@ -88,7 +87,7 @@ class HostStubGen(val options: HostStubGenOptions) {
     /**
      * Convert a JAR file into "stub" and "impl" JAR files.
      */
-    private fun convert(
+    private fun ZipFile.convert(
         inJar: String,
         outJar: String?,
         processor: HostStubGenClassProcessor,
@@ -100,45 +99,39 @@ class HostStubGen(val options: HostStubGenOptions) {
         log.i("ASM CheckClassAdapter is %s", if (enableChecker) "enabled" else "disabled")
 
         log.iTime("Transforming jar") {
-            var itemIndex = 0
             var numItemsProcessed = 0
             var numItems = -1 // == Unknown
 
             log.withIndent {
-                // Open the input jar file and process each entry.
-                ZipFile(inJar).use { inZip ->
-
-                    numItems = inZip.size()
-                    val shardStart = numItems * shard / numShards
-                    val shardNextStart = numItems * (shard + 1) / numShards
-
-                    maybeWithZipOutputStream(outJar) { outStream ->
-                        val inEntries = inZip.entries()
-                        while (inEntries.hasMoreElements()) {
-                            val entry = inEntries.nextElement()
-                            val inShard = (shardStart <= itemIndex)
-                                    && (itemIndex < shardNextStart)
-                            itemIndex++
-                            if (!inShard) {
-                                continue
-                            }
-                            convertSingleEntry(inZip, entry, outStream, processor)
-                            numItemsProcessed++
+                val entries = entries.toList()
+
+                numItems = entries.size
+                val shardStart = numItems * shard / numShards
+                val shardNextStart = numItems * (shard + 1) / numShards
+
+                maybeWithZipOutputStream(outJar) { outStream ->
+                    entries.forEachIndexed { itemIndex, entry ->
+                        val inShard = (shardStart <= itemIndex)
+                                && (itemIndex < shardNextStart)
+                        if (!inShard) {
+                            return@forEachIndexed
                         }
-                        log.i("Converted all entries.")
+                        convertSingleEntry(this, entry, outStream, processor)
+                        numItemsProcessed++
                     }
-                    outJar?.let { log.i("Created: $it") }
+                    log.i("Converted all entries.")
                 }
+                outJar?.let { log.i("Created: $it") }
             }
             log.i("%d / %d item(s) processed.", numItemsProcessed, numItems)
         }
     }
 
-    private fun <T> maybeWithZipOutputStream(filename: String?, block: (ZipOutputStream?) -> T): T {
+    private fun <T> maybeWithZipOutputStream(filename: String?, block: (ZipArchiveOutputStream?) -> T): T {
         if (filename == null) {
             return block(null)
         }
-        return ZipOutputStream(BufferedOutputStream(FileOutputStream(filename))).use(block)
+        return ZipArchiveOutputStream(FileOutputStream(filename).buffered()).use(block)
     }
 
     /**
@@ -146,8 +139,8 @@ class HostStubGen(val options: HostStubGenOptions) {
      */
     private fun convertSingleEntry(
         inZip: ZipFile,
-        entry: ZipEntry,
-        outStream: ZipOutputStream?,
+        entry: ZipArchiveEntry,
+        outStream: ZipArchiveOutputStream?,
         processor: HostStubGenClassProcessor
     ) {
         log.d("Entry: %s", entry.name)
@@ -181,32 +174,12 @@ class HostStubGen(val options: HostStubGenOptions) {
     }
 
     /**
-     * Copy a single ZIP entry to the output.
-     */
-    private fun copyZipEntry(
-        inZip: ZipFile,
-        entry: ZipEntry,
-        out: ZipOutputStream,
-    ) {
-        // TODO: It seems like copying entries this way is _very_ slow,
-        // even with out.setLevel(0). Look for other ways to do it.
-
-        inZip.getInputStream(entry).use { ins ->
-            // Copy unknown entries as is to the impl out. (but not to the stub out.)
-            val outEntry = ZipEntry(entry.name)
-            out.putNextEntry(outEntry)
-            ins.transferTo(out)
-            out.closeEntry()
-        }
-    }
-
-    /**
      * Convert a single class.
      */
     private fun processSingleClass(
         inZip: ZipFile,
-        entry: ZipEntry,
-        outStream: ZipOutputStream?,
+        entry: ZipArchiveEntry,
+        outStream: ZipArchiveOutputStream?,
         processor: HostStubGenClassProcessor
     ) {
         val classInternalName = entry.name.replaceFirst("\\.class$".toRegex(), "")
@@ -227,12 +200,10 @@ class HostStubGen(val options: HostStubGenOptions) {
         if (outStream != null) {
             log.v("Creating class: %s Policy: %s", classInternalName, classPolicy)
             log.withIndent {
-                BufferedInputStream(inZip.getInputStream(entry)).use { bis ->
-                    val newEntry = ZipEntry(newName)
-                    outStream.putNextEntry(newEntry)
-                    val classBytecode = bis.readAllBytes()
-                    outStream.write(processor.processClassBytecode(classBytecode))
-                    outStream.closeEntry()
+                inZip.getInputStream(entry).use { zis ->
+                    var classBytecode = zis.readAllBytes()
+                    classBytecode = processor.processClassBytecode(classBytecode)
+                    outStream.addBytesEntry(newName, classBytecode)
                 }
             }
         }
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
index 04e3bda2ba27..8e36323fd495 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
@@ -17,17 +17,17 @@ package com.android.platform.test.ravenwood.ravenizer
 
 import com.android.hoststubgen.GeneralUserErrorException
 import com.android.hoststubgen.HostStubGenClassProcessor
+import com.android.hoststubgen.addBytesEntry
 import com.android.hoststubgen.asm.ClassNodes
 import com.android.hoststubgen.asm.zipEntryNameToClassName
+import com.android.hoststubgen.copyZipEntry
 import com.android.hoststubgen.executableName
 import com.android.hoststubgen.log
 import com.android.platform.test.ravenwood.ravenizer.adapter.RunnerRewritingAdapter
-import java.io.BufferedInputStream
-import java.io.BufferedOutputStream
 import java.io.FileOutputStream
-import java.util.zip.ZipEntry
-import java.util.zip.ZipFile
-import java.util.zip.ZipOutputStream
+import org.apache.commons.compress.archivers.zip.ZipArchiveEntry
+import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream
+import org.apache.commons.compress.archivers.zip.ZipFile
 import org.objectweb.asm.ClassReader
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
@@ -93,13 +93,14 @@ class Ravenizer {
         val stats = RavenizerStats()
 
         stats.totalTime = log.nTime {
-            val allClasses = ClassNodes.loadClassStructures(options.inJar.get) {
+            val inJar = ZipFile(options.inJar.get)
+            val allClasses = ClassNodes.loadClassStructures(inJar, options.inJar.get) {
                 stats.loadStructureTime = it
             }
             val processor = HostStubGenClassProcessor(options, allClasses)
 
-            process(
-                options.inJar.get,
+            inJar.process(
+                options.outJar.get,
                 options.outJar.get,
                 options.enableValidation.get,
                 options.fatalValidation.get,
@@ -111,7 +112,7 @@ class Ravenizer {
         log.i(stats.toString())
     }
 
-    private fun process(
+    private fun ZipFile.process(
         inJar: String,
         outJar: String,
         enableValidation: Boolean,
@@ -138,40 +139,34 @@ class Ravenizer {
         }
 
         stats.totalProcessTime = log.vTime("$executableName processing $inJar") {
-            ZipFile(inJar).use { inZip ->
-                val inEntries = inZip.entries()
-
-                stats.totalEntries = inZip.size()
-
-                ZipOutputStream(BufferedOutputStream(FileOutputStream(outJar))).use { outZip ->
-                    while (inEntries.hasMoreElements()) {
-                        val entry = inEntries.nextElement()
-
-                        if (entry.name.endsWith(".dex")) {
-                            // Seems like it's an ART jar file. We can't process it.
-                            // It's a fatal error.
-                            throw GeneralUserErrorException(
-                                "$inJar is not a desktop jar file. It contains a *.dex file."
-                            )
-                        }
+            ZipArchiveOutputStream(FileOutputStream(outJar).buffered()).use { outZip ->
+                entries.asSequence().forEach { entry ->
+                    stats.totalEntries++
+                    if (entry.name.endsWith(".dex")) {
+                        // Seems like it's an ART jar file. We can't process it.
+                        // It's a fatal error.
+                        throw GeneralUserErrorException(
+                            "$inJar is not a desktop jar file. It contains a *.dex file."
+                        )
+                    }
 
-                        if (stripMockito && entry.name.isMockitoFile()) {
-                            // Skip this entry
-                            continue
-                        }
+                    if (stripMockito && entry.name.isMockitoFile()) {
+                        // Skip this entry
+                        return@forEach
+                    }
 
-                        val className = zipEntryNameToClassName(entry.name)
+                    val className = zipEntryNameToClassName(entry.name)
 
-                        if (className != null) {
-                            stats.totalClasses += 1
-                        }
+                    if (className != null) {
+                        stats.totalClasses += 1
+                    }
 
-                        if (className != null &&
-                            shouldProcessClass(processor.allClasses, className)) {
-                            processSingleClass(inZip, entry, outZip, processor, stats)
-                        } else {
-                            // Too slow, let's use merge_zips to bring back the original classes.
-                            copyZipEntry(inZip, entry, outZip, stats)
+                    if (className != null &&
+                        shouldProcessClass(processor.allClasses, className)) {
+                        processSingleClass(this, entry, outZip, processor, stats)
+                    } else {
+                        stats.totalCopyTime += log.nTime {
+                            copyZipEntry(this, entry, outZip)
                         }
                     }
                 }
@@ -179,53 +174,25 @@ class Ravenizer {
         }
     }
 
-    /**
-     * Copy a single ZIP entry to the output.
-     */
-    private fun copyZipEntry(
-        inZip: ZipFile,
-        entry: ZipEntry,
-        out: ZipOutputStream,
-        stats: RavenizerStats,
-    ) {
-        stats.totalCopyTime += log.nTime {
-            inZip.getInputStream(entry).use { ins ->
-                // Copy unknown entries as is to the impl out. (but not to the stub out.)
-                val outEntry = ZipEntry(entry.name)
-                outEntry.method = 0
-                outEntry.size = entry.size
-                outEntry.crc = entry.crc
-                out.putNextEntry(outEntry)
-
-                ins.transferTo(out)
-
-                out.closeEntry()
-            }
-        }
-    }
-
     private fun processSingleClass(
         inZip: ZipFile,
-        entry: ZipEntry,
-        outZip: ZipOutputStream,
+        entry: ZipArchiveEntry,
+        outZip: ZipArchiveOutputStream,
         processor: HostStubGenClassProcessor,
         stats: RavenizerStats,
     ) {
         stats.processedClasses += 1
-        val newEntry = ZipEntry(entry.name)
-        outZip.putNextEntry(newEntry)
-
-        BufferedInputStream(inZip.getInputStream(entry)).use { bis ->
-            var classBytes = bis.readBytes()
+        inZip.getInputStream(entry).use { zis ->
+            var classBytes = zis.readAllBytes()
             stats.totalRavenizeTime += log.vTime("Ravenize ${entry.name}") {
                 classBytes = ravenizeSingleClass(entry, classBytes, processor.allClasses)
             }
             stats.totalHostStubGenTime += log.vTime("HostStubGen ${entry.name}") {
                 classBytes = processor.processClassBytecode(classBytes)
             }
-            outZip.write(classBytes)
+            // TODO: if the class does not change, use copyZipEntry
+            outZip.addBytesEntry(entry.name, classBytes)
         }
-        outZip.closeEntry()
     }
 
     /**
@@ -237,7 +204,7 @@ class Ravenizer {
     }
 
     private fun ravenizeSingleClass(
-        entry: ZipEntry,
+        entry: ZipArchiveEntry,
         input: ByteArray,
         allClasses: ClassNodes,
     ): ByteArray {