Add collector API for multi-round transform

Author: johnsonlee

booster-transform-util/src/main/kotlin/com/didiglobal/booster/transform/util/Collector.kt

@@ -0,0 +1,69 @@
+package com.didiglobal.booster.transform.util
+
+import com.didiglobal.booster.kotlinx.search
+import com.didiglobal.booster.transform.TransformContext
+import org.apache.commons.compress.archivers.ArchiveEntry
+import org.apache.commons.compress.archivers.ArchiveStreamFactory
+import java.io.IOException
+
+interface Collector<R> {
+    fun accept(name: String): Boolean
+    fun collect(name: String, data: () -> ByteArray): R
+}
+
+sealed class Collectors {
+
+    object ClassNameCollector : Collector<String> {
+        override fun accept(name: String): Boolean {
+            return name.endsWith(".class", true)
+        }
+
+        override fun collect(name: String, data: () -> ByteArray): String {
+            return name.substringBeforeLast('.').replace('/', '.')
+        }
+    }
+
+    object ServiceCollector : Collector<Pair<String, Collection<String>>> {
+        override fun accept(name: String): Boolean {
+            return name.substringBeforeLast('/') == "META-INF/services"
+        }
+
+        override fun collect(name: String, data: () -> ByteArray): Pair<String, Collection<String>> {
+            return name.substringAfterLast('/').replace('/', '.') to data().inputStream().bufferedReader().lineSequence().filterNot {
+                it.isBlank() || it.startsWith('#')
+            }.toMutableSet()
+        }
+    }
+
+}
+
+fun <R> TransformContext.collect(collector: Collector<R>): List<R> = compileClasspath.map { file ->
+    when {
+        file.isDirectory -> {
+            val base = file.toURI()
+            file.search { f ->
+                f.isFile && collector.accept(base.relativize(f.toURI()).normalize().path)
+            }.map { f ->
+                collector.collect(base.relativize(f.toURI()).normalize().path, f::readBytes)
+            }
+        }
+        file.isFile -> {
+            file.inputStream().buffered().use {
+                ArchiveStreamFactory().createArchiveInputStream(it).let { archive ->
+                    generateSequence {
+                        try {
+                            archive.nextEntry
+                        } catch (e: IOException) {
+                            null
+                        }
+                    }.filterNot(ArchiveEntry::isDirectory).filter { entry ->
+                        collector.accept(entry.name)
+                    }.map { entry ->
+                        collector.collect(entry.name, archive::readBytes)
+                    }.toList()
+                }
+            }
+        }
+        else -> emptyList()
+    }
+}.flatten()

booster-transform-util/src/main/kotlin/com/didiglobal/booster/transform/util/transform.kt

@@ -134,14 +134,14 @@ fun ZipInputStream.transform(
 private const val DEFAULT_BUFFER_SIZE = 8 * 1024
 
 private fun InputStream.readBytes(estimatedSize: Int = DEFAULT_BUFFER_SIZE): ByteArray {
-    val buffer = ByteArrayOutputStream(estimatedSize.coerceAtLeast(this.available()))
+    val buffer = ByteArrayOutputStream(maxOf(estimatedSize, this.available()))
     copyTo(buffer)
     return buffer.toByteArray()
 }
 
 private fun InputStream.copyTo(out: OutputStream, bufferSize: Int = DEFAULT_BUFFER_SIZE): Long {
     var bytesCopied: Long = 0
-    val buffer = ByteArray(bufferSize)
+    val buffer = ByteArray(maxOf(bufferSize, this.available()))
     var bytes = read(buffer)
     while (bytes >= 0) {
         out.write(buffer, 0, bytes)

booster-transform-util/src/test/kotlin/com/didiglobal/booster/transform/util/CollectorTest.kt

@@ -0,0 +1,32 @@
+package com.didiglobal.booster.transform.util
+
+import com.didiglobal.booster.transform.AbstractTransformContext
+import java.io.File
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+class CollectorTest {
+
+    @Test
+    fun `collect class names from jar`() {
+        val classpath = listOf(File(javaClass.classLoader.getResource("pinyin4j-2.5.0.jar")!!.file))
+        val context = object : AbstractTransformContext("", "", emptyList(), classpath) {}
+        val classes = context.collect(Collectors.ClassNameCollector)
+        println(classes)
+        assertTrue(classes.isNotEmpty())
+    }
+
+    @Test
+    fun `collect services from dir`() {
+        val classpath = listOf(
+                File(javaClass.protectionDomain.codeSource.location.file),
+                File(javaClass.classLoader.getResource("pinyin4j-2.5.0.jar")!!.file).parentFile
+        )
+        val context = object : AbstractTransformContext("", "", emptyList(), classpath) {}
+        val (api, impl) = context.collect(Collectors.ServiceCollector).single()
+        assertEquals(api, Collector::class.java.name)
+        assertEquals(2, impl.size)
+    }
+
+}

booster-transform-util/src/test/kotlin/com/didiglobal/booster/transform/util/FileTransformTest.kt

@@ -8,8 +8,8 @@ import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
 
-private val PWD: String = System.getProperty("user.dir")
-private val BUILD_DIR = File(PWD, "build")
+internal val PWD: String = System.getProperty("user.dir")
+internal val BUILD_DIR = File(PWD, "build")
 
 class FileTransformTest {
 

booster-transform-util/src/test/resources/META-INF/services/com.didiglobal.booster.transform.util.Collector

@@ -0,0 +1,3 @@
+# This is a comment
+com.didiglobal.booster.transform.util.Collectors$ClassNameCollector
+com.didiglobal.booster.transform.util.Collectors$ServiceCollector