Repository: yandex/yatagan (kotlin)

Issue:
[KSP2] HeavyTest failures because of wildcards
In code of generated component there are a few places where instead of
```kotlin
  @Override
  public Lazy<Node_122> getGet_24() {
    return new CachingProviderImpl(this, 23);
  }
```
generates this:
```kotlin
  @Override
  public Lazy<? extends Node_122> getGet_24() {
    return new CachingProviderImpl(this, 23);
  }
```

Reproduces only with KSP2 even with latest version at the moment of issue creation (`2.1.20-2.0.0`).

Pull request description:
Fixes #184

Patch:
diff --git a/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspLexicalScope.kt b/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspLexicalScope.kt
index 6926228e..18ad381c 100644
--- a/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspLexicalScope.kt
+++ b/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspLexicalScope.kt
@@ -1,6 +1,7 @@
 package com.yandex.yatagan.lang.ksp
 
 import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.yandex.yatagan.lang.LangModelFactory
 import com.yandex.yatagan.lang.TypeDeclaration
@@ -15,10 +16,14 @@ import com.yandex.yatagan.lang.scope.CachingMetaFactory
  */
 class KspLexicalScope(
     resolver: Resolver,
+    environment: SymbolProcessorEnvironment,
 ) : LexicalScopeBase() {
     init {
         ext[CachingMetaFactory] = CachingFactorySimple
-        ext[ProcessingUtils] = ProcessingUtils(resolver)
+        ext[ProcessingUtils] = ProcessingUtils(
+            resolver = resolver,
+            isKsp2 = environment.kspVersion >= KotlinVersion(2, 0),
+        )
         ext[LangModelFactory] = KspModelFactoryImpl(this)
     }
 
diff --git a/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspTypeImpl.kt b/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspTypeImpl.kt
index 4f50f7e4..20e5fe68 100644
--- a/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspTypeImpl.kt
+++ b/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/KspTypeImpl.kt
@@ -20,6 +20,7 @@ import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSType
 import com.google.devtools.ksp.symbol.KSTypeParameter
 import com.google.devtools.ksp.symbol.KSTypeReference
+import com.google.devtools.ksp.symbol.Variance
 import com.yandex.yatagan.lang.Type
 import com.yandex.yatagan.lang.TypeDeclaration
 import com.yandex.yatagan.lang.common.NoDeclaration
@@ -34,9 +35,9 @@ import kotlin.LazyThreadSafetyMode.PUBLICATION
 
 internal class KspTypeImpl private constructor(
     lexicalScope: LexicalScope,
-    resolvedTypeInfo: Pair<JvmTypeInfo, KSType>,
+    resolvedTypeInfo: Pair<JvmTypeInfo, KSTypeEquivalence>,
 ) : CtTypeBase(), LexicalScope by lexicalScope {
-    val impl: KSType = resolvedTypeInfo.second
+    val impl: KSType = resolvedTypeInfo.second.type
     private val jvmType: JvmTypeInfo = resolvedTypeInfo.first
 
     override val nameModel: CtTypeNameModel by lazy {
@@ -93,7 +94,44 @@ internal class KspTypeImpl private constructor(
     )
 
     companion object Factory : FactoryKey<ResolveTypeInfo, Type> {
-        private object Caching : FactoryKey<Pair<JvmTypeInfo, KSType>, KspTypeImpl> {
+        /**
+         * A polymorphic wrapper over KSType required to correctly cache the KspTypeImpl instances.
+         */
+        private sealed interface KSTypeEquivalence {
+            val type: KSType
+
+            /**
+             * Use KSType directly - its equals/hashcode suit us.
+             */
+            @JvmInline value class Ksp1(override val type: KSType) : KSTypeEquivalence
+
+            /**
+             * KSP2 uses semantic type equivalence, which ignores "wildcards" - redundant projections.
+             * Given `interface Lazy<out T> {}` and `interface Foo`,
+             * types `Lazy<Foo>` and `Lazy<out Foo>` are semantically equivalent,
+             * but we need to distinguish between them.
+             *
+             * So we capture type projections and compare them additionally.
+             */
+            class Ksp2(
+                override val type: KSType,
+            ) : KSTypeEquivalence {
+                private val flatProjections: List<Variance> = buildList { flatten(type) }
+                private val hashCode by lazy { type.hashCode() + 31 * flatProjections.hashCode() }
+                private fun MutableList<Variance>.flatten(type: KSType?) {
+                    type?.arguments?.forEach {
+                        add(it.variance)
+                        flatten(it.type?.resolve())
+                    }
+                }
+
+                override fun hashCode(): Int = hashCode
+                override fun equals(other: Any?) =
+                    other === this || other is Ksp2 && type == other.type && flatProjections == other.flatProjections
+            }
+        }
+
+        private object Caching : FactoryKey<Pair<JvmTypeInfo, KSTypeEquivalence>, KspTypeImpl> {
             override fun LexicalScope.factory() = caching(createWithScope = ::KspTypeImpl)
         }
 
@@ -121,7 +159,11 @@ internal class KspTypeImpl private constructor(
                         nameModel = InvalidNameModel.TypeVariable(typeVar = type.declaration.simpleName.asString())
                     )
                 }
-                else -> Caching(JvmTypeInfo(jvmSignature = jvmSignatureHint, type = type) to type)
+                else -> {
+                    val equivalenceWrapped =
+                        if (Utils.isKsp2) KSTypeEquivalence.Ksp2(type) else KSTypeEquivalence.Ksp1(type)
+                    Caching(JvmTypeInfo(jvmSignature = jvmSignatureHint, type = type) to equivalenceWrapped)
+                }
             }
         }
     }
diff --git a/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/ProcessingUtils.kt b/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/ProcessingUtils.kt
index 1155429d..a9927284 100644
--- a/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/ProcessingUtils.kt
+++ b/lang/ksp/src/main/kotlin/com/yandex/yatagan/lang/ksp/ProcessingUtils.kt
@@ -24,6 +24,7 @@ import com.yandex.yatagan.lang.scope.LexicalScope
 
 internal class ProcessingUtils(
     val resolver: Resolver,
+    val isKsp2: Boolean,
 ) {
     val classType by lazy {
         resolver.getClassDeclarationByName("java.lang.Class")!!
diff --git a/processor/ksp/src/main/kotlin/com/yandex/yatagan/processor/ksp/KspYataganProcessor.kt b/processor/ksp/src/main/kotlin/com/yandex/yatagan/processor/ksp/KspYataganProcessor.kt
index d64b3037..05285e0b 100644
--- a/processor/ksp/src/main/kotlin/com/yandex/yatagan/processor/ksp/KspYataganProcessor.kt
+++ b/processor/ksp/src/main/kotlin/com/yandex/yatagan/processor/ksp/KspYataganProcessor.kt
@@ -42,7 +42,7 @@ internal class KspYataganProcessor(
     override lateinit var lexicalScope: KspLexicalScope
 
     override fun process(resolver: Resolver): List<KSAnnotated> {
-        lexicalScope = KspLexicalScope(resolver)
+        lexicalScope = KspLexicalScope(resolver, environment)
         initScopedOptions(lexicalScope, this)
         process(
             sources = sequenceOf(


Test patch:
diff --git a/lang/api/src/test/kotlin/com/yandex/yatagan/lang/LangTestDriver.kt b/lang/api/src/test/kotlin/com/yandex/yatagan/lang/LangTestDriver.kt
index 5c27d9cc..d4aaf90b 100644
--- a/lang/api/src/test/kotlin/com/yandex/yatagan/lang/LangTestDriver.kt
+++ b/lang/api/src/test/kotlin/com/yandex/yatagan/lang/LangTestDriver.kt
@@ -20,6 +20,7 @@ import androidx.room.compiler.processing.util.compiler.TestCompilationArguments
 import androidx.room.compiler.processing.util.compiler.compile
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.SymbolProcessor
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
 import com.google.devtools.ksp.processing.SymbolProcessorProvider
 import com.google.devtools.ksp.symbol.KSAnnotated
 import com.yandex.yatagan.base.mapToArray
@@ -146,14 +147,15 @@ interface LangTestDriver : SourceSet {
                 compilation: TestCompilationArguments,
                 block: InspectBlock,
             ) = compilation.copy(
-                symbolProcessorProviders = listOf(SymbolProcessorProvider { Inspector(block) })
+                symbolProcessorProviders = listOf(SymbolProcessorProvider { Inspector(block, it) })
             )
 
             private class Inspector(
                 private val inspectBlock: InspectBlock,
+                private val environment: SymbolProcessorEnvironment,
             ) : SymbolProcessor {
                 override fun process(resolver: Resolver): List<KSAnnotated> {
-                    val lexicalScope = KspLexicalScope(resolver)
+                    val lexicalScope = KspLexicalScope(resolver, environment)
                     inspectBlock(lexicalScope.ext.langFactory)
                     return emptyList()
                 }
diff --git a/testing/tests/src/test/kotlin/com/yandex/yatagan/testing/tests/HeavyTest.kt b/testing/tests/src/test/kotlin/com/yandex/yatagan/testing/tests/HeavyTest.kt
index de97d984..020d0a04 100644
--- a/testing/tests/src/test/kotlin/com/yandex/yatagan/testing/tests/HeavyTest.kt
+++ b/testing/tests/src/test/kotlin/com/yandex/yatagan/testing/tests/HeavyTest.kt
@@ -20,7 +20,6 @@ import com.yandex.yatagan.testing.procedural.Distribution
 import com.yandex.yatagan.testing.procedural.ExperimentalGenerationApi
 import com.yandex.yatagan.testing.procedural.GenerationParams
 import com.yandex.yatagan.testing.procedural.generate
-import org.junit.Assume.assumeFalse
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
@@ -64,8 +63,6 @@ class HeavyTest(
 
     @Test
     fun `procedural - Tall & Thin`() {
-        // TODO(#184): fix and enable test back
-        assumeFalse(backendUnderTest == Backend.Ksp2)
         generate(
             params = baseParams,
             testMethodName = "test",
@@ -77,8 +74,6 @@ class HeavyTest(
 
     @Test
     fun `procedural - Fat & Flat`() {
-        // TODO(#184): fix and enable test back
-        assumeFalse(backendUnderTest == Backend.Ksp2)
         generate(
             params = baseParams.copy(
                 componentTreeMaxDepth = 1,