From 207f510ccaeefb09d68c15dfe513a1127d447c93 Mon Sep 17 00:00:00 2001
From: "sebastian.sellmair" <sebastian.sellmair@jetbrains.com>
Date: Fri, 18 Feb 2022 16:31:00 +0100
Subject: [PATCH] [KPM] IdeaKotlinProjectModelObjectGraphTest: Implement
 WriteReplacedModel annotation

KT-51386
---
 .../kpm/KotlinExternalModelContainerImpl.kt   |  5 ++-
 .../gradle/kpm/idea/WriteReplacedModel.kt     | 10 ++++++
 .../IdeaKotlinProjectModelObjectGraphTest.kt  | 32 +++++++++++--------
 3 files changed, 33 insertions(+), 14 deletions(-)
 create mode 100644 libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/WriteReplacedModel.kt

diff --git a/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/KotlinExternalModelContainerImpl.kt b/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/KotlinExternalModelContainerImpl.kt
index 524b2a4fe6d..711f24be37f 100644
--- a/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/KotlinExternalModelContainerImpl.kt
+++ b/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/KotlinExternalModelContainerImpl.kt
@@ -5,11 +5,13 @@
 
 package org.jetbrains.kotlin.gradle.kpm
 
+import org.jetbrains.kotlin.gradle.kpm.idea.WriteReplacedModel
 import java.io.Serializable
 
+@WriteReplacedModel(SerializedKotlinExternalModelContainerCarrier::class)
 internal class KotlinMutableExternalModelContainerImpl private constructor(
     private val values: MutableMap<KotlinExternalModelKey<*>, Any>
-) : KotlinMutableExternalModelContainer(), Serializable {
+) : KotlinMutableExternalModelContainer() {
 
     constructor() : this(mutableMapOf())
 
@@ -42,6 +44,7 @@ internal class KotlinMutableExternalModelContainerImpl private constructor(
     }
 }
 
+@WriteReplacedModel(SerializedKotlinExternalModelContainerCarrier::class)
 private class SerializedKotlinExternalModelContainer(
     private val serializedValues: MutableMap<KotlinExternalModelId<*>, ByteArray>
 ) : KotlinExternalModelContainer(), Serializable {
diff --git a/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/WriteReplacedModel.kt b/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/WriteReplacedModel.kt
new file mode 100644
index 00000000000..0b19c83bdb8
--- /dev/null
+++ b/libraries/tools/kotlin-gradle-plugin-idea/src/main/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/WriteReplacedModel.kt
@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.gradle.kpm.idea
+
+import kotlin.reflect.KClass
+
+internal annotation class WriteReplacedModel(val replacedBy: KClass<*>)
diff --git a/libraries/tools/kotlin-gradle-plugin-idea/src/test/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/IdeaKotlinProjectModelObjectGraphTest.kt b/libraries/tools/kotlin-gradle-plugin-idea/src/test/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/IdeaKotlinProjectModelObjectGraphTest.kt
index 9cc277b409c..1993d755680 100644
--- a/libraries/tools/kotlin-gradle-plugin-idea/src/test/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/IdeaKotlinProjectModelObjectGraphTest.kt
+++ b/libraries/tools/kotlin-gradle-plugin-idea/src/test/kotlin/org/jetbrains/kotlin/gradle/kpm/idea/IdeaKotlinProjectModelObjectGraphTest.kt
@@ -11,8 +11,8 @@ import org.reflections.Reflections
 import java.lang.reflect.Field
 import java.lang.reflect.Modifier
 import kotlin.reflect.KClass
-import kotlin.reflect.KVisibility
 import kotlin.reflect.KVisibility.PUBLIC
+import kotlin.reflect.full.findAnnotation
 import kotlin.reflect.full.memberProperties
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -80,20 +80,26 @@ class IdeaKotlinProjectModelObjectGraphTest(private val node: KClass<*>, @Suppre
         @JvmStatic
         @Parameterized.Parameters(name = "{1}")
         fun findNodes(): List<Array<Any>> {
-            val classes = mutableSetOf<KClass<*>>(IdeaKotlinProjectModel::class)
-            val resolveQueue = ArrayDeque(classes)
+            val classes = mutableSetOf<KClass<*>>()
+            val resolveQueue = ArrayDeque<KClass<*>>(listOf(IdeaKotlinProjectModel::class))
 
             while (resolveQueue.isNotEmpty()) {
-                val children = resolveQueue.removeFirst().resolveReachableClasses()
-                children.forEach { child ->
-                    if (classes.add(child)) {
-                        resolveQueue.add(child)
-                        if (child.java.isInterface || Modifier.isAbstract(child.java.modifiers)) {
-                            val subtypes = reflections.getSubTypesOf(child.java).map { it.kotlin }
-                            assertTrue(subtypes.isNotEmpty(), "Missing implementations for $child")
-                            classes.addAll(subtypes)
-                            resolveQueue.addAll(subtypes)
-                        }
+                val next = resolveQueue.removeFirst()
+
+                /* Model gets replaced by other class */
+                val writeReplacedModelAnnotation = next.findAnnotation<WriteReplacedModel>()
+                if (writeReplacedModelAnnotation != null) {
+                    resolveQueue.add(writeReplacedModelAnnotation.replacedBy)
+                    continue
+                }
+
+                if (!classes.add(next)) continue
+                next.resolveReachableClasses().forEach { child ->
+                    resolveQueue.add(child)
+                    if (child.java.isInterface || Modifier.isAbstract(child.java.modifiers)) {
+                        val subtypes = reflections.getSubTypesOf(child.java).map { it.kotlin }
+                        assertTrue(subtypes.isNotEmpty(), "Missing implementations for $child")
+                        resolveQueue.addAll(subtypes)
                     }
                 }
             }