[LL API] Add basic tests for 'ContextCollector'

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.test.framework.services
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.util.PrivateForInline
@@ -15,6 +16,8 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.model.TestFile
 import org.jetbrains.kotlin.test.model.TestModule
@@ -155,6 +158,27 @@ class ExpressionMarkerProvider : TestService {
         }
     }
 
+    /**
+     * Find the bottommost element of an [elementType] or its subtype located precisely in the [range].
+     */
+    private fun <T : PsiElement> getBottommostElementOfTypeInRange(file: KtFile, range: TextRange, elementType: Class<T>): T {
+        var candidate = PsiTreeUtil.findElementOfClassAtOffset(file, range.startOffset, elementType, true)
+        while (candidate != null && candidate.endOffset < range.endOffset) {
+            candidate = PsiTreeUtil.getParentOfType(candidate, elementType)?.takeIf { it.startOffset == range.startOffset }
+        }
+
+        return candidate?.takeIf { it.endOffset == range.endOffset }
+            ?: error("Cannot find '${elementType.name}' in range $range")
+    }
+
+    /**
+     * Find the bottommost element of [E] or its subtype wrapped in an '<expr>' selection tag.
+     */
+    fun <E : KtElement> getBottommostSelectedElementOfType(file: KtFile, elementType: Class<E>): E {
+        val range = selected[file.name] ?: error("No selected expression found in file")
+        return getBottommostElementOfTypeInRange(file, range, elementType)
+    }
+
     private fun List<PsiElement>.trimWhitespaces(): List<PsiElement> =
         dropWhile { it is PsiWhiteSpace }
             .dropLastWhile { it is PsiWhiteSpace }

analysis/low-level-api-fir/testdata/contextCollector/contextReceivers.kt

@@ -0,0 +1,12 @@
+interface Foo {
+    fun foo(): Int
+}
+
+interface Bar {
+    fun bar(): Int
+}
+
+context(Foo, Bar)
+fun test() {
+    <expr>foo() + bar()</expr>
+}

analysis/low-level-api-fir/testdata/contextCollector/contextReceivers.txt

@@ -0,0 +1,29 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+    Element 8
+        Context receivers:
+            FirNamedFunctionSymbol context(R|Foo|, R|Bar|)
+            public final fun test(): R|kotlin/Unit|
+                Type: Foo
+                Label: Foo
+            FirNamedFunctionSymbol context(R|Foo|, R|Bar|)
+            public final fun test(): R|kotlin/Unit|
+                Type: Bar
+                Label: Bar
+    Element 9
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/contextReceiversClass.kt

@@ -0,0 +1,10 @@
+interface Foo {
+    fun foo(): Int
+}
+
+context(Foo)
+class Test {
+    fun test() {
+        <expr>foo()</expr>
+    }
+}

analysis/low-level-api-fir/testdata/contextCollector/contextReceiversClass.txt

@@ -0,0 +1,30 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Implicit receiver:
+            FirRegularClassSymbol context(R|Foo|)
+            public final class Test : R|kotlin/Any|
+                Type: Test
+    Element 8
+        Context receivers:
+            FirRegularClassSymbol context(R|Foo|)
+            public final class Test : R|kotlin/Any|
+                Type: Foo
+                Label: Foo
+    Element 9
+        Scope: FirLocalScope
+    Element 10
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/extensionFunction.kt

@@ -0,0 +1,5 @@
+fun String.test() {
+    <expr>consume(this)</expr>
+}
+
+fun consume(obj: Any) {}

analysis/low-level-api-fir/testdata/contextCollector/extensionFunction.txt

@@ -0,0 +1,23 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+    Element 8
+        Implicit receiver:
+            FirNamedFunctionSymbol public final fun R|kotlin/String|.test(): R|kotlin/Unit|
+                Type: kotlin/String
+    Element 9
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/extensionLambdas.kt

@@ -0,0 +1,13 @@
+fun test() {
+    block("foo") foo@{
+        block(42) num@{
+            <expr>consume(this@foo.length + this@num)</expr>
+        }
+    }
+}
+
+fun <T> block(reciever: T, block: T.() -> Unit) {
+    receiver.block()
+}
+
+fun consume(num: Int) {}

analysis/low-level-api-fir/testdata/contextCollector/extensionLambdas.txt

@@ -0,0 +1,35 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+    Element 10
+        Implicit receiver:
+            FirAnonymousFunctionSymbol foo@fun R|kotlin/String|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>
+                Type: kotlin/String
+    Element 11
+        Scope: FirLocalScope
+    Element 12
+        Scope: FirLocalScope
+    Element 13
+        Implicit receiver:
+            FirAnonymousFunctionSymbol num@fun R|kotlin/Int|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>
+                Type: kotlin/Int
+    Element 14
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/innerClasses.kt

@@ -0,0 +1,11 @@
+package test
+
+class Foo<F> {
+    inner class Bar<B> {
+        fun test(f: F, b: B) {
+            <expr>consume(f)</expr>
+        }
+    }
+}
+
+fun consume(obj: Any) {}

analysis/low-level-api-fir/testdata/contextCollector/innerClasses.txt

@@ -0,0 +1,42 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirMemberTypeParameterScope
+            Classifiers:
+                FirTypeParameterSymbol F
+    Element 8
+        Scope: FirNestedClassifierScopeImpl
+            Classifiers:
+                FirRegularClassSymbol public final inner class Bar<B, F> : R|kotlin/Any|
+    Element 9
+        Implicit receiver:
+            FirRegularClassSymbol public final class Foo<F> : R|kotlin/Any|
+                Type: test/Foo<F>
+    Element 10
+        Scope: FirMemberTypeParameterScope
+            Classifiers:
+                FirTypeParameterSymbol B
+    Element 11
+        Implicit receiver:
+            FirRegularClassSymbol public final inner class Bar<B, F> : R|kotlin/Any|
+                Type: test/Foo.Bar<B, F>
+    Element 12
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol b: R|B|
+                FirValueParameterSymbol f: R|F|
+    Element 13
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/lambdaArguments.kt

@@ -0,0 +1,11 @@
+fun test() {
+    block("foo") { s ->
+        block(42) { n ->
+            <expr>s.length + n</expr>
+        }
+    }
+}
+
+fun <T> block(obj: T, block: (T) -> Unit) {
+    block(obj)
+}

analysis/low-level-api-fir/testdata/contextCollector/lambdaArguments.txt

@@ -0,0 +1,31 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol s: R|kotlin/String|
+    Element 10
+        Scope: FirLocalScope
+    Element 11
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol n: R|kotlin/Int|
+    Element 12
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/nestedClasses.kt

@@ -0,0 +1,11 @@
+package test
+
+class Foo<F> {
+    class Bar<B> {
+        fun test(b: B) {
+            <expr>consume(b)</expr>
+        }
+    }
+}
+
+fun consume(obj: Any) {}

analysis/low-level-api-fir/testdata/contextCollector/nestedClasses.txt

@@ -0,0 +1,33 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirNestedClassifierScopeImpl
+            Classifiers:
+                FirRegularClassSymbol public final class Bar<B> : R|kotlin/Any|
+    Element 8
+        Scope: FirMemberTypeParameterScope
+            Classifiers:
+                FirTypeParameterSymbol B
+    Element 9
+        Implicit receiver:
+            FirRegularClassSymbol public final class Bar<B> : R|kotlin/Any|
+                Type: test/Foo.Bar<B>
+    Element 10
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol b: R|B|
+    Element 11
+        Scope: FirLocalScope

analysis/low-level-api-fir/testdata/contextCollector/simple.kt

@@ -0,0 +1,10 @@
+package foo
+
+class Foo {
+    val a: Int
+
+    fun foo(p: String) {
+        val x = 1.0f
+        <expr>print(x)</expr>
+    }
+}

analysis/low-level-api-fir/testdata/contextCollector/simple.txt

@@ -0,0 +1,27 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Implicit receiver:
+            FirRegularClassSymbol public final class Foo : R|kotlin/Any|
+                Type: foo/Foo
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol p: R|kotlin/String|
+    Element 9
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval x: R|kotlin/Float|

analysis/low-level-api-fir/testdata/contextCollector/smartCastArgument.kt

@@ -0,0 +1,9 @@
+interface Foo {}
+
+fun test(obj: Any) {
+    if (obj is Foo) {
+        consume(<expr>obj</expr>)
+    }
+}
+
+fun consume(obj: Any) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCastArgument.txt

@@ -0,0 +1,28 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+Smart Casts:
+    FirValueParameterSymbol obj: R|kotlin/Any|
+    Types:
+        Foo
+        kotlin/Any

analysis/low-level-api-fir/testdata/contextCollector/smartCastReceiverArgument.kt

@@ -0,0 +1,12 @@
+interface Foo {}
+interface Bar {}
+
+fun test(obj: Any) {
+    if (obj is Foo) {
+        if (obj is Bar) {
+            obj.consume(<expr>obj</expr>)
+        }
+    }
+}
+
+fun Bar.consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCastReceiverArgument.txt

@@ -0,0 +1,31 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+    Element 10
+        Scope: FirLocalScope
+Smart Casts:
+    FirValueParameterSymbol obj: R|kotlin/Any|
+    Types:
+        Foo
+        kotlin/Any
+        Bar

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractContextCollectorTest.kt

@@ -0,0 +1,242 @@
+/*
+ * Copyright 2010-2023 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.analysis.low.level.api.fir
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedSingleModuleTest
+import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
+import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
+import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
+import org.jetbrains.kotlin.fir.renderer.FirRenderer
+import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
+import org.jetbrains.kotlin.fir.scopes.*
+import org.jetbrains.kotlin.fir.scopes.impl.*
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.renderReadableWithFqNames
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.model.TestModule
+import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.assertions
+import java.lang.IllegalArgumentException
+
+abstract class AbstractContextCollectorTest : AbstractAnalysisApiBasedSingleModuleTest() {
+    override val configurator = AnalysisApiFirSourceTestConfigurator(analyseInDependentSession = false)
+
+    override fun doTestByFileStructure(ktFiles: List<KtFile>, module: TestModule, testServices: TestServices) {
+        val mainKtFile = ktFiles.singleOrNull() ?: ktFiles.single { it.name == "main.kt" }
+
+        val project = mainKtFile.project
+        val sourceModule = ProjectStructureProvider.getModule(project, mainKtFile, contextualModule = null)
+
+        val resolveSession = sourceModule.getFirResolveSession(project)
+        val session = resolveSession.useSiteFirSession
+        val sessionHolder = SessionHolderImpl(session, session.getScopeSession())
+
+        val firFile = mainKtFile.getOrBuildFirFile(resolveSession)
+
+        val targetElement = testServices.expressionMarkerProvider
+            .getBottommostSelectedElementOfType(mainKtFile, KtElement::class.java)
+
+        val elementContext = ContextCollector.process(firFile, sessionHolder, targetElement) { it === targetElement }
+            ?: error("Context not found for element $targetElement")
+
+        val actualText = ElementContextRenderer.render(elementContext)
+        testServices.assertions.assertEqualsToTestDataFileSibling(actualText)
+    }
+}
+
+private object ElementContextRenderer {
+    fun render(context: ContextCollector.Context): String {
+        return buildString {
+            renderTowerDataContext(context.towerDataContext)
+            renderSmartCasts(context.smartCasts)
+        }.trim()
+    }
+
+    private fun StringBuilder.renderTowerDataContext(towerDataContext: FirTowerDataContext) {
+        appendBlock("Tower Data Context:") {
+            for ((index, towerDataElement) in towerDataContext.towerDataElements.withIndex()) {
+                appendBlock("Element $index") {
+                    for (scope in towerDataElement.scope?.flatten().orEmpty()) {
+                        appendBlock("Scope: " + scope.javaClass.simpleName) {
+                            renderScope(scope)
+                        }
+                    }
+
+                    towerDataElement.implicitReceiver?.let { implicitReceiver ->
+                        appendBlock("Implicit receiver:") {
+                            appendSymbol(implicitReceiver.boundSymbol).appendLine()
+
+                            appendBlock {
+                                append("Type: ").appendType(implicitReceiver.type).appendLine()
+                                if (implicitReceiver.isContextReceiver) {
+                                    append("Context receiver index: ").appendLine(implicitReceiver.contextReceiverNumber)
+                                }
+                            }
+                        }
+                    }
+
+                    towerDataElement.contextReceiverGroup?.takeIf { it.isNotEmpty() }?.let { contextReceiverValues ->
+                        appendBlock("Context receivers:") {
+                            for (contextReceiverValue in contextReceiverValues) {
+                                appendSymbol(contextReceiverValue.boundSymbol).appendLine()
+
+                                appendBlock {
+                                    append("Type: ").appendType(contextReceiverValue.type).appendLine()
+                                    contextReceiverValue.labelName?.let { labelName ->
+                                        append("Label: ").appendLine(labelName)
+                                    }
+                                }
+                            }
+                        }
+                    }
+
+                    towerDataElement.staticScopeOwnerSymbol?.let { staticScopeOwnerSymbol ->
+                        append("Static scope owner symbol: ").appendSymbol(staticScopeOwnerSymbol).appendLine()
+                    }
+                }
+            }
+        }
+    }
+
+    private fun StringBuilder.renderScope(scope: FirScope) {
+        when (scope) {
+            is FirDefaultSimpleImportingScope, is FirDefaultStarImportingScope -> {
+                Unit
+                // Skip to avoid fixing default imports in an unrelated test
+            }
+            is FirPackageMemberScope -> {
+                // Skip as the scope can be huge
+            }
+            is FirAbstractSimpleImportingScope -> {
+                for (import in scope.simpleImports.flatMap { it.value }) {
+                    renderImport(import)
+                }
+            }
+            is FirAbstractStarImportingScope -> {
+                for (import in scope.starImports) {
+                    renderImport(import)
+                }
+            }
+            is FirContainingNamesAwareScope -> {
+                val classifierNames = scope.getClassifierNames().sorted()
+                val callableNames = scope.getCallableNames().sorted()
+
+                fun <T : FirBasedSymbol<*>> appendDeclarations(title: String, names: List<Name>, collector: (Name, (T) -> Unit) -> Unit) {
+                    collect(names, collector).takeIf { it.isNotEmpty() }?.let { classifiers ->
+                        appendBlock(title) {
+                            classifiers.forEach { appendSymbol(it).appendLine() }
+                        }
+                    }
+                }
+
+                appendDeclarations("Classifiers:", classifierNames, scope::processClassifiersByName)
+                appendDeclarations("Functions", callableNames, scope::processFunctionsByName)
+                appendDeclarations("Properties:", callableNames, scope::processPropertiesByName)
+            }
+            else -> {
+                throw IllegalArgumentException("Unexpected scope type: " + scope.javaClass.name)
+            }
+        }
+    }
+
+    private fun <T> collect(names: List<Name>, collector: (Name, (T) -> Unit) -> Unit): List<T> {
+        return buildList {
+            for (name in names) {
+                collector(name) { add(it) }
+            }
+        }
+    }
+
+    private fun StringBuilder.renderImport(import: FirResolvedImport) {
+        val name = import.importedName ?: SpecialNames.NO_NAME_PROVIDED
+
+        appendBlock("Import name:" + name.asString()) {
+            import.importedFqName?.let { importedFqName ->
+                append("Qualified name: ").append(importedFqName.asString()).appendLine()
+            }
+
+            append("Is all under: ").append(import.isAllUnder).appendLine()
+
+            import.aliasName?.let { aliasName ->
+                append("Alias: ").append(aliasName.asString()).appendLine()
+            }
+        }
+    }
+
+    private fun StringBuilder.renderSmartCasts(smartCasts: Map<FirBasedSymbol<*>, Set<ConeKotlinType>>) {
+        if (smartCasts.isEmpty()) {
+            return
+        }
+
+        appendBlock("Smart Casts:") {
+            for ((symbol, types) in smartCasts) {
+                appendSymbol(symbol).appendLine()
+
+                appendBlock("Types:") {
+                    for (type in types) {
+                        appendType(type).appendLine()
+                    }
+                }
+
+            }
+        }
+    }
+
+    private fun StringBuilder.appendBlock(title: String? = null, block: StringBuilder.() -> Unit): StringBuilder {
+        if (title != null) {
+            appendLine(title)
+        }
+
+        val nestedText = StringBuilder().apply(block).toString().trim()
+
+        if (nestedText.isNotBlank()) {
+            for (line in nestedText.lineSequence()) {
+                appendIndent().appendLine(line)
+            }
+        }
+
+        return this
+    }
+
+    private fun StringBuilder.appendSymbol(symbol: FirBasedSymbol<*>): StringBuilder {
+        val renderer = FirRenderer(
+            annotationRenderer = null,
+            bodyRenderer = null,
+            classMemberRenderer = null,
+            contractRenderer = null
+        )
+
+        val type = symbol.javaClass.simpleName
+        val text = renderer.renderElementAsString(symbol.fir, trim = true)
+        return append(type).append(' ').append(text)
+    }
+
+    private fun StringBuilder.appendType(type: ConeKotlinType): StringBuilder {
+        return append(type.renderReadableWithFqNames())
+    }
+
+    private fun StringBuilder.appendIndent(): StringBuilder {
+        append("    ")
+        return this
+    }
+}
+
+private fun FirScope.flatten(): List<FirScope> {
+    return when (this) {
+        is FirCompositeScope -> scopes.flatMap { it.flatten() }
+        is FirNameAwareCompositeScope -> scopes.flatMap { it.flatten() }
+        else -> listOf(this)
+    }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorTestGenerated.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2010-2023 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.analysis.low.level.api.fir;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.analysis.api.GenerateAnalysisApiTestsKt}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("analysis/low-level-api-fir/testdata/contextCollector")
+@TestDataPath("$PROJECT_ROOT")
+public class ContextCollectorTestGenerated extends AbstractContextCollectorTest {
+    @Test
+    public void testAllFilesPresentInContextCollector() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/contextCollector"), Pattern.compile("^(.+)\\.kt$"), null, true);
+    }
+
+    @Test
+    @TestMetadata("contextReceivers.kt")
+    public void testContextReceivers() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/contextReceivers.kt");
+    }
+
+    @Test
+    @TestMetadata("contextReceiversClass.kt")
+    public void testContextReceiversClass() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/contextReceiversClass.kt");
+    }
+
+    @Test
+    @TestMetadata("extensionFunction.kt")
+    public void testExtensionFunction() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/extensionFunction.kt");
+    }
+
+    @Test
+    @TestMetadata("extensionLambdas.kt")
+    public void testExtensionLambdas() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/extensionLambdas.kt");
+    }
+
+    @Test
+    @TestMetadata("innerClasses.kt")
+    public void testInnerClasses() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/innerClasses.kt");
+    }
+
+    @Test
+    @TestMetadata("lambdaArguments.kt")
+    public void testLambdaArguments() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/lambdaArguments.kt");
+    }
+
+    @Test
+    @TestMetadata("nestedClasses.kt")
+    public void testNestedClasses() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/nestedClasses.kt");
+    }
+
+    @Test
+    @TestMetadata("simple.kt")
+    public void testSimple() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/simple.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastArgument.kt")
+    public void testSmartCastArgument() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCastArgument.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastReceiverArgument.kt")
+    public void testSmartCastReceiverArgument() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCastReceiverArgument.kt");
+    }
+}

generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/firLowLevel.kt

@@ -179,6 +179,10 @@ internal fun TestGroupSuite.generateFirLowLevelApiTests() {
         testClass<AbstractCompilationPeerAnalysisTest> {
             model("compilationPeers")
         }
+
+        testClass<AbstractContextCollectorTest> {
+            model("contextCollector")
+        }
     }
 
     testGroup("analysis/low-level-api-fir/tests", "analysis/analysis-api/testData") {