[Analysis API] add tests for KtSubstitutor

analysis/analysis-test-framework/build.gradle.kts

@@ -33,3 +33,7 @@ projectTest(jUnitMode = JUnitMode.JUnit5) {
 
 testsJar()
 
+
+tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>().configureEach {
+    kotlinOptions.freeCompilerArgs += "-Xcontext-receivers"
+}

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt

@@ -104,7 +104,6 @@ abstract class AbstractAnalysisApiBasedTest : TestWithDisposable() {
 
         useDirectives(*AbstractKotlinCompilerTest.defaultDirectiveContainers.toTypedArray())
         useDirectives(JvmEnvironmentConfigurationDirectives)
-        useDirectives(SubstitutionParser.Directives)
 
         useSourcePreprocessor(::ExpressionMarkersSourceFilePreprocessor)
         useAdditionalService { ExpressionMarkerProvider() }

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/SubstitutionParser.kt

@@ -5,45 +5,62 @@
 
 package org.jetbrains.kotlin.analysis.test.framework.services
 
+import com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.components.buildSubstitutor
 import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtTypeParameter
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
-import org.jetbrains.kotlin.test.model.TestModule
 
 object SubstitutionParser {
-    fun KtAnalysisSession.parseSubstitutors(module: TestModule, file: KtFile): List<KtSubstitutor> {
-        val directives = module.directives[Directives.SUBSTITUTOR]
-        if (directives.isEmpty()) return emptyList()
-        val allTypeParameterSymbols = file
+    context(KtAnalysisSession)
+    fun parseSubstitutor(declaration: KtCallableDeclaration): KtSubstitutor {
+        val comment = declaration.firstChild as PsiComment
+        return parseSubstitutor(comment, declaration)
+    }
+
+    context(KtAnalysisSession)
+    fun parseSubstitutor(ktFile: KtFile, declaration: KtCallableDeclaration): KtSubstitutor {
+        val comment = ktFile.children.filterIsInstance<PsiComment>().single { it.text.startsWith(SUBSTITUTOR_PREFIX) }
+        return parseSubstitutor(comment, declaration)
+    }
+
+
+    context(KtAnalysisSession)
+    fun parseSubstitutor(comment: PsiComment, scopeForTypeParameters: KtElement): KtSubstitutor {
+        val directivesAsString = comment.text.trim()
+        check(directivesAsString.startsWith(SUBSTITUTOR_PREFIX))
+        val substitutorAsMap = parseSubstitutions(directivesAsString.removePrefix(SUBSTITUTOR_PREFIX))
+
+        val allTypeParameterSymbols = scopeForTypeParameters
             .collectDescendantsOfType<KtTypeParameter>()
             .map { it.getTypeParameterSymbol() }
             .groupBy { it.name.asString() }
             .mapValues { it.value.single() }
 
-        return directives.map { subsitutor ->
-            buildSubstitutor {
-                subsitutor.forEach { (typeParameterName, typeString) ->
-                    val typeParameterSymbol = allTypeParameterSymbols.getValue(typeParameterName)
-                    val type = with(TypeParser) { parseTypeFromString(typeString, file, allTypeParameterSymbols) }
-                    substitution(typeParameterSymbol, type)
-                }
+        return buildSubstitutor {
+            substitutorAsMap.forEach { (typeParameterName, typeString) ->
+                val typeParameterSymbol = allTypeParameterSymbols.getValue(typeParameterName)
+                val type = TypeParser.parseTypeFromString(typeString, scopeForTypeParameters, allTypeParameterSymbols)
+                substitution(typeParameterSymbol, type)
             }
         }
     }
 
-    object Directives : SimpleDirectivesContainer() {
-        val SUBSTITUTOR by valueDirective("") { stringValue ->
-            stringValue.trim().split(";").map { it.trim() }.map { substitution ->
+    private fun parseSubstitutions(substitutionsAsString: String): List<Pair<String, String>> =
+        substitutionsAsString.trim().split(",")
+            .map { it.trim() }
+            .filter { it.isNotEmpty() }
+            .map { substitution ->
                 val asList = substitution.split("->").map { it.trim() }
                 check(asList.size == 2) {
                     "Substitution should look like `x -> y` but was `$substitution`"
                 }
                 asList[0] to asList[1]
             }
-        }
-    }
+
+    const val SUBSTITUTOR_PREFIX = "// SUBSTITUTOR:"
 }

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/TypeParser.kt

@@ -15,16 +15,18 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 
 object TypeParser {
-    fun KtAnalysisSession.parseTypeFromString(
+    context (KtAnalysisSession)
+    fun parseTypeFromString(
         stringType: String,
-        ktFile: KtFile,
+        contextElement: KtElement,
         typeParameterByName: Map<String, KtTypeParameterSymbol>
     ): KtType {
-        val type = KtPsiFactory(ktFile).createType(stringType)
+        val type = KtPsiFactory(contextElement).createType(stringType)
         return convertType(type.typeElement ?: incorrectType(type), typeParameterByName)
     }
 
-    private fun KtAnalysisSession.convertType(type: KtTypeElement, typeParameterByName: Map<String, KtTypeParameterSymbol>): KtType =
+    context (KtAnalysisSession)
+            private fun convertType(type: KtTypeElement, typeParameterByName: Map<String, KtTypeParameterSymbol>): KtType =
         when (type) {
             is KtUserType -> {
                 val qualifier = fullQualifier(type)