[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
@@ -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
      }
    ]
  },