[KPM] KotlinExternalModelContainerImpl: Fix corrupt container after access without deserializer

KT-51262
KT-51220

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

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.gradle.kpm
 import org.jetbrains.kotlin.gradle.kpm.idea.InternalKotlinGradlePluginApi
 import java.io.Serializable
 
-sealed class KotlinExternalModelContainer() : Serializable {
+sealed class KotlinExternalModelContainer : Serializable {
     abstract val ids: Set<KotlinExternalModelId<*>>
     abstract operator fun <T : Any> contains(key: KotlinExternalModelKey<T>): Boolean
     abstract operator fun <T : Any> get(key: KotlinExternalModelKey<T>): T?

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

@@ -57,9 +57,10 @@ private class SerializedKotlinExternalModelContainer(
     @Suppress("unchecked_cast")
     override fun <T : Any> get(key: KotlinExternalModelKey<T>): T? {
         deserializedValues[key]?.let { return it as T }
-        val serializedValue = serializedValues.remove(key.id) ?: return null
+        val serializedValue = serializedValues[key.id] ?: return null
         val deserializedValue = key.serializer?.deserialize(serializedValue) ?: return null
         deserializedValues[key] = deserializedValue
+        serializedValues.remove(key.id)
         return deserializedValue
     }
 

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

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.gradle.kpm
 
 import deserialize
 import serialize
-import java.io.*
+import java.io.Serializable
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNull
@@ -84,4 +84,26 @@ class KotlinExternalModelContainerTest {
         assertEquals(RetainedModel(1), twiceDeserializedContainer[retainedModelKey])
         assertNull(twiceDeserializedContainer[unretainedModelKey])
     }
+
+    @Test
+    fun `test - accessing deserialized container without serializer`() {
+        val container = KotlinExternalModelContainer.mutable()
+        container[retainedModelKey] = RetainedModel(1)
+        container[retainedModelKeyFoo] = RetainedModel(2)
+
+        val deserializedContainer = container.serialize().deserialize<KotlinExternalModelContainer>()
+        assertEquals(RetainedModel(1), deserializedContainer[retainedModelKey])
+
+        /* Accessing something already deserialized without deserializer: Returning null, because the keys do not match */
+        assertNull(deserializedContainer[KotlinExternalModelKey(retainedModelKey.id)])
+
+        /* Accessing something not deserialized without deserializer: Cannot deserialize: Null to be lenient */
+        assertNull(deserializedContainer[KotlinExternalModelKey(retainedModelKeyFoo.id)])
+
+        /* Accessing something previously requested but now with deserializer */
+        assertEquals(RetainedModel(2), deserializedContainer[retainedModelKeyFoo])
+
+        /* Now accessing retainedModelFoo without serializer should behave like above */
+        assertNull(deserializedContainer[KotlinExternalModelKey(retainedModelKeyFoo.id)])
+    }
 }