[New IC] Optimize Java class snapshotting with ASM ClassWriter

To snapshot a Java class (+ its fields and methods), previously we used Gson to serialize a class field/method to a string via reflection, and hash that string. We now use an ASM ClassWriter to write a placeholder class containing the field/method of interest and hash the bytecode of that class. One experiment showed that this new approach is ~10 times faster than the previous approach (140s down to 16s when snapshotting 600 jars). Test: Updated expectation files for JavaClassSnapshotterTest unit tests + Existing integration tests to prevent regression ^KT-52141 In Progress
2022-04-21 16:43:58 +01:00
parent 5a0c3920a5
commit 9eb3c7ed76
8 changed files with 74 additions and 96 deletions
@@ -30,13 +30,13 @@ import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache
 import org.jetbrains.kotlin.load.kotlin.incremental.components.JvmPackagePartProto
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.jvm.deserialization.BitEncoding
+import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMemberSignature
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
 import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping
 import org.jetbrains.kotlin.metadata.jvm.serialization.JvmStringTable
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.ClassReader.*
@@ -728,13 +728,13 @@ private fun getConstantsAndInlineFunctions(
    classContents: ByteArray
 ): Pair<LinkedHashMap<String, Any>, LinkedHashMap<String, Long>> {
    val constantsClassVisitor = ConstantsClassVisitor()
-    val inlineFunctionNames = inlineFunctionsJvmNames(classHeader)
+    val inlineFunctionSignatures = inlineFunctionsJvmNames(classHeader)

-    return if (inlineFunctionNames.isEmpty()) {
+    return if (inlineFunctionSignatures.isEmpty()) {
        ClassReader(classContents).accept(constantsClassVisitor, SKIP_CODE or SKIP_DEBUG or SKIP_FRAMES)
        Pair(constantsClassVisitor.getResult(), LinkedHashMap())
    } else {
-        val inlineFunctionsClassVisitor = InlineFunctionsClassVisitor(inlineFunctionNames, constantsClassVisitor)
+        val inlineFunctionsClassVisitor = InlineFunctionsClassVisitor(inlineFunctionSignatures, constantsClassVisitor)
        ClassReader(classContents).accept(inlineFunctionsClassVisitor, 0)
        Pair(constantsClassVisitor.getResult(), inlineFunctionsClassVisitor.getResult())
    }
@@ -757,33 +757,35 @@ private class ConstantsClassVisitor : ClassVisitor(Opcodes.API_VERSION) {
 }

 private class InlineFunctionsClassVisitor(
-    private val inlineFunctionNames: Set<String>,
+    private val inlineFunctionSignatures: Set<String>,
    cv: ConstantsClassVisitor // Note: cv must not override the visitMethod (it will not be called with the current implementation below)
 ) : ClassVisitor(Opcodes.API_VERSION, cv) {

    private val result = LinkedHashMap<String, Long>()
-    private var dummyVersion: Int = -1
+    private var classVersion: Int? = null

-    override fun visit(version: Int, access: Int, name: String?, signature: String?, superName: String?, interfaces: Array<out String>?) {
+    override fun visit(version: Int, access: Int, name: String, signature: String?, superName: String?, interfaces: Array<out String>?) {
        super.visit(version, access, name, signature, superName, interfaces)
-        dummyVersion = version
+        classVersion = version
    }

    override fun visitMethod(access: Int, name: String, desc: String, signature: String?, exceptions: Array<out String>?): MethodVisitor? {
        if (access and Opcodes.ACC_PRIVATE == Opcodes.ACC_PRIVATE) return null

-        val dummyClassWriter = ClassWriter(0)
-        dummyClassWriter.visit(dummyVersion, 0, "dummy", null, AsmTypes.OBJECT_TYPE.internalName, null)
+        // Note: Here, functionSignature = name + descriptor.
+        // It is different from the `signature` parameter above, which is essentially a more detailed descriptor when generics are used
+        // (or null otherwise).
+        val functionSignature = JvmMemberSignature.Method(name, desc).asString()
+        if (functionSignature !in inlineFunctionSignatures) return null

-        return object : MethodVisitor(Opcodes.API_VERSION, dummyClassWriter.visitMethod(0, name, desc, null, exceptions)) {
+        val classWriter = ClassWriter(0)
+
+        // The `version` and `name` parameters are important (see KT-38857), the others can be null.
+        classWriter.visit(/* version */ classVersion!!, /* access */ 0, /* name */ "ClassWithOneMethod", null, null, null)
+
+        return object : MethodVisitor(Opcodes.API_VERSION, classWriter.visitMethod(access, name, desc, signature, exceptions)) {
            override fun visitEnd() {
-                val jvmName = name + desc
-                if (jvmName !in inlineFunctionNames) return
-
-                val dummyBytes = dummyClassWriter.toByteArray()!!
-
-                val hash = dummyBytes.md5()
-                result[jvmName] = hash
+                result[functionSignature] = classWriter.toByteArray().md5()
            }
        }
    }
@@ -18,7 +18,6 @@ dependencies {
    api(project(":compiler:backend.jvm.entrypoint"))
    api(project(":kotlin-build-common"))
    api(project(":daemon-common"))
-    implementation(commonDependency("com.google.code.gson:gson"))
    compileOnly(intellijCore())

    testApi(commonDependency("junit:junit"))
@@ -30,6 +29,7 @@ dependencies {
    testApi(commonDependency("org.jetbrains.intellij.deps:log4j"))
    testApi(commonDependency("org.jetbrains.intellij.deps:jdom"))

+    testImplementation(commonDependency("com.google.code.gson:gson"))
    testRuntimeOnly(project(":kotlin-reflect"))
    testRuntimeOnly(project(":core:descriptors.runtime"))
 }
@@ -114,7 +114,7 @@ class JavaClassMemberLevelSnapshot(
 )

 /** Snapshot of a Java class or a Java class member (field or method). */
-open class JavaElementSnapshot(
+class JavaElementSnapshot(

    /** The name of the Java element. It is part of the Java element's ABI. */
    val name: String,
@@ -123,17 +123,6 @@ open class JavaElementSnapshot(
    val abiHash: Long
 )

-/** TEST-ONLY: A [JavaElementSnapshot] that is used for testing only and must not be used in production code. */
-class JavaElementSnapshotForTests(
-    name: String,
-    abiHash: Long,
-
-    /** The Java element's ABI, captured in a [String]. */
-    @Suppress("unused") // Used by Gson reflection
-    val abiValue: String
-
-) : JavaElementSnapshot(name, abiHash)
-
 /**
 * [ClassSnapshot] of an inaccessible class.
 *
@@ -59,7 +59,6 @@ object ClassSnapshotter {
    fun snapshot(
        classes: List<ClassFileWithContents>,
        granularity: ClassSnapshotGranularity = CLASS_MEMBER_LEVEL,
-        includeDebugInfoInJavaSnapshot: Boolean = false,
        metrics: BuildMetricsReporter = DoNothingBuildMetricsReporter
    ): List<ClassSnapshot> {
        val classesInfo: List<BasicClassInfo> = metrics.measure(BuildTime.READ_CLASSES_BASIC_INFO) {
@@ -75,7 +74,7 @@ object ClassSnapshotter {
                    snapshotKotlinClass(it, granularity)
                }
                else -> metrics.measure(BuildTime.SNAPSHOT_JAVA_CLASSES) {
-                    JavaClassSnapshotter.snapshot(it, granularity, includeDebugInfoInJavaSnapshot)
+                    JavaClassSnapshotter.snapshot(it, granularity)
                }
            }
        }
@@ -86,7 +85,7 @@ object ClassSnapshotter {
        val kotlinClassInfo =
            KotlinClassInfo.createFrom(classFile.classInfo.classId, classFile.classInfo.kotlinClassHeader!!, classFile.contents)
        val classId = kotlinClassInfo.classId
-        val classAbiHash = KotlinClassInfoExternalizer.toByteArray(kotlinClassInfo).md5()
+        val classAbiHash = KotlinClassInfoExternalizer.toByteArray(kotlinClassInfo).hashToLong()
        val classMemberLevelSnapshot = kotlinClassInfo.takeIf { granularity == CLASS_MEMBER_LEVEL }

        return when (kotlinClassInfo.classKind) {
@@ -209,3 +208,9 @@ private object DirectoryOrJarContentsReader {
        return relativePathsToContents.toSortedMap().toMap(LinkedHashMap())
    }
 }
+
+internal fun ByteArray.hashToLong(): Long {
+    // Note: md5 is 128-bit while Long is 64-bit.
+    // Use md5 for now until we find a better 64-bit hash function.
+    return md5()
+}
@@ -5,22 +5,18 @@

 package org.jetbrains.kotlin.incremental.classpathDiff

-import com.google.gson.GsonBuilder
 import org.jetbrains.kotlin.incremental.classpathDiff.ClassSnapshotGranularity.CLASS_MEMBER_LEVEL
-import org.jetbrains.kotlin.incremental.md5
-import org.jetbrains.kotlin.incremental.storage.toByteArray
 import org.jetbrains.org.objectweb.asm.ClassReader
+import org.jetbrains.org.objectweb.asm.ClassWriter
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.tree.ClassNode
+import org.jetbrains.org.objectweb.asm.tree.FieldNode
+import org.jetbrains.org.objectweb.asm.tree.MethodNode

 /** Computes a [JavaClassSnapshot] of a Java class. */
 object JavaClassSnapshotter {

-    fun snapshot(
-        classFile: ClassFileWithContents,
-        granularity: ClassSnapshotGranularity,
-        includeDebugInfoInSnapshot: Boolean
-    ): JavaClassSnapshot {
+    fun snapshot(classFile: ClassFileWithContents, granularity: ClassSnapshotGranularity): JavaClassSnapshot {
        // We will extract ABI information from the given class and store it into the `abiClass` variable.
        // It is acceptable to collect more info than required, but it is incorrect to collect less info than required.
        // There are 2 approaches:
@@ -51,50 +47,49 @@ object JavaClassSnapshotter {
        abiClass.methods.sortWith(compareBy({ it.name }, { it.desc }))

        // Snapshot the class
-        val fieldsAbi = abiClass.fields.map { snapshotJavaElement(it, it.name, includeDebugInfoInSnapshot) }
-        val methodsAbi = abiClass.methods.map { snapshotJavaElement(it, it.name, includeDebugInfoInSnapshot) }
+        val classAbiHash = snapshotClass(abiClass)
+        val classMemberLevelSnapshot = if (granularity == CLASS_MEMBER_LEVEL) {
+            val fieldsAbi = abiClass.fields.map { JavaElementSnapshot(it.name, snapshotField(it)) }
+            val methodsAbi = abiClass.methods.map { JavaElementSnapshot(it.name, snapshotMethod(it)) }
+            val classAbiExcludingMembers = abiClass.let {
+                it.fields.clear()
+                it.methods.clear()
+                JavaElementSnapshot(it.name, snapshotClass(it))
+            }
+            JavaClassMemberLevelSnapshot(classAbiExcludingMembers, fieldsAbi, methodsAbi)
+        } else null

-        abiClass.fields.clear()
-        abiClass.methods.clear()
-        val classAbiExcludingMembers = abiClass.let { snapshotJavaElement(it, it.name, includeDebugInfoInSnapshot) }
-
-        val detailedSnapshot = JavaClassMemberLevelSnapshot(classAbiExcludingMembers, fieldsAbi, methodsAbi)
        return JavaClassSnapshot(
            classId = classFile.classInfo.classId,
-            classAbiHash = JavaClassMemberLevelSnapshotExternalizer.toByteArray(detailedSnapshot).md5(),
-            classMemberLevelSnapshot = detailedSnapshot.takeIf { granularity == CLASS_MEMBER_LEVEL },
+            classAbiHash = classAbiHash,
+            classMemberLevelSnapshot = classMemberLevelSnapshot,
            supertypes = classFile.classInfo.supertypes
        )
    }

-    private val gson by lazy {
-        // Use serializeSpecialFloatingPointValues() to avoid this error
-        //    "java.lang.IllegalArgumentException: NaN is not a valid double value as per JSON specification. To override this behavior, use
-        //    GsonBuilder.serializeSpecialFloatingPointValues() method."
-        // on jars such as ~/.gradle/kotlin-build-dependencies/repo/kotlin.build/ideaIC/203.8084.24/artifacts/lib/rhino-1.7.12.jar.
-        GsonBuilder().serializeSpecialFloatingPointValues().create()
+    private fun snapshotClass(classNode: ClassNode): Long {
+        val classWriter = ClassWriter(0)
+        classNode.accept(classWriter)
+        return classWriter.toByteArray().hashToLong()
    }

-    // Same as above but with `setPrettyPrinting()`
-    private val gsonForDebug by lazy {
-        GsonBuilder().serializeSpecialFloatingPointValues()
-            .setPrettyPrinting()
-            .create()
+    private fun snapshotField(fieldNode: FieldNode): Long {
+        val classNode = emptyClass()
+        classNode.fields.add(fieldNode)
+        return snapshotClass(classNode)
    }

-    private fun snapshotJavaElement(
-        javaElement: Any,
-        javaElementName: String,
-        includeDebugInfoInSnapshot: Boolean
-    ): JavaElementSnapshot {
-        return if (includeDebugInfoInSnapshot) {
-            val abiValue = gsonForDebug.toJson(javaElement)
-            val abiHash = abiValue.toByteArray().md5()
-            JavaElementSnapshotForTests(javaElementName, abiHash, abiValue)
-        } else {
-            val abiValue = gson.toJson(javaElement)
-            val abiHash = abiValue.toByteArray().md5()
-            JavaElementSnapshot(javaElementName, abiHash)
-        }
+    private fun snapshotMethod(methodNode: MethodNode): Long {
+        val classNode = emptyClass()
+        classNode.methods.add(methodNode)
+        return snapshotClass(classNode)
+    }
+
+    private fun emptyClass() = ClassNode().also {
+        // We need to provide some minimal info to the class:
+        //   - Name is required.
+        //   - Class version is required if method bodies are considered, but we have removed method bodies in this class, so it's optional.
+        //   - Other info is optional.
+        it.name = "EmptyClass"
    }
 }
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.incremental.classpathDiff

-import org.jetbrains.kotlin.incremental.classpathDiff.ClasspathSnapshotTestCommon.ClassFileUtil.readBytes
 import org.jetbrains.kotlin.incremental.classpathDiff.ClasspathSnapshotTestCommon.ClassFileUtil.snapshot
 import org.jetbrains.kotlin.incremental.classpathDiff.ClasspathSnapshotTestCommon.SourceFile.JavaSourceFile
 import org.jetbrains.kotlin.incremental.classpathDiff.ClasspathSnapshotTestCommon.SourceFile.KotlinSourceFile
@@ -101,18 +100,10 @@ class JavaOnlyClasspathSnapshotterTest : ClasspathSnapshotTestCommon() {
        JavaSourceFile(baseDir = File("$testDataDir/java/$testName/src"), relativePath = relativePath), tmpDir
    )

-    private fun TestSourceFile.compileAndSnapshotWithDebugInfo(): ClassSnapshot {
-        val classFile = compileSingle()
-        return ClassSnapshotter.snapshot(
-            listOf(ClassFileWithContents(classFile, classFile.readBytes())),
-            includeDebugInfoInJavaSnapshot = true
-        ).single()
-    }
-
    @Test
    fun testSimpleClass() {
        val sourceFile = getSourceFile("testSimpleClass", "com/example/SimpleClass.java")
-        val actualSnapshot = sourceFile.compileAndSnapshotWithDebugInfo().toGson()
+        val actualSnapshot = sourceFile.compileAndSnapshot().toGson()
        val expectedSnapshot = sourceFile.getExpectedSnapshotFile().readText()

        assertEquals(expectedSnapshot, actualSnapshot)
@@ -12,7 +12,7 @@
    },
    "local": false
  },
-  "classAbiHash": 2419143040117344894,
+  "classAbiHash": -6515999856905133685,
  "supertypes": [
    {
      "internalName": "java/lang/Object"
@@ -12,30 +12,26 @@
    },
    "local": false
  },
-  "classAbiHash": 7253633932031152470,
+  "classAbiHash": -6515999856905133685,
  "classMemberLevelSnapshot": {
    "classAbiExcludingMembers": {
-      "abiValue": "{\n  \"version\": 52,\n  \"access\": 33,\n  \"name\": \"com/example/SimpleClass\",\n  \"superName\": \"java/lang/Object\",\n  \"interfaces\": [],\n  \"sourceFile\": \"SimpleClass.java\",\n  \"innerClasses\": [],\n  \"fields\": [],\n  \"methods\": [],\n  \"api\": 589824\n}",
      "name": "com/example/SimpleClass",
-      "abiHash": -1986494095366785349
+      "abiHash": 8734303838782665320
    },
    "fieldsAbi": [
      {
-        "abiValue": "{\n  \"access\": 1,\n  \"name\": \"publicField\",\n  \"desc\": \"Ljava/lang/String;\",\n  \"api\": 589824\n}",
        "name": "publicField",
-        "abiHash": 1021768301411937004
+        "abiHash": 8370167348140025367
      }
    ],
    "methodsAbi": [
      {
-        "abiValue": "{\n  \"access\": 1,\n  \"name\": \"\\u003cinit\\u003e\",\n  \"desc\": \"()V\",\n  \"exceptions\": [],\n  \"visibleAnnotableParameterCount\": 0,\n  \"invisibleAnnotableParameterCount\": 0,\n  \"instructions\": {\n    \"size\": 0\n  },\n  \"tryCatchBlocks\": [],\n  \"maxStack\": 0,\n  \"maxLocals\": 0,\n  \"localVariables\": [],\n  \"visited\": false,\n  \"api\": 589824\n}",
        "name": "\u003cinit\u003e",
-        "abiHash": 5725188035643409224
+        "abiHash": 2413123887428571534
      },
      {
-        "abiValue": "{\n  \"access\": 1,\n  \"name\": \"publicMethod\",\n  \"desc\": \"()Ljava/lang/String;\",\n  \"exceptions\": [],\n  \"visibleAnnotableParameterCount\": 0,\n  \"invisibleAnnotableParameterCount\": 0,\n  \"instructions\": {\n    \"size\": 0\n  },\n  \"tryCatchBlocks\": [],\n  \"maxStack\": 0,\n  \"maxLocals\": 0,\n  \"localVariables\": [],\n  \"visited\": false,\n  \"api\": 589824\n}",
        "name": "publicMethod",
-        "abiHash": 8210873051092995586
+        "abiHash": 7574889098198027162
      }
    ]
  },