[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

build-common/src/org/jetbrains/kotlin/incremental/IncrementalJvmCache.kt

@@ -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)
+        // Note: Here, functionSignature = name + descriptor.
-        dummyClassWriter.visit(dummyVersion, 0, "dummy", null, AsmTypes.OBJECT_TYPE.internalName, null)
+        // 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
+                result[functionSignature] = classWriter.toByteArray().md5()
-                if (jvmName !in inlineFunctionNames) return
-
-                val dummyBytes = dummyClassWriter.toByteArray()!!
-
-                val hash = dummyBytes.md5()
-                result[jvmName] = hash
             }
         }
     }

compiler/incremental-compilation-impl/build.gradle.kts

@@ -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"))
 }

compiler/incremental-compilation-impl/src/org/jetbrains/kotlin/incremental/classpathDiff/ClasspathSnapshot.kt

@@ -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.
  *

compiler/incremental-compilation-impl/src/org/jetbrains/kotlin/incremental/classpathDiff/ClasspathSnapshotter.kt

@@ -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()
+}

compiler/incremental-compilation-impl/src/org/jetbrains/kotlin/incremental/classpathDiff/JavaClassSnapshotter.kt

@@ -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(
+    fun snapshot(classFile: ClassFileWithContents, granularity: ClassSnapshotGranularity): JavaClassSnapshot {
-        classFile: ClassFileWithContents,
-        granularity: ClassSnapshotGranularity,
-        includeDebugInfoInSnapshot: Boolean
-    ): 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 classAbiHash = snapshotClass(abiClass)
-        val methodsAbi = abiClass.methods.map { snapshotJavaElement(it, it.name, includeDebugInfoInSnapshot) }
+        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(),
+            classAbiHash = classAbiHash,
-            classMemberLevelSnapshot = detailedSnapshot.takeIf { granularity == CLASS_MEMBER_LEVEL },
+            classMemberLevelSnapshot = classMemberLevelSnapshot,
             supertypes = classFile.classInfo.supertypes
         )
     }
 
-    private val gson by lazy {
+    private fun snapshotClass(classNode: ClassNode): Long {
-        // Use serializeSpecialFloatingPointValues() to avoid this error
+        val classWriter = ClassWriter(0)
-        //    "java.lang.IllegalArgumentException: NaN is not a valid double value as per JSON specification. To override this behavior, use
+        classNode.accept(classWriter)
-        //    GsonBuilder.serializeSpecialFloatingPointValues() method."
+        return classWriter.toByteArray().hashToLong()
-        // 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()
     }
 
-    // Same as above but with `setPrettyPrinting()`
+    private fun snapshotField(fieldNode: FieldNode): Long {
-    private val gsonForDebug by lazy {
+        val classNode = emptyClass()
-        GsonBuilder().serializeSpecialFloatingPointValues()
+        classNode.fields.add(fieldNode)
-            .setPrettyPrinting()
+        return snapshotClass(classNode)
-            .create()
     }
 
-    private fun snapshotJavaElement(
+    private fun snapshotMethod(methodNode: MethodNode): Long {
-        javaElement: Any,
+        val classNode = emptyClass()
-        javaElementName: String,
+        classNode.methods.add(methodNode)
-        includeDebugInfoInSnapshot: Boolean
+        return snapshotClass(classNode)
-    ): JavaElementSnapshot {
+    }
-        return if (includeDebugInfoInSnapshot) {
+
-            val abiValue = gsonForDebug.toJson(javaElement)
+    private fun emptyClass() = ClassNode().also {
-            val abiHash = abiValue.toByteArray().md5()
+        // We need to provide some minimal info to the class:
-            JavaElementSnapshotForTests(javaElementName, abiHash, abiValue)
+        //   - Name is required.
-        } else {
+        //   - Class version is required if method bodies are considered, but we have removed method bodies in this class, so it's optional.
-            val abiValue = gson.toJson(javaElement)
+        //   - Other info is optional.
-            val abiHash = abiValue.toByteArray().md5()
+        it.name = "EmptyClass"
-            JavaElementSnapshot(javaElementName, abiHash)
-        }
     }
 }

compiler/incremental-compilation-impl/test/org/jetbrains/kotlin/incremental/classpathDiff/ClasspathSnapshotterTest.kt

@@ -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)

compiler/incremental-compilation-impl/testData/org/jetbrains/kotlin/incremental/classpathDiff/ClasspathSnapshotterTest/java/testSimpleClass/expected-snapshot-CLASS_LEVEL/com/example/SimpleClass.json

@@ -12,7 +12,7 @@
     },
     "local": false
   },
-  "classAbiHash": 2419143040117344894,
+  "classAbiHash": -6515999856905133685,
   "supertypes": [
     {
       "internalName": "java/lang/Object"

compiler/incremental-compilation-impl/testData/org/jetbrains/kotlin/incremental/classpathDiff/ClasspathSnapshotterTest/java/testSimpleClass/expected-snapshot/com/example/SimpleClass.json

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