[KPM] KotlinExternalModelContainerImpl: Use carrier class for serialization

Without this carrier the `SerializedKotlinExternalModelContainer`
itself would fail serialization if values got accessed.

KT-51262
KT-51220

libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/KotlinExternalModelContainerImpl.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.gradle.kpm
 
-import org.jetbrains.kotlin.gradle.kpm.idea.InternalKotlinGradlePluginApi
 import java.io.Serializable
 
 internal class KotlinMutableExternalModelContainerImpl : KotlinMutableExternalModelContainer(), Serializable {
@@ -32,15 +31,11 @@ internal class KotlinMutableExternalModelContainerImpl : KotlinMutableExternalMo
 
     @Synchronized
     private fun writeReplace(): Any {
-        val serializedValues = values.filterKeys { it.serializer != null }
-            .mapValues { (key, value) ->
-                @Suppress("unchecked_cast")
-                val serializer = checkNotNull(key.serializer) as KotlinExternalModelSerializer<Any>
-                serializer.serialize(value)
-            }.mapKeys { (key, _) -> key.id }
-            .toMutableMap()
+        return SerializedKotlinExternalModelContainerCarrier(serialize(values))
+    }
 
-        return SerializedKotlinExternalModelContainer(serializedValues)
+    companion object {
+        private const val serialVersionUID = 0L
     }
 }
 
@@ -48,28 +43,54 @@ private class SerializedKotlinExternalModelContainer(
     private val serializedValues: MutableMap<KotlinExternalModelId<*>, ByteArray>
 ) : KotlinExternalModelContainer(), Serializable {
 
-    private val deserializedValues = mutableMapOf<KotlinExternalModelId<*>, Any>()
+    private val deserializedValues = mutableMapOf<KotlinExternalModelKey<*>, Any>()
 
     override val ids: Set<KotlinExternalModelId<*>>
-        @Synchronized get() = serializedValues.keys + deserializedValues.keys
+        @Synchronized get() = serializedValues.keys + deserializedValues.keys.map { it.id }
 
     @Synchronized
     override fun <T : Any> contains(key: KotlinExternalModelKey<T>): Boolean {
-        return key.id in deserializedValues || key.id in serializedValues
+        return key.id in serializedValues || key in deserializedValues
     }
 
     @Synchronized
     @Suppress("unchecked_cast")
     override fun <T : Any> get(key: KotlinExternalModelKey<T>): T? {
-        deserializedValues[key.id]?.let { return it as T }
+        deserializedValues[key]?.let { return it as T }
         val serializedValue = serializedValues.remove(key.id) ?: return null
         val deserializedValue = key.serializer?.deserialize(serializedValue) ?: return null
-        deserializedValues[key.id] = deserializedValue
+        deserializedValues[key] = deserializedValue
         return deserializedValue
     }
 
-    @InternalKotlinGradlePluginApi
+    @Synchronized
+    private fun writeReplace(): Any {
+        return SerializedKotlinExternalModelContainerCarrier(serializedValues + serialize(deserializedValues))
+    }
+
     companion object {
         private const val serialVersionUID = 0L
     }
 }
+
+private class SerializedKotlinExternalModelContainerCarrier(
+    private val serializedValues: Map<KotlinExternalModelId<*>, ByteArray>
+) : Serializable {
+
+    private fun readResolve(): Any {
+        return SerializedKotlinExternalModelContainer(serializedValues.toMutableMap())
+    }
+
+    companion object {
+        private const val serialVersionUID = 0L
+    }
+}
+
+private fun serialize(values: Map<KotlinExternalModelKey<*>, Any>): Map<KotlinExternalModelId<*>, ByteArray> {
+    return values.filterKeys { it.serializer != null }
+        .mapValues { (key, value) ->
+            @Suppress("unchecked_cast")
+            val serializer = checkNotNull(key.serializer) as KotlinExternalModelSerializer<Any>
+            serializer.serialize(value)
+        }.mapKeys { (key, _) -> key.id }
+}

libraries/tools/kotlin-gradle-plugin-idea/src/test/kotlin/SerializationUtil.kt

@@ -0,0 +1,26 @@
+import org.gradle.internal.io.ClassLoaderObjectInputStream
+import java.io.*
+
+/*
+ * Copyright 2010-2022 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.
+ */
+
+internal fun Any.serialize(): ByteArray {
+    return ByteArrayOutputStream().use { byteArrayOutputStream ->
+        ObjectOutputStream(byteArrayOutputStream).writeObject(this)
+        byteArrayOutputStream.toByteArray()
+    }
+}
+
+internal inline fun <reified T : Serializable> ByteArray.deserialize(): T {
+    val inputStream = ByteArrayInputStream(this)
+    val objectInputStream = ObjectInputStream(inputStream)
+    return objectInputStream.use { it.readObject() } as T
+}
+
+internal fun ByteArray.deserialize(classLoader: ClassLoader): Any {
+    val inputStream = ByteArrayInputStream(this)
+    val objectInputStream = ClassLoaderObjectInputStream(inputStream, classLoader)
+    return objectInputStream.use { it.readObject() }
+}

libraries/tools/kotlin-gradle-plugin-idea/src/test/kotlin/org/jetbrains/kotlin/gradle/kpm/KotlinExternalModelContainerTest.kt

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.gradle.kpm
 
+import deserialize
+import serialize
 import java.io.*
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -58,14 +60,7 @@ class KotlinExternalModelContainerTest {
         container[unretainedModelKeyFoo] = UnretainedModel(5)
         container[unretainedModelKeyBar] = UnretainedModel(6)
 
-        val serializedContainer = ByteArrayOutputStream().use { baos ->
-            ObjectOutputStream(baos).use { oos -> oos.writeObject(container) }
-            baos.toByteArray()
-        }
-
-        val deserializedContainer = ObjectInputStream(ByteArrayInputStream(serializedContainer)).use { ois ->
-            ois.readObject() as KotlinExternalModelContainer
-        }
+        val deserializedContainer = container.serialize().deserialize<KotlinExternalModelContainer>()
 
         assertEquals(RetainedModel(1), deserializedContainer[retainedModelKey])
         assertEquals(RetainedModel(2), deserializedContainer[retainedModelKeyFoo])
@@ -74,4 +69,19 @@ class KotlinExternalModelContainerTest {
         assertNull(deserializedContainer[unretainedModelKeyFoo])
         assertNull(deserializedContainer[unretainedModelKeyBar])
     }
+
+    @Test
+    fun `test - serializing container twice`() {
+        val container = KotlinExternalModelContainer.mutable()
+        container[retainedModelKey] = RetainedModel(1)
+        container[unretainedModelKey] = UnretainedModel(4)
+
+        val deserializedContainer = container.serialize().deserialize<KotlinExternalModelContainer>()
+        assertEquals(RetainedModel(1), deserializedContainer[retainedModelKey])
+        assertNull(deserializedContainer[unretainedModelKey])
+
+        val twiceDeserializedContainer = deserializedContainer.serialize().deserialize<KotlinExternalModelContainer>()
+        assertEquals(RetainedModel(1), twiceDeserializedContainer[retainedModelKey])
+        assertNull(twiceDeserializedContainer[unretainedModelKey])
+    }
 }