KT-45777: Add custom serialization for classpath snapshot

Also add a few commonly-used classes/methods to externalizers.kt to
allow reuse.

Bug: KT-45777
Test: New ClasspathSnapshotSerializerTest

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

@@ -118,10 +118,10 @@ open class IncrementalJvmCache(
     }
 
     /**
-     * Saves information about the given (kotlinc-generated) class to this cache, and stores changes between this class and its previous
-     * version into the given [ChangesCollector].
+     * Saves information about the given (Kotlin) class to this cache, and stores changes between this class and its previous version into
+     * the given [ChangesCollector].
      *
-     * @param kotlinClassInfo A kotlin-generated class
+     * @param kotlinClassInfo Information about a Kotlin class
      * @param sourceFiles The source files that the given class was generated from, or `null` if this information is not available
      * @param changesCollector A [ChangesCollector]
      */
@@ -129,11 +129,13 @@ open class IncrementalJvmCache(
         val className = kotlinClassInfo.className
 
         dirtyOutputClassesMap.notDirty(className)
-        sourceFiles?.forEach {
+
+        if (sourceFiles != null) {
+            sourceFiles.forEach {
                 sourceToClassesMap.add(it, className)
             }
-
-        sourceFiles?.let { internalNameToSource[className.internalName] = it }
+            internalNameToSource[className.internalName] = sourceFiles
+        }
 
         if (kotlinClassInfo.classId.isLocal) return
 
@@ -149,8 +151,8 @@ open class IncrementalJvmCache(
                 inlineFunctionsMap.process(kotlinClassInfo, changesCollector)
             }
             KotlinClassHeader.Kind.MULTIFILE_CLASS -> {
-                val partNames = kotlinClassInfo.classHeaderData?.toList()
-                    ?: throw AssertionError("Multifile class has no parts: $className")
+                val partNames = kotlinClassInfo.classHeaderData.toList()
+                check(partNames.isNotEmpty()) { "Multifile class has no parts: $className" }
                 multifileFacadeToParts[className] = partNames
                 // When a class is replaced with a facade with the same name,
                 // the class' proto wouldn't ever be deleted,
@@ -315,8 +317,8 @@ open class IncrementalJvmCache(
             val oldData = storage[key]
             val newData = ProtoMapValue(
                 kotlinClassInfo.classKind != KotlinClassHeader.Kind.CLASS,
-                BitEncoding.decodeBytes(kotlinClassInfo.classHeaderData!!),
-                kotlinClassInfo.classHeaderStrings!!
+                BitEncoding.decodeBytes(kotlinClassInfo.classHeaderData),
+                kotlinClassInfo.classHeaderStrings
             )
             storage[key] = newData
 
@@ -380,7 +382,8 @@ open class IncrementalJvmCache(
     }
 
     // todo: reuse code with InlineFunctionsMap?
-    private inner class ConstantsMap(storageFile: File) : BasicStringMap<Map<String, Any>>(storageFile, ConstantsMapExternalizer) {
+    private inner class ConstantsMap(storageFile: File) :
+        BasicStringMap<LinkedHashMap<String, Any>>(storageFile, LinkedHashMapExternalizer(StringExternalizer, ConstantExternalizer)) {
 
         operator fun contains(className: JvmClassName): Boolean =
             className.internalName in storage
@@ -419,7 +422,7 @@ open class IncrementalJvmCache(
             storage.remove(className.internalName)
         }
 
-        override fun dumpValue(value: Map<String, Any>): String =
+        override fun dumpValue(value: LinkedHashMap<String, Any>): String =
             value.dumpMap(Any::toString)
     }
 
@@ -499,12 +502,12 @@ open class IncrementalJvmCache(
     }
 
     private fun addToClassStorage(classInfo: KotlinClassInfo, srcFile: File) {
-        val (nameResolver, proto) = JvmProtoBufUtil.readClassDataFrom(classInfo.classHeaderData!!, classInfo.classHeaderStrings!!)
+        val (nameResolver, proto) = JvmProtoBufUtil.readClassDataFrom(classInfo.classHeaderData, classInfo.classHeaderStrings)
         addToClassStorage(proto, nameResolver, srcFile)
     }
 
     private inner class InlineFunctionsMap(storageFile: File) :
-        BasicStringMap<Map<String, Long>>(storageFile, StringToLongMapExternalizer) {
+        BasicStringMap<LinkedHashMap<String, Long>>(storageFile, LinkedHashMapExternalizer(StringExternalizer, LongExternalizer)) {
 
         @Synchronized
         fun process(kotlinClassInfo: KotlinClassInfo, changesCollector: ChangesCollector) {
@@ -537,7 +540,7 @@ open class IncrementalJvmCache(
             storage.remove(className.internalName)
         }
 
-        override fun dumpValue(value: Map<String, Long>): String =
+        override fun dumpValue(value: LinkedHashMap<String, Long>): String =
             value.dumpMap { java.lang.Long.toHexString(it) }
     }
 }
@@ -596,18 +599,18 @@ fun <T : Comparable<T>> Collection<T>.dumpCollection(): String =
     "[${sorted().joinToString(", ", transform = Any::toString)}]"
 
 /**
- * Minimal information about a kotlinc-generated class that will be used to compute recompilation-triggered changes to support incremental
- * compilation (see [IncrementalJvmCache.saveClassToCache]).
+ * Minimal information about a Kotlin class to compute recompilation-triggering changes during an incremental run of the `KotlinCompile`
+ * task (see [IncrementalJvmCache.saveClassToCache]).
  *
  * It's important that this class contain only the minimal required information, as it will be part of the classpath snapshot of the
  * `KotlinCompile` task and the task needs to support compile avoidance. For example, this class should contain public method signatures,
  * and should not contain private method signatures, or method implementations.
  */
-class KotlinClassInfo private constructor(
+class KotlinClassInfo constructor(
     val classId: ClassId,
     val classKind: KotlinClassHeader.Kind,
-    val classHeaderData: Array<String>?,
-    val classHeaderStrings: Array<String>?,
+    val classHeaderData: Array<String>, // Can be empty
+    val classHeaderStrings: Array<String>, // Can be empty
     @Suppress("SpellCheckingInspection") val multifileClassName: String?,
     val constantsMap: LinkedHashMap<String, Any>,
     val inlineFunctionsMap: LinkedHashMap<String, Long>
@@ -628,22 +631,22 @@ class KotlinClassInfo private constructor(
             return KotlinClassInfo(
                 kotlinClass.classId,
                 kotlinClass.classHeader.kind,
-                kotlinClass.classHeader.data,
-                kotlinClass.classHeader.strings,
+                kotlinClass.classHeader.data ?: emptyArray(),
+                kotlinClass.classHeader.strings ?: emptyArray(),
                 kotlinClass.classHeader.multifileClassName,
                 getConstantsMap(kotlinClass.fileContents),
                 getInlineFunctionsMap(kotlinClass.classHeader, kotlinClass.fileContents)
             )
         }
 
-        /** Creates [KotlinClassInfo] from the given classContents, or returns `null` if the class is not a kotlinc-generated class. */
+        /** Creates [KotlinClassInfo] from the given classContents, or returns `null` if the class is not a Kotlin class. */
         fun tryCreateFrom(classContents: ByteArray): KotlinClassInfo? {
             return FileBasedKotlinClass.create(classContents) { classId, _, classHeader, _ ->
                 KotlinClassInfo(
                     classId,
                     classHeader.kind,
-                    classHeader.data,
-                    classHeader.strings,
+                    classHeader.data ?: emptyArray(),
+                    classHeader.strings ?: emptyArray(),
                     classHeader.multifileClassName,
                     getConstantsMap(classContents),
                     getInlineFunctionsMap(classHeader, classContents)

build-common/src/org/jetbrains/kotlin/incremental/storage/externalizers.kt

@@ -94,13 +94,12 @@ object ProtoMapValueExternalizer : DataExternalizer<ProtoMapValue> {
     }
 }
 
-
 abstract class StringMapExternalizer<T> : DataExternalizer<Map<String, T>> {
     override fun save(output: DataOutput, map: Map<String, T>?) {
         output.writeInt(map!!.size)
 
         for ((key, value) in map.entries) {
-            IOUtil.writeString(key, output)
+            output.writeString(key)
             writeValue(output, value)
         }
     }
@@ -110,7 +109,7 @@ abstract class StringMapExternalizer<T> : DataExternalizer<Map<String, T>> {
         val map = HashMap<String, T>(size)
 
         repeat(size) {
-            val name = IOUtil.readString(input)!!
+            val name = input.readString()
             map[name] = readValue(input)
         }
 
@@ -121,7 +120,6 @@ abstract class StringMapExternalizer<T> : DataExternalizer<Map<String, T>> {
     protected abstract fun readValue(input: DataInput): T
 }
 
-
 object StringToLongMapExternalizer : StringMapExternalizer<Long>() {
     override fun readValue(input: DataInput): Long = input.readLong()
 
@@ -130,12 +128,10 @@ object StringToLongMapExternalizer : StringMapExternalizer<Long>() {
     }
 }
 
-object ConstantsMapExternalizer : DataExternalizer<Map<String, Any>> {
-    override fun save(output: DataOutput, map: Map<String, Any>?) {
-        output.writeInt(map!!.size)
-        for (name in map.keys.sorted()) {
-            IOUtil.writeString(name, output)
-            val value = map[name]!!
+/** [DataExternalizer] for a Kotlin constant. */
+object ConstantExternalizer : DataExternalizer<Any> {
+
+    override fun save(output: DataOutput, value: Any) {
         when (value) {
             is Int -> {
                 output.writeByte(Kind.INT.ordinal)
@@ -155,33 +151,20 @@ object ConstantsMapExternalizer : DataExternalizer<Map<String, Any>> {
             }
             is String -> {
                 output.writeByte(Kind.STRING.ordinal)
-                    IOUtil.writeString(value, output)
+                output.writeString(value)
             }
             else -> throw IllegalStateException("Unexpected constant class: ${value::class.java}")
         }
     }
-    }
 
-    override fun read(input: DataInput): Map<String, Any> {
-        val size = input.readInt()
-        val map = HashMap<String, Any>(size)
-
-        repeat(size) {
-            val name = IOUtil.readString(input)!!
-            val kind = Kind.values()[input.readByte().toInt()]
-
-            val value: Any = when (kind) {
+    override fun read(input: DataInput): Any {
+        return when (Kind.values()[input.readByte().toInt()]) {
             Kind.INT -> input.readInt()
             Kind.FLOAT -> input.readFloat()
             Kind.LONG -> input.readLong()
             Kind.DOUBLE -> input.readDouble()
-                Kind.STRING -> IOUtil.readString(input)!!
+            Kind.STRING -> input.readString()
         }
-
-            map[name] = value
-        }
-
-        return map
     }
 
     private enum class Kind {
@@ -190,11 +173,18 @@ object ConstantsMapExternalizer : DataExternalizer<Map<String, Any>> {
 }
 
 object IntExternalizer : DataExternalizer<Int> {
+    override fun save(output: DataOutput, value: Int) = output.writeInt(value)
     override fun read(input: DataInput): Int = input.readInt()
-
-    override fun save(output: DataOutput, value: Int) {
-        output.writeInt(value)
 }
+
+object LongExternalizer : DataExternalizer<Long> {
+    override fun save(output: DataOutput, value: Long) = output.writeLong(value)
+    override fun read(input: DataInput): Long = input.readLong()
+}
+
+object StringExternalizer : DataExternalizer<String> {
+    override fun save(output: DataOutput, value: String) = IOUtil.writeString(value, output)
+    override fun read(input: DataInput): String = IOUtil.readString(input)
 }
 
 
@@ -244,3 +234,69 @@ open class CollectionExternalizer<T>(
 object StringCollectionExternalizer : CollectionExternalizer<String>(EnumeratorStringDescriptor(), { HashSet() })
 
 object IntCollectionExternalizer : CollectionExternalizer<Int>(IntExternalizer, { HashSet() })
+
+fun DataOutput.writeString(value: String) = StringExternalizer.save(this, value)
+
+fun DataInput.readString(): String = StringExternalizer.read(this)
+
+class ListExternalizer<T>(
+    private val elementExternalizer: DataExternalizer<T>
+) : DataExternalizer<List<T>> {
+
+    override fun save(output: DataOutput, value: List<T>) {
+        output.writeInt(value.size)
+        value.forEach {
+            elementExternalizer.save(output, it)
+        }
+    }
+
+    override fun read(input: DataInput): List<T> {
+        val size = input.readInt()
+        val list = ArrayList<T>(size)
+        repeat(size) {
+            list.add(elementExternalizer.read(input))
+        }
+        return list
+    }
+}
+
+class LinkedHashMapExternalizer<K, V>(
+    private val keyExternalizer: DataExternalizer<K>,
+    private val valueExternalizer: DataExternalizer<V>
+) : DataExternalizer<LinkedHashMap<K, V>> {
+
+    override fun save(output: DataOutput, map: LinkedHashMap<K, V>) {
+        output.writeInt(map.size)
+        for ((key, value) in map) {
+            keyExternalizer.save(output, key)
+            valueExternalizer.save(output, value)
+        }
+    }
+
+    override fun read(input: DataInput): LinkedHashMap<K, V> {
+        val size = input.readInt()
+        val map = LinkedHashMap<K, V>(size)
+        repeat(size) {
+            val key = keyExternalizer.read(input)
+            val value = valueExternalizer.read(input)
+            map[key] = value
+        }
+        return map
+    }
+}
+
+class NullableValueExternalizer<T>(private val valueExternalizer: DataExternalizer<T>) : DataExternalizer<T> {
+
+    override fun save(output: DataOutput, value: T?) {
+        output.writeBoolean(value != null)
+        value?.let {
+            valueExternalizer.save(output, it)
+        }
+    }
+
+    override fun read(input: DataInput): T? {
+        return if (input.readBoolean()) {
+            valueExternalizer.read(input)
+        } else null
+    }
+}

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/incremental/ClasspathSnapshot.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.gradle.incremental
 
 import org.jetbrains.kotlin.incremental.KotlinClassInfo
-import java.io.*
 
 /** Snapshot of a classpath. It consists of a list of [ClasspathEntrySnapshot]s. */
 class ClasspathSnapshot(val classpathEntrySnapshots: List<ClasspathEntrySnapshot>)
@@ -19,64 +18,20 @@ class ClasspathEntrySnapshot(
      * jar).
      */
     val classSnapshots: LinkedHashMap<String, ClassSnapshot>
-) : Serializable {
-
-    companion object {
-        private const val serialVersionUID = 0L
-    }
-}
+)
 
 /**
- * Snapshot of a class. It contains information to compute the source files that need to be recompiled during an incremental run of the
- * `KotlinCompile` task.
+ * Snapshot of a class. It contains minimal information about a class to compute the source files that need to be recompiled during an
+ * incremental run of the `KotlinCompile` task.
+ *
+ * It's important that this class contain only the minimal required information, as it will be part of the classpath snapshot of the
+ * `KotlinCompile` task and the task needs to support compile avoidance. For example, this class should contain public method signatures,
+ * and should not contain private method signatures, or method implementations.
  */
-abstract class ClassSnapshot : Serializable {
+sealed class ClassSnapshot
 
-    companion object {
-        private const val serialVersionUID = 0L
-    }
-}
+/** [ClassSnapshot] of a Kotlin class. */
+class KotlinClassSnapshot(val classInfo: KotlinClassInfo) : ClassSnapshot()
 
-/** [ClassSnapshot] of a kotlinc-generated class. */
-class KotlinClassSnapshot(val classInfo: KotlinClassInfo) : ClassSnapshot(), Serializable {
-
-    companion object {
-        private const val serialVersionUID = 0L
-    }
-}
-
-/** [ClassSnapshot] of a non-kotlinc-generated class. */
-class JavaClassSnapshot : ClassSnapshot(), Serializable {
-
-    // TODO WORK-IN-PROGRESS
-
-    companion object {
-        private const val serialVersionUID = 0L
-    }
-}
-
-/** Utility to read/write a [ClasspathSnapshot] from/to a file. */
-object ClasspathSnapshotSerializer {
-
-    fun readFromFiles(classpathEntrySnapshotFiles: List<File>): ClasspathSnapshot {
-        return ClasspathSnapshot(classpathEntrySnapshotFiles.map { ClasspathEntrySnapshotSerializer.readFromFile(it) })
-    }
-}
-
-/** Utility to read/write a [ClasspathEntrySnapshot] from/to a file. */
-object ClasspathEntrySnapshotSerializer {
-
-    fun readFromFile(classpathEntrySnapshotFile: File): ClasspathEntrySnapshot {
-        check(classpathEntrySnapshotFile.isFile) { "`${classpathEntrySnapshotFile.path}` does not exist (or is a directory)." }
-        return ObjectInputStream(FileInputStream(classpathEntrySnapshotFile).buffered()).use {
-            it.readObject() as ClasspathEntrySnapshot
-        }
-    }
-
-    fun writeToFile(classpathEntrySnapshot: ClasspathEntrySnapshot, classpathEntrySnapshotFile: File) {
-        check(classpathEntrySnapshotFile.parentFile.exists()) { "Parent dir of `${classpathEntrySnapshotFile.path}` does not exist." }
-        ObjectOutputStream(FileOutputStream(classpathEntrySnapshotFile).buffered()).use {
-            it.writeObject(classpathEntrySnapshot)
-        }
-    }
-}
+/** [ClassSnapshot] of a Java class. */
+object JavaClassSnapshot : ClassSnapshot()

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/incremental/ClasspathSnapshotSerializer.kt

@@ -0,0 +1,171 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.incremental
+
+import com.intellij.util.io.DataExternalizer
+import org.jetbrains.kotlin.incremental.KotlinClassInfo
+import org.jetbrains.kotlin.incremental.storage.*
+import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import java.io.*
+
+/** Utility to serialize a [ClasspathSnapshot]. */
+object ClasspathSnapshotSerializer {
+
+    fun load(classpathEntrySnapshotFiles: List<File>): ClasspathSnapshot {
+        return ClasspathSnapshot(classpathEntrySnapshotFiles.map {
+            ClasspathEntrySnapshotSerializer.load(it)
+        })
+    }
+}
+
+object ClasspathEntrySnapshotSerializer : DataSerializer<ClasspathEntrySnapshot> {
+
+    override fun save(output: DataOutput, snapshot: ClasspathEntrySnapshot) {
+        LinkedHashMapExternalizer(StringExternalizer, ClassSnapshotDataSerializer).save(output, snapshot.classSnapshots)
+    }
+
+    override fun read(input: DataInput): ClasspathEntrySnapshot {
+        return ClasspathEntrySnapshot(
+            classSnapshots = LinkedHashMapExternalizer(StringExternalizer, ClassSnapshotDataSerializer).read(input)
+        )
+    }
+}
+
+object ClassSnapshotDataSerializer : DataSerializer<ClassSnapshot> {
+
+    override fun save(output: DataOutput, snapshot: ClassSnapshot) {
+        output.writeString(snapshot.javaClass.name)
+        when (snapshot) {
+            is KotlinClassSnapshot -> KotlinClassSnapshotExternalizer.save(output, snapshot)
+            is JavaClassSnapshot -> JavaClassSnapshotExternalizer.save(output, snapshot)
+        }
+    }
+
+    override fun read(input: DataInput): ClassSnapshot {
+        return when (val className = input.readString()) {
+            KotlinClassSnapshot::class.java.name -> KotlinClassSnapshotExternalizer.read(input)
+            JavaClassSnapshot::class.java.name -> JavaClassSnapshotExternalizer.read(input)
+            else -> error("Unrecognized class: $className")
+        }
+    }
+}
+
+object KotlinClassSnapshotExternalizer : DataExternalizer<KotlinClassSnapshot> {
+
+    override fun save(output: DataOutput, snapshot: KotlinClassSnapshot) {
+        KotlinClassInfoExternalizer.save(output, snapshot.classInfo)
+    }
+
+    override fun read(input: DataInput): KotlinClassSnapshot {
+        return KotlinClassSnapshot(classInfo = KotlinClassInfoExternalizer.read(input))
+    }
+}
+
+object KotlinClassInfoExternalizer : DataExternalizer<KotlinClassInfo> {
+
+    override fun save(output: DataOutput, info: KotlinClassInfo) {
+        ClassIdExternalizer.save(output, info.classId)
+        output.writeInt(info.classKind.id)
+        ListExternalizer(StringExternalizer).save(output, info.classHeaderData.toList())
+        ListExternalizer(StringExternalizer).save(output, info.classHeaderStrings.toList())
+        NullableValueExternalizer(StringExternalizer).save(output, info.multifileClassName)
+        LinkedHashMapExternalizer(StringExternalizer, ConstantExternalizer).save(output, info.constantsMap)
+        LinkedHashMapExternalizer(StringExternalizer, LongExternalizer).save(output, info.inlineFunctionsMap)
+    }
+
+    override fun read(input: DataInput): KotlinClassInfo {
+        return KotlinClassInfo(
+            classId = ClassIdExternalizer.read(input),
+            classKind = KotlinClassHeader.Kind.getById(input.readInt()),
+            classHeaderData = ListExternalizer(StringExternalizer).read(input).toTypedArray(),
+            classHeaderStrings = ListExternalizer(StringExternalizer).read(input).toTypedArray(),
+            multifileClassName = NullableValueExternalizer(StringExternalizer).read(input),
+            constantsMap = LinkedHashMapExternalizer(StringExternalizer, ConstantExternalizer).read(input),
+            inlineFunctionsMap = LinkedHashMapExternalizer(StringExternalizer, LongExternalizer).read(input)
+        )
+    }
+}
+
+object ClassIdExternalizer : DataExternalizer<ClassId> {
+
+    override fun save(output: DataOutput, classId: ClassId) {
+        FqNameExternalizer.save(output, classId.packageFqName)
+        FqNameExternalizer.save(output, classId.relativeClassName)
+        output.writeBoolean(classId.isLocal)
+    }
+
+    override fun read(input: DataInput): ClassId {
+        return ClassId(
+            /* packageFqName */ FqNameExternalizer.read(input),
+            /* relativeClassName */ FqNameExternalizer.read(input),
+            /* isLocal */ input.readBoolean()
+        )
+    }
+}
+
+object FqNameExternalizer : DataExternalizer<FqName> {
+
+    override fun save(output: DataOutput, fqName: FqName) {
+        output.writeString(fqName.asString())
+    }
+
+    override fun read(input: DataInput): FqName {
+        return FqName(input.readString())
+    }
+}
+
+object JavaClassSnapshotExternalizer : DataExternalizer<JavaClassSnapshot> {
+
+    override fun save(output: DataOutput, snapshot: JavaClassSnapshot) {
+    }
+
+    override fun read(input: DataInput): JavaClassSnapshot {
+        return JavaClassSnapshot
+    }
+}
+
+interface DataSerializer<T> : DataExternalizer<T> {
+
+    fun save(file: File, value: T) {
+        return FileOutputStream(file).buffered().use {
+            it.writeValue(value)
+        }
+    }
+
+    fun load(file: File): T {
+        return FileInputStream(file).buffered().use {
+            it.readValue()
+        }
+    }
+
+    fun toByteArray(value: T): ByteArray {
+        val byteArrayOutputStream = ByteArrayOutputStream()
+        byteArrayOutputStream.buffered().use {
+            it.writeValue(value)
+        }
+        return byteArrayOutputStream.toByteArray()
+    }
+
+    fun fromByteArray(byteArray: ByteArray): T {
+        return ByteArrayInputStream(byteArray).buffered().use {
+            it.readValue()
+        }
+    }
+
+    private fun OutputStream.writeValue(value: T) {
+        DataOutputStream(this).use {
+            save(it, value)
+        }
+    }
+
+    private fun InputStream.readValue(): T {
+        return DataInputStream(this).use {
+            read(it)
+        }
+    }
+}

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/incremental/ClasspathSnapshotter.kt

@@ -32,9 +32,8 @@ object ClasspathEntrySnapshotter {
 object ClassSnapshotter {
 
     fun snapshot(classContents: ByteArray): ClassSnapshot {
-        // TODO Add custom serialization to KotlinClassInfo then return it here. For now, use JavaClassSnapshot.
-        KotlinClassInfo.tryCreateFrom(classContents)?.let { KotlinClassSnapshot(it) }
-        return JavaClassSnapshot()
+        return KotlinClassInfo.tryCreateFrom(classContents)?.let { KotlinClassSnapshot(it) }
+            ?: JavaClassSnapshot
     }
 }
 

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/transforms/ClasspathEntrySnapshotTransform.kt

@@ -9,10 +9,10 @@ import org.gradle.api.artifacts.transform.*
 import org.gradle.api.file.FileSystemLocation
 import org.gradle.api.provider.Provider
 import org.gradle.api.tasks.Classpath
-import org.jetbrains.kotlin.gradle.incremental.ClasspathEntrySnapshotter
 import org.jetbrains.kotlin.gradle.incremental.ClasspathEntrySnapshotSerializer
+import org.jetbrains.kotlin.gradle.incremental.ClasspathEntrySnapshotter
 
-/** Transform to create a snapshot ([CLASSPATH_ENTRY_SNAPSHOT_ARTIFACT_TYPE]) of a classpath entry (directory or jar). */
+/** Transform to create a snapshot of a classpath entry (directory or jar). */
 @CacheableTransform
 abstract class ClasspathEntrySnapshotTransform : TransformAction<TransformParameters.None> {
 
@@ -25,7 +25,7 @@ abstract class ClasspathEntrySnapshotTransform : TransformAction<TransformParame
         val snapshotFile = outputs.file(CLASSPATH_ENTRY_SNAPSHOT_FILE_NAME)
 
         val snapshot = ClasspathEntrySnapshotter.snapshot(classpathEntry)
-        ClasspathEntrySnapshotSerializer.writeToFile(snapshot, snapshotFile)
+        ClasspathEntrySnapshotSerializer.save(snapshotFile, snapshot)
     }
 }
 

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/Tasks.kt

@@ -778,8 +778,8 @@ abstract class KotlinCompile @Inject constructor(
         val currentSnapshotFiles = classpathSnapshotProperties.classpathSnapshot.files.toList()
         val previousSnapshotFiles = getClasspathSnapshotFilesInDir(classpathSnapshotProperties.classpathSnapshotDir.get().asFile)
 
-        val currentSnapshot = ClasspathSnapshotSerializer.readFromFiles(currentSnapshotFiles)
-        val previousSnapshot = ClasspathSnapshotSerializer.readFromFiles(previousSnapshotFiles)
+        val currentSnapshot = ClasspathSnapshotSerializer.load(currentSnapshotFiles)
+        val previousSnapshot = ClasspathSnapshotSerializer.load(previousSnapshotFiles)
 
         return ClasspathChangesComputer.getChanges(currentSnapshot, previousSnapshot)
     }

libraries/tools/kotlin-gradle-plugin/src/test/kotlin/org/jetbrains/kotlin/gradle/incremental/ClasspathSnapshotSerializerTest.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.incremental
+
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+abstract class ClasspathSnapshotSerializerTest : ClasspathSnapshotTestCommon() {
+
+    protected abstract val testSourceFile: ChangeableTestSourceFile
+
+    @Test
+    fun `test ClassSnapshotDataSerializer`() {
+        val originalSnapshot = testSourceFile.compileAndSnapshot()
+        val serializedSnapshot = ClassSnapshotDataSerializer.toByteArray(originalSnapshot)
+        val deserializedSnapshot = ClassSnapshotDataSerializer.fromByteArray(serializedSnapshot)
+
+        assertEquals(originalSnapshot.toGson(), deserializedSnapshot.toGson())
+    }
+}
+
+class KotlinClassesClasspathSnapshotSerializerTest : ClasspathSnapshotSerializerTest() {
+    override val testSourceFile = SimpleKotlinClass(tmpDir)
+}

libraries/tools/kotlin-gradle-plugin/src/test/kotlin/org/jetbrains/kotlin/gradle/incremental/ClasspathSnapshotTestCommon.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.gradle.incremental
 
 import com.google.gson.GsonBuilder
-import org.jetbrains.kotlin.incremental.KotlinClassInfo
 import org.junit.Rule
 import org.junit.rules.TemporaryFolder
 import java.io.File
@@ -35,7 +34,7 @@ abstract class ClasspathSnapshotTestCommon {
             check(relativePath.endsWith(".class"))
         }
 
-        fun snapshot() = KotlinClassSnapshot(KotlinClassInfo.tryCreateFrom(asFile().readBytes())!!)
+        fun snapshot() = ClassSnapshotter.snapshot(asFile().readBytes())
     }
 
     open class TestSourceFile(val sourceFile: SourceFile, val tmpDir: TemporaryFolder) {