[FIR] Introduce CheckerContext for call and declaration checkers

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/call/FirExpressionChecker.kt

@@ -5,19 +5,15 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.call
 
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 
 abstract class FirExpressionChecker<in E : FirExpression> {
-    abstract fun check(functionCall: E, reporter: DiagnosticReporter)
+    abstract fun check(functionCall: E, context: CheckerContext, reporter: DiagnosticReporter)
 }
 
 typealias FirFunctionCallChecker = FirExpressionChecker<FirFunctionCall>
-typealias FirQualifiedAccessChecker = FirExpressionChecker<FirQualifiedAccessExpression>
+typealias FirQualifiedAccessChecker = FirExpressionChecker<FirQualifiedAccessExpression>
-
-class FirExpressionCheckerContext(
-    val implicitReceiverStack: ImplicitReceiverStack
-)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/CheckerContext.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2020 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.fir.analysis.checkers.context
+
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.persistentListOf
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.name.Name
+
+abstract class CheckerContext {
+    abstract val implicitReceiverStack: ImplicitReceiverStack
+    abstract val containingDeclarations: List<FirDeclaration>
+}
+
+class PersistentCheckerContext(
+    override val implicitReceiverStack: PersistentImplicitReceiverStack,
+    override val containingDeclarations: PersistentList<FirDeclaration>
+) : CheckerContext() {
+    constructor() : this(
+        PersistentImplicitReceiverStack(),
+        persistentListOf()
+    )
+
+    fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext {
+        return PersistentCheckerContext(
+            implicitReceiverStack.add(name, value),
+            containingDeclarations
+        )
+    }
+
+    fun addDeclaration(declaration: FirDeclaration): PersistentCheckerContext {
+        return PersistentCheckerContext(
+            implicitReceiverStack,
+            containingDeclarations.add(declaration)
+        )
+    }
+}

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/context/PersistentImplicitReceiverStack.kt

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2010-2020 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.fir.analysis.checkers.context
+
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.persistentListOf
+import kotlinx.collections.immutable.persistentMapOf
+import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
+import org.jetbrains.kotlin.fir.resolve.PersistentSetMultimap
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.name.Name
+
+class PersistentImplicitReceiverStack private constructor(
+    private val stack: PersistentList<ImplicitReceiverValue<*>>,
+    // This multi-map holds indexes of the stack ^
+    private val indexesPerLabel: PersistentSetMultimap<Name, Int>,
+    private val indexesPerSymbol: PersistentMap<FirBasedSymbol<*>, Int>
+) : ImplicitReceiverStack(), Iterable<ImplicitReceiverValue<*>> {
+    val size: Int get() = stack.size
+
+    constructor() : this(
+        persistentListOf(),
+        PersistentSetMultimap(),
+        persistentMapOf()
+    )
+
+    fun add(name: Name?, value: ImplicitReceiverValue<*>): PersistentImplicitReceiverStack {
+        val stack = stack.add(value)
+        val index = stack.size - 1
+        val indexesPerLabel = name?.let { indexesPerLabel.put(it, index) } ?: indexesPerLabel
+        val indexesPerSymbol = indexesPerSymbol.put(value.boundSymbol, index)
+        return PersistentImplicitReceiverStack(
+            stack,
+            indexesPerLabel,
+            indexesPerSymbol
+        )
+    }
+
+    override operator fun get(name: String?): ImplicitReceiverValue<*>? {
+        if (name == null) return stack.lastOrNull()
+        return indexesPerLabel[Name.identifier(name)].lastOrNull()?.let { stack[it] }
+    }
+
+    override fun lastDispatchReceiver(): ImplicitDispatchReceiverValue? {
+        return stack.filterIsInstance<ImplicitDispatchReceiverValue>().lastOrNull()
+    }
+
+    override fun receiversAsReversed(): List<ImplicitReceiverValue<*>> = stack.asReversed()
+
+    override operator fun iterator(): Iterator<ImplicitReceiverValue<*>> {
+        return stack.iterator()
+    }
+}

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirDeclarationChecker.kt

@@ -5,9 +5,10 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 
 abstract class FirDeclarationChecker<in D : FirDeclaration> {
-    abstract fun check(declaration: D, reporter: DiagnosticReporter)
+    abstract fun check(declaration: D, context: CheckerContext, reporter: DiagnosticReporter)
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirInfixFunctionDeclarationChecker.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.onSource
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
@@ -14,7 +15,7 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.isInfix
 
 object FirInfixFunctionDeclarationChecker : FirDeclarationChecker<FirMemberDeclaration>() {
-    override fun check(declaration: FirMemberDeclaration, reporter: DiagnosticReporter) {
+    override fun check(declaration: FirMemberDeclaration, context: CheckerContext, reporter: DiagnosticReporter) {
         if (declaration is FirSimpleFunction && declaration.isInfix) {
             if (declaration.valueParameters.size != 1 || declaration.receiverTypeRef == null) {
                 reporter.report(declaration.source)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollector.kt

@@ -6,16 +6,30 @@
 package org.jetbrains.kotlin.fir.analysis.collectors
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
 import org.jetbrains.kotlin.fir.analysis.collectors.components.DeclarationCheckersDiagnosticComponent
 import org.jetbrains.kotlin.fir.analysis.collectors.components.ErrorNodeDiagnosticCollectorComponent
 import org.jetbrains.kotlin.fir.analysis.collectors.components.ExpressionCheckersDiagnosticComponent
 import org.jetbrains.kotlin.fir.analysis.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.SessionHolder
+import org.jetbrains.kotlin.fir.resolve.collectImplicitReceivers
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.name.Name
 
-abstract class AbstractDiagnosticCollector {
+abstract class AbstractDiagnosticCollector(
+    override val session: FirSession,
+    override val scopeSession: ScopeSession = ScopeSession()
+) : SessionHolder {
     fun collectDiagnostics(firFile: FirFile): Iterable<ConeDiagnostic> {
         if (!componentsInitialized) {
             throw IllegalStateException("Components are not initialized")
@@ -33,6 +47,8 @@ abstract class AbstractDiagnosticCollector {
     private var componentsInitialized = false
     private val visitor = Visitor()
 
+    private var context = PersistentCheckerContext()
+
     fun initializeComponents(vararg components: AbstractDiagnosticCollectorComponent) {
         if (componentsInitialized) {
             throw IllegalStateException()
@@ -44,7 +60,7 @@ abstract class AbstractDiagnosticCollector {
     private inner class Visitor : FirVisitorVoid() {
         private fun <T : FirElement> T.runComponents() {
             components.forEach {
-                this.accept(it)
+                this.accept(it, context)
             }
         }
 
@@ -52,6 +68,102 @@ abstract class AbstractDiagnosticCollector {
             element.runComponents()
             element.acceptChildren(this)
         }
+
+        override fun visitRegularClass(regularClass: FirRegularClass) {
+            val typeRef = buildResolvedTypeRef {
+                type = regularClass.defaultType()
+            }
+            visitWithDeclarationAndReceiver(regularClass, regularClass.name, typeRef)
+        }
+
+        override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) {
+            visitWithDeclarationAndReceiver(simpleFunction, simpleFunction.name, simpleFunction.receiverTypeRef)
+        }
+
+        override fun visitConstructor(constructor: FirConstructor) {
+            visitWithDeclaration(constructor)
+        }
+
+        override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+            val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }
+            visitWithDeclarationAndReceiver(
+                anonymousFunction,
+                labelName,
+                anonymousFunction.receiverTypeRef
+            )
+        }
+
+        override fun visitProperty(property: FirProperty) {
+            visitWithDeclaration(property)
+        }
+
+        override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor) {
+            val property = context.containingDeclarations.last() as FirProperty
+            visitWithDeclarationAndReceiver(propertyAccessor, property.name, property.receiverTypeRef)
+        }
+
+        override fun visitValueParameter(valueParameter: FirValueParameter) {
+            visitWithDeclaration(valueParameter)
+        }
+
+        override fun visitEnumEntry(enumEntry: FirEnumEntry) {
+            visitWithDeclaration(enumEntry)
+        }
+
+        override fun visitFile(file: FirFile) {
+            visitWithDeclaration(file)
+        }
+
+        private fun visitWithDeclaration(declaration: FirDeclaration) {
+            declaration.runComponents()
+            withDeclaration(declaration) {
+                declaration.acceptChildren(this)
+            }
+        }
+
+        private fun visitWithDeclarationAndReceiver(declaration: FirDeclaration, labelName: Name?, receiverTypeRef: FirTypeRef?) {
+            declaration.runComponents()
+            withDeclaration(declaration) {
+                withLabelAndReceiverType(
+                    labelName,
+                    declaration,
+                    receiverTypeRef?.coneTypeSafe()
+                ) {
+                    declaration.acceptChildren(this)
+                }
+            }
+        }
+    }
+
+    private inline fun <R> withDeclaration(declaration: FirDeclaration, block: () -> R): R {
+        val existingContext = context
+        context = context.addDeclaration(declaration)
+        try {
+            return block()
+        } finally {
+            context = existingContext
+        }
+    }
+
+    private inline fun <R> withLabelAndReceiverType(
+        labelName: Name?,
+        owner: FirDeclaration,
+        type: ConeKotlinType?,
+        block: () -> R
+    ): R {
+        val (implicitReceiverValue, implicitCompanionValues) = collectImplicitReceivers(type, owner)
+        val existingContext = context
+        implicitCompanionValues.forEach { value ->
+            context = context.addImplicitReceiver(null, value)
+        }
+        implicitReceiverValue?.let {
+            context = context.addImplicitReceiver(labelName, it)
+        }
+        try {
+            return block()
+        } finally {
+            context = existingContext
+        }
     }
 }
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/FirDiagnosticsCollector.kt

@@ -5,17 +5,19 @@
 
 package org.jetbrains.kotlin.fir.analysis.collectors
 
+import org.jetbrains.kotlin.fir.FirSession
+
 object FirDiagnosticsCollector {
-    fun create(): AbstractDiagnosticCollector {
+    fun create(session: FirSession): AbstractDiagnosticCollector {
-        val collector = SimpleDiagnosticsCollector()
+        val collector = SimpleDiagnosticsCollector(session)
         collector.registerAllComponents()
         return collector
     }
 
     // Use in CLI compiler
     @Suppress("unused")
-    fun createParallel(): AbstractDiagnosticCollector {
+    fun createParallel(session: FirSession): AbstractDiagnosticCollector {
-        val collector = ParallelDiagnosticsCollector(4)
+        val collector = ParallelDiagnosticsCollector(session, numberOfThreads = 4)
         collector.registerAllComponents()
         return collector
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/ParallelDiagnosticsCollector.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.collectors
 
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.SimpleDiagnosticReporter
@@ -13,7 +14,7 @@ import java.util.concurrent.Executors
 import java.util.concurrent.Future
 import java.util.concurrent.atomic.AtomicInteger
 
-class ParallelDiagnosticsCollector(private val numberOfThreads: Int) : AbstractDiagnosticCollector() {
+class ParallelDiagnosticsCollector(session: FirSession, private val numberOfThreads: Int) : AbstractDiagnosticCollector(session) {
     init {
         require(numberOfThreads >= 1) {
             "Number of threads should be at least 1"

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/SimpleDiagnosticsCollector.kt

@@ -5,11 +5,12 @@
 
 package org.jetbrains.kotlin.fir.analysis.collectors
 
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.SimpleDiagnosticReporter
 
-class SimpleDiagnosticsCollector : AbstractDiagnosticCollector() {
+class SimpleDiagnosticsCollector(session: FirSession) : AbstractDiagnosticCollector(session) {
     private var reporter = SimpleDiagnosticReporter()
 
     override fun initializeCollector() {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/AbstractDiagnosticCollectorComponent.kt

@@ -6,12 +6,13 @@
 package org.jetbrains.kotlin.fir.analysis.collectors.components
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
-abstract class AbstractDiagnosticCollectorComponent(private val collector: AbstractDiagnosticCollector) : FirVisitorVoid() {
+abstract class AbstractDiagnosticCollectorComponent(private val collector: AbstractDiagnosticCollector) : FirVisitor<Unit, CheckerContext>() {
-    override fun visitElement(element: FirElement) {}
+    override fun visitElement(element: FirElement, data: CheckerContext) {}
 
     protected fun runCheck(block: (DiagnosticReporter) -> Unit) {
         collector.runCheck(block)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/DeclarationCheckersDiagnosticComponent.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.collectors.components
 
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirDeclarationChecker
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
@@ -12,45 +13,45 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.declarations.*
 
 class DeclarationCheckersDiagnosticComponent(collector: AbstractDiagnosticCollector) : AbstractDiagnosticCollectorComponent(collector) {
-    override fun visitProperty(property: FirProperty) {
+    override fun visitProperty(property: FirProperty, data: CheckerContext) {
-        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(property, it) }
+        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(property, data, it) }
     }
 
-    override fun visitRegularClass(regularClass: FirRegularClass) {
+    override fun visitRegularClass(regularClass: FirRegularClass, data: CheckerContext) {
-        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(regularClass, it) }
+        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(regularClass, data, it) }
     }
 
-    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) {
+    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: CheckerContext) {
-        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(simpleFunction, it) }
+        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(simpleFunction, data, it) }
     }
 
-    override fun visitTypeAlias(typeAlias: FirTypeAlias) {
+    override fun visitTypeAlias(typeAlias: FirTypeAlias, data: CheckerContext) {
-        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(typeAlias, it) }
+        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(typeAlias, data, it) }
     }
 
-    override fun visitConstructor(constructor: FirConstructor) {
+    override fun visitConstructor(constructor: FirConstructor, data: CheckerContext) {
-        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(constructor, it) }
+        runCheck { DeclarationCheckers.MEMBER_DECLARATIONS.check(constructor, data, it) }
     }
 
-    override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+    override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: CheckerContext) {
-        runCheck { DeclarationCheckers.DECLARATIONS.check(anonymousFunction, it) }
+        runCheck { DeclarationCheckers.DECLARATIONS.check(anonymousFunction, data, it) }
     }
 
-    override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor) {
+    override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: CheckerContext) {
-        runCheck { DeclarationCheckers.DECLARATIONS.check(propertyAccessor, it) }
+        runCheck { DeclarationCheckers.DECLARATIONS.check(propertyAccessor, data, it) }
     }
 
-    override fun visitValueParameter(valueParameter: FirValueParameter) {
+    override fun visitValueParameter(valueParameter: FirValueParameter, data: CheckerContext) {
-        runCheck { DeclarationCheckers.DECLARATIONS.check(valueParameter, it) }
+        runCheck { DeclarationCheckers.DECLARATIONS.check(valueParameter, data, it) }
     }
 
-    override fun visitEnumEntry(enumEntry: FirEnumEntry) {
+    override fun visitEnumEntry(enumEntry: FirEnumEntry, data: CheckerContext) {
-        runCheck { DeclarationCheckers.DECLARATIONS.check(enumEntry, it) }
+        runCheck { DeclarationCheckers.DECLARATIONS.check(enumEntry, data, it) }
     }
 
-    private fun <D : FirDeclaration> List<FirDeclarationChecker<D>>.check(declaration: D, reporter: DiagnosticReporter) {
+    private fun <D : FirDeclaration> List<FirDeclarationChecker<D>>.check(declaration: D, context: CheckerContext, reporter: DiagnosticReporter) {
         for (checker in this) {
-            checker.check(declaration, reporter)
+            checker.check(declaration, context, reporter)
         }
     }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/ErrorNodeDiagnosticCollectorComponent.kt

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory0
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
@@ -16,8 +17,8 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.onSource
 import org.jetbrains.kotlin.fir.declarations.FirErrorFunction
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnostic
-import org.jetbrains.kotlin.fir.diagnostics.FirStubDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.FirSimpleDiagnostic
+import org.jetbrains.kotlin.fir.diagnostics.FirStubDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirErrorExpression
 import org.jetbrains.kotlin.fir.expressions.FirErrorLoop
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
@@ -25,27 +26,27 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
 
 class ErrorNodeDiagnosticCollectorComponent(collector: AbstractDiagnosticCollector) : AbstractDiagnosticCollectorComponent(collector) {
-    override fun visitErrorLoop(errorLoop: FirErrorLoop) {
+    override fun visitErrorLoop(errorLoop: FirErrorLoop, data: CheckerContext) {
         val source = errorLoop.source ?: return
         runCheck { reportFirDiagnostic(errorLoop.diagnostic, source, it) }
     }
 
-    override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef) {
+    override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef, data: CheckerContext) {
         val source = errorTypeRef.source ?: return
         runCheck { reportFirDiagnostic(errorTypeRef.diagnostic, source, it) }
     }
 
-    override fun visitErrorNamedReference(errorNamedReference: FirErrorNamedReference) {
+    override fun visitErrorNamedReference(errorNamedReference: FirErrorNamedReference, data: CheckerContext) {
         val source = errorNamedReference.source ?: return
         runCheck { reportFirDiagnostic(errorNamedReference.diagnostic, source, it) }
     }
 
-    override fun visitErrorExpression(errorExpression: FirErrorExpression) {
+    override fun visitErrorExpression(errorExpression: FirErrorExpression, data: CheckerContext) {
         val source = errorExpression.source ?: return
         runCheck { reportFirDiagnostic(errorExpression.diagnostic, source, it) }
     }
 
-    override fun visitErrorFunction(errorFunction: FirErrorFunction) {
+    override fun visitErrorFunction(errorFunction: FirErrorFunction, data: CheckerContext) {
         val source = errorFunction.source ?: return
         runCheck { reportFirDiagnostic(errorFunction.diagnostic, source, it) }
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/ExpressionCheckersDiagnosticComponent.kt

@@ -7,78 +7,79 @@ package org.jetbrains.kotlin.fir.analysis.collectors.components
 
 import org.jetbrains.kotlin.fir.analysis.checkers.call.CallCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.call.FirExpressionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.expressions.*
 
 class ExpressionCheckersDiagnosticComponent(collector: AbstractDiagnosticCollector) : AbstractDiagnosticCollectorComponent(collector) {
-    override fun visitTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall) {
+    override fun visitTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(typeOperatorCall, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(typeOperatorCall, data, it) }
     }
 
-    override fun <T> visitConstExpression(constExpression: FirConstExpression<T>) {
+    override fun <T> visitConstExpression(constExpression: FirConstExpression<T>, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(constExpression, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(constExpression, data, it) }
     }
 
-    override fun visitAnnotationCall(annotationCall: FirAnnotationCall) {
+    override fun visitAnnotationCall(annotationCall: FirAnnotationCall, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(annotationCall, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(annotationCall, data, it) }
     }
 
-    override fun visitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {
+    override fun visitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression, data: CheckerContext) {
-        runCheck { CallCheckers.QUALIFIED_ACCESS.check(qualifiedAccessExpression, it) }
+        runCheck { CallCheckers.QUALIFIED_ACCESS.check(qualifiedAccessExpression, data, it) }
     }
 
-    override fun visitFunctionCall(functionCall: FirFunctionCall) {
+    override fun visitFunctionCall(functionCall: FirFunctionCall, data: CheckerContext) {
-        runCheck { CallCheckers.FUNCTION_CALLS.check(functionCall, it) }
+        runCheck { CallCheckers.FUNCTION_CALLS.check(functionCall, data, it) }
     }
 
-    override fun visitCallableReferenceAccess(callableReferenceAccess: FirCallableReferenceAccess) {
+    override fun visitCallableReferenceAccess(callableReferenceAccess: FirCallableReferenceAccess, data: CheckerContext) {
-        runCheck { CallCheckers.QUALIFIED_ACCESS.check(callableReferenceAccess, it) }
+        runCheck { CallCheckers.QUALIFIED_ACCESS.check(callableReferenceAccess, data, it) }
     }
 
-    override fun visitThisReceiverExpression(thisReceiverExpression: FirThisReceiverExpression) {
+    override fun visitThisReceiverExpression(thisReceiverExpression: FirThisReceiverExpression, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(thisReceiverExpression, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(thisReceiverExpression, data, it) }
     }
 
-    override fun visitResolvedQualifier(resolvedQualifier: FirResolvedQualifier) {
+    override fun visitResolvedQualifier(resolvedQualifier: FirResolvedQualifier, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(resolvedQualifier, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(resolvedQualifier, data, it) }
     }
 
-    override fun visitWhenExpression(whenExpression: FirWhenExpression) {
+    override fun visitWhenExpression(whenExpression: FirWhenExpression, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(whenExpression, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(whenExpression, data, it) }
     }
 
-    override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression) {
+    override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(binaryLogicExpression, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(binaryLogicExpression, data, it) }
     }
 
-    override fun visitArrayOfCall(arrayOfCall: FirArrayOfCall) {
+    override fun visitArrayOfCall(arrayOfCall: FirArrayOfCall, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(arrayOfCall, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(arrayOfCall, data, it) }
     }
 
-    override fun visitStringConcatenationCall(stringConcatenationCall: FirStringConcatenationCall) {
+    override fun visitStringConcatenationCall(stringConcatenationCall: FirStringConcatenationCall, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(stringConcatenationCall, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(stringConcatenationCall, data, it) }
     }
 
-    override fun visitCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall) {
+    override fun visitCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(checkNotNullCall, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(checkNotNullCall, data, it) }
     }
 
-    override fun visitTryExpression(tryExpression: FirTryExpression) {
+    override fun visitTryExpression(tryExpression: FirTryExpression, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(tryExpression, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(tryExpression, data, it) }
     }
 
-    override fun visitClassReferenceExpression(classReferenceExpression: FirClassReferenceExpression) {
+    override fun visitClassReferenceExpression(classReferenceExpression: FirClassReferenceExpression, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(classReferenceExpression, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(classReferenceExpression, data, it) }
     }
 
-    override fun visitGetClassCall(getClassCall: FirGetClassCall) {
+    override fun visitGetClassCall(getClassCall: FirGetClassCall, data: CheckerContext) {
-        runCheck { CallCheckers.EXPRESSIONS.check(getClassCall, it) }
+        runCheck { CallCheckers.EXPRESSIONS.check(getClassCall, data, it) }
     }
 
-    private fun <E : FirExpression> List<FirExpressionChecker<E>>.check(expression: E, reporter: DiagnosticReporter) {
+    private fun <E : FirExpression> List<FirExpressionChecker<E>>.check(expression: E, context: CheckerContext, reporter: DiagnosticReporter) {
         for (checker in this) {
-            checker.check(expression, reporter)
+            checker.check(expression, context, reporter)
         }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt

@@ -9,11 +9,8 @@ import kotlinx.collections.immutable.PersistentList
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSymbolOwner
-import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.resolve.calls.*
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
 import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
@@ -22,19 +19,13 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 
-typealias FirLocalScopes = PersistentList<FirLocalScope>
-
-class FirLocalContext(
-    val localScopes: FirLocalScopes,
-    val implicitReceiverStack: MutableImplicitReceiverStack
-)
-
-typealias LocalContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol, FirLocalContext>
-
 interface SessionHolder {
     val session: FirSession
+    val scopeSession: ScopeSession
 }
 
 interface BodyResolveComponents : SessionHolder {
@@ -49,7 +40,6 @@ interface BodyResolveComponents : SessionHolder {
     val container: FirDeclaration
     val inferenceComponents: InferenceComponents
     val resolutionStageRunner: ResolutionStageRunner
-    val scopeSession: ScopeSession
     val samResolver: FirSamResolver
     val callResolver: FirCallResolver
     val callCompleter: FirCallCompleter
@@ -64,4 +54,59 @@ interface BodyResolveComponents : SessionHolder {
 
     fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction)
     fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction)
-}
+}
+
+typealias FirLocalScopes = PersistentList<FirLocalScope>
+
+class FirLocalContext(
+    val localScopes: FirLocalScopes,
+    val implicitReceiverStack: MutableImplicitReceiverStack
+)
+
+typealias LocalContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol, FirLocalContext>
+
+// --------------------------------------- Utils ---------------------------------------
+
+data class ImplicitReceivers(
+    val implicitReceiverValue: ImplicitReceiverValue<*>?,
+    val implicitCompanionValues: List<ImplicitReceiverValue<*>>
+)
+
+fun SessionHolder.collectImplicitReceivers(
+    type: ConeKotlinType?,
+    owner: FirDeclaration
+): ImplicitReceivers {
+    if (type == null) return ImplicitReceivers(null, emptyList())
+
+    val implicitCompanionValues = mutableListOf<ImplicitReceiverValue<*>>()
+    val implicitReceiverValue = when (owner) {
+        is FirClass<*> -> {
+            // Questionable: performance
+            (owner as? FirRegularClass)?.companionObject?.let { companion ->
+                implicitCompanionValues += ImplicitDispatchReceiverValue(
+                    companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION
+                )
+            }
+            lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session).mapNotNull {
+                val superClass = (it as? ConeClassLikeType)?.lookupTag?.toSymbol(session)?.fir as? FirRegularClass
+                superClass?.companionObject?.let { companion ->
+                    implicitCompanionValues += ImplicitDispatchReceiverValue(
+                        companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION_FROM_SUPERTYPE
+                    )
+                }
+            }
+            // ---
+            ImplicitDispatchReceiverValue(owner.symbol, type, session, scopeSession)
+        }
+        is FirFunction<*> -> {
+            ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+        }
+        is FirVariable<*> -> {
+            ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+        }
+        else -> {
+            throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
+        }
+    }
+    return ImplicitReceivers(implicitReceiverValue, implicitCompanionValues.asReversed())
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildUnitExpression
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.inference.FirDelegatedPropertyInferenceSession
 import org.jetbrains.kotlin.fir.resolve.inference.extractLambdaInfoFromFunctionalType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -621,53 +621,19 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
         type: ConeKotlinType?,
         block: () -> T
     ): T {
-        val implicitCompanionValues: List<ImplicitReceiverValue<*>>
+        val (implicitReceiverValue, implicitCompanionValues) = components.collectImplicitReceivers(type, owner)
-        if (type != null) {
+        implicitCompanionValues.forEach { value ->
-            implicitCompanionValues = mutableListOf()
+            implicitReceiverStack.add(null, value)
-            val implicitReceiverValue = when (owner) {
-                is FirClass<*> -> {
-                    // Questionable: performance
-                    (owner as? FirRegularClass)?.companionObject?.let { companion ->
-                        implicitCompanionValues += ImplicitDispatchReceiverValue(
-                            companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION
-                        )
-                    }
-                    lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session).mapNotNull {
-                        val superClass = (it as? ConeClassLikeType)?.lookupTag?.toSymbol(session)?.fir as? FirRegularClass
-                        superClass?.companionObject?.let { companion ->
-                            implicitCompanionValues += ImplicitDispatchReceiverValue(
-                                companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION_FROM_SUPERTYPE
-                            )
-                        }
-                    }
-                    // ---
-                    ImplicitDispatchReceiverValue(owner.symbol, type, session, scopeSession)
-                }
-                is FirFunction<*> -> {
-                    ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
-                }
-                is FirVariable<*> -> {
-                    ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
-                }
-                else -> {
-                    throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
-                }
-            }
-            for (implicitCompanionValue in implicitCompanionValues.asReversed()) {
-                implicitReceiverStack.add(null, implicitCompanionValue)
-            }
-            implicitReceiverStack.add(labelName, implicitReceiverValue)
-        } else {
-            implicitCompanionValues = emptyList()
         }
+        implicitReceiverValue?.let { implicitReceiverStack.add(labelName, it) }
+
         try {
             return block()
         } finally {
             if (type != null) {
                 implicitReceiverStack.pop(labelName)
-                for (implicitCompanionValue in implicitCompanionValues) {
+                for (i in implicitCompanionValues.indices)
                     implicitReceiverStack.pop(null)
-                }
             }
         }
     }

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt

@@ -78,11 +78,13 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
     }
 
     protected fun checkDiagnostics(file: File, testFiles: List<TestFile>, firFiles: List<FirFile>) {
-        val collector = createCollector()
+        val collectors = mutableMapOf<FirSession, AbstractDiagnosticCollector>()
         val actualText = StringBuilder()
         for (testFile in testFiles) {
             val firFile = firFiles.firstOrNull { it.psi == testFile.ktFile }
             if (firFile != null) {
+                val session = firFile.session
+                val collector = collectors.computeIfAbsent(session) { createCollector(session) }
                 val coneDiagnostics = collector.collectDiagnostics(firFile)
                 testFile.getActualText(coneDiagnostics, actualText)
             } else {
@@ -92,8 +94,8 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
         KotlinTestUtils.assertEqualsToFile(file, actualText.toString())
     }
 
-    protected fun createCollector(): AbstractDiagnosticCollector {
+    protected fun createCollector(session: FirSession): AbstractDiagnosticCollector {
-        return FirDiagnosticsCollector.create()
+        return FirDiagnosticsCollector.create(session)
     }
 
     private fun checkCfg(testDataFile: File, firFiles: List<FirFile>) {

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeDiagnosticsCollector.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir
 
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.collectors.registerAllComponents
 import org.jetbrains.kotlin.fir.analysis.diagnostics.ConeDiagnostic
@@ -12,7 +13,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.KtElement
 
-class FirIdeDiagnosticsCollector(private val resolveState: FirModuleResolveState) : AbstractDiagnosticCollector() {
+class FirIdeDiagnosticsCollector(session: FirSession, private val resolveState: FirModuleResolveState) : AbstractDiagnosticCollector(session) {
 
     init {
         registerAllComponents()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt

@@ -157,7 +157,7 @@ fun KtFile.getOrBuildFirWithDiagnostics(state: FirModuleResolveState): FirFile {
     ProgressIndicatorProvider.checkCanceled() // ???
     if (state.hasDiagnosticsForFile(this)) return firFile
 
-    FirIdeDiagnosticsCollector(state).collectDiagnostics(firFile)
+    FirIdeDiagnosticsCollector(firFile.session, state).collectDiagnostics(firFile)
     state.setDiagnosticsForFile(this, firFile)
     return firFile
 }