FIR: introduce MutableCheckerContext to save performance

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

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2010-2021 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 org.jetbrains.kotlin.fir.resolve.SessionHolder
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+
+abstract class AbstractCheckerContext(
+    override val sessionHolder: SessionHolder,
+    override val returnTypeCalculator: ReturnTypeCalculator,
+    override val allInfosSuppressed: Boolean,
+    override val allWarningsSuppressed: Boolean,
+    override val allErrorsSuppressed: Boolean
+) : CheckerContext()

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

@@ -14,7 +14,9 @@ import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.name.Name
 
 abstract class CheckerContext {
     // Services
@@ -41,7 +43,21 @@ abstract class CheckerContext {
         allInfosSuppressed: Boolean,
         allWarningsSuppressed: Boolean,
         allErrorsSuppressed: Boolean
-    ): PersistentCheckerContext
+    ): CheckerContext
+
+    abstract fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): CheckerContext
+
+    abstract fun addDeclaration(declaration: FirDeclaration): CheckerContext
+
+    abstract fun dropDeclaration()
+
+    abstract fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): CheckerContext
+
+    abstract fun dropQualifiedAccessOrAnnotationCall()
+
+    abstract fun addGetClassCall(getClassCall: FirGetClassCall): CheckerContext
+
+    abstract fun dropGetClassCall()
 
     fun isDiagnosticSuppressed(diagnostic: FirDiagnostic): Boolean {
         val factory = diagnostic.factory

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

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2010-2021 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.PersistentSet
+import kotlinx.collections.immutable.persistentSetOf
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.expressions.FirGetClassCall
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
+import org.jetbrains.kotlin.fir.resolve.SessionHolder
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.name.Name
+
+class MutableCheckerContext private constructor(
+    override val implicitReceiverStack: PersistentImplicitReceiverStack,
+    override val containingDeclarations: MutableList<FirDeclaration>,
+    override val qualifiedAccessOrAnnotationCalls: MutableList<FirStatement>,
+    override val getClassCalls: MutableList<FirGetClassCall>,
+    sessionHolder: SessionHolder,
+    returnTypeCalculator: ReturnTypeCalculator,
+    override val suppressedDiagnostics: PersistentSet<String>,
+    allInfosSuppressed: Boolean,
+    allWarningsSuppressed: Boolean,
+    allErrorsSuppressed: Boolean
+) : AbstractCheckerContext(sessionHolder, returnTypeCalculator, allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed) {
+    constructor(sessionHolder: SessionHolder, returnTypeCalculator: ReturnTypeCalculator) : this(
+        PersistentImplicitReceiverStack(),
+        mutableListOf(),
+        mutableListOf(),
+        mutableListOf(),
+        sessionHolder,
+        returnTypeCalculator,
+        persistentSetOf(),
+        allInfosSuppressed = false,
+        allWarningsSuppressed = false,
+        allErrorsSuppressed = false
+    )
+
+    override fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): MutableCheckerContext {
+        return MutableCheckerContext(
+            implicitReceiverStack.add(name, value),
+            containingDeclarations,
+            qualifiedAccessOrAnnotationCalls,
+            getClassCalls,
+            sessionHolder,
+            returnTypeCalculator,
+            suppressedDiagnostics,
+            allInfosSuppressed,
+            allWarningsSuppressed,
+            allErrorsSuppressed
+        )
+    }
+
+    override fun addDeclaration(declaration: FirDeclaration): MutableCheckerContext {
+        containingDeclarations.add(declaration)
+        return this
+    }
+
+    override fun dropDeclaration() {
+        containingDeclarations.removeAt(containingDeclarations.size - 1)
+    }
+
+    override fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): MutableCheckerContext {
+        qualifiedAccessOrAnnotationCalls.add(qualifiedAccessOrAnnotationCall)
+        return this
+    }
+
+    override fun dropQualifiedAccessOrAnnotationCall() {
+        qualifiedAccessOrAnnotationCalls.removeAt(qualifiedAccessOrAnnotationCalls.size - 1)
+    }
+
+    override fun addGetClassCall(getClassCall: FirGetClassCall): MutableCheckerContext {
+        getClassCalls.add(getClassCall)
+        return this
+    }
+
+    override fun dropGetClassCall() {
+        getClassCalls.removeAt(getClassCalls.size - 1)
+    }
+
+    override fun addSuppressedDiagnostics(
+        diagnosticNames: Collection<String>,
+        allInfosSuppressed: Boolean,
+        allWarningsSuppressed: Boolean,
+        allErrorsSuppressed: Boolean
+    ): MutableCheckerContext {
+        if (diagnosticNames.isEmpty()) return this
+        return MutableCheckerContext(
+            implicitReceiverStack,
+            containingDeclarations,
+            qualifiedAccessOrAnnotationCalls,
+            getClassCalls,
+            sessionHolder,
+            returnTypeCalculator,
+            suppressedDiagnostics.addAll(diagnosticNames),
+            this.allInfosSuppressed || allInfosSuppressed,
+            this.allWarningsSuppressed || allWarningsSuppressed,
+            this.allErrorsSuppressed || allErrorsSuppressed
+        )
+    }
+}

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

@@ -23,13 +23,13 @@ class PersistentCheckerContext private constructor(
     override val containingDeclarations: PersistentList<FirDeclaration>,
     override val qualifiedAccessOrAnnotationCalls: PersistentList<FirStatement>,
     override val getClassCalls: PersistentList<FirGetClassCall>,
-    override val sessionHolder: SessionHolder,
-    override val returnTypeCalculator: ReturnTypeCalculator,
+    sessionHolder: SessionHolder,
+    returnTypeCalculator: ReturnTypeCalculator,
     override val suppressedDiagnostics: PersistentSet<String>,
-    override val allInfosSuppressed: Boolean,
-    override val allWarningsSuppressed: Boolean,
-    override val allErrorsSuppressed: Boolean
-) : CheckerContext() {
+    allInfosSuppressed: Boolean,
+    allWarningsSuppressed: Boolean,
+    allErrorsSuppressed: Boolean
+) : AbstractCheckerContext(sessionHolder, returnTypeCalculator, allInfosSuppressed, allWarningsSuppressed, allErrorsSuppressed) {
     constructor(sessionHolder: SessionHolder, returnTypeCalculator: ReturnTypeCalculator) : this(
         PersistentImplicitReceiverStack(),
         persistentListOf(),
@@ -43,7 +43,7 @@ class PersistentCheckerContext private constructor(
         allErrorsSuppressed = false
     )
 
-    fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext {
+    override fun addImplicitReceiver(name: Name?, value: ImplicitReceiverValue<*>): PersistentCheckerContext {
         return PersistentCheckerContext(
             implicitReceiverStack.add(name, value),
             containingDeclarations,
@@ -58,7 +58,7 @@ class PersistentCheckerContext private constructor(
         )
     }
 
-    fun addDeclaration(declaration: FirDeclaration): PersistentCheckerContext {
+    override fun addDeclaration(declaration: FirDeclaration): PersistentCheckerContext {
         return PersistentCheckerContext(
             implicitReceiverStack,
             containingDeclarations.add(declaration),
@@ -73,7 +73,10 @@ class PersistentCheckerContext private constructor(
         )
     }
 
-    fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): PersistentCheckerContext {
+    override fun dropDeclaration() {
+    }
+
+    override fun addQualifiedAccessOrAnnotationCall(qualifiedAccessOrAnnotationCall: FirStatement): PersistentCheckerContext {
         return PersistentCheckerContext(
             implicitReceiverStack,
             containingDeclarations,
@@ -88,7 +91,10 @@ class PersistentCheckerContext private constructor(
         )
     }
 
-    fun addGetClassCall(getClassCall: FirGetClassCall): PersistentCheckerContext {
+    override fun dropQualifiedAccessOrAnnotationCall() {
+    }
+
+    override fun addGetClassCall(getClassCall: FirGetClassCall): PersistentCheckerContext {
         return PersistentCheckerContext(
             implicitReceiverStack,
             containingDeclarations,
@@ -103,6 +109,9 @@ class PersistentCheckerContext private constructor(
         )
     }
 
+    override fun dropGetClassCall() {
+    }
+
     override fun addSuppressedDiagnostics(
         diagnosticNames: Collection<String>,
         allInfosSuppressed: Boolean,

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/util/firCheckersUtils.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.analysis.checkers.util
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorVisitor
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
@@ -15,7 +14,7 @@ fun FirElement.checkChildrenWithCustomVisitor(
     parentContext: CheckerContext,
     visitorVoid: FirVisitor<Unit, CheckerContext>
 ) {
-    val collectingVisitor = object : AbstractDiagnosticCollectorVisitor(parentContext as PersistentCheckerContext) {
+    val collectingVisitor = object : AbstractDiagnosticCollectorVisitor(parentContext) {
         override fun checkElement(element: FirElement) {
             element.accept(visitorVoid, context)
         }

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

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.PrivateForInline
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.*
@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.name.Name
 
 abstract class AbstractDiagnosticCollectorVisitor(
-    @set:PrivateForInline var context: PersistentCheckerContext,
+    @set:PrivateForInline var context: CheckerContext,
 ) : FirDefaultVisitor<Unit, Nothing?>() {
 
     protected open fun shouldVisitDeclaration(declaration: FirDeclaration) = true
@@ -256,6 +256,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
         try {
             return block()
         } finally {
+            existingContext.dropQualifiedAccessOrAnnotationCall()
             context = existingContext
         }
     }
@@ -268,6 +269,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
         try {
             return block()
         } finally {
+            existingContext.dropGetClassCall()
             context = existingContext
         }
     }
@@ -280,6 +282,7 @@ abstract class AbstractDiagnosticCollectorVisitor(
         try {
             return block()
         } finally {
+            existingContext.dropDeclaration()
             context = existingContext
         }
     }

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

@@ -6,11 +6,11 @@
 package org.jetbrains.kotlin.fir.analysis.collectors
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
 
 open class CheckerRunningDiagnosticCollectorVisitor(
-    context: PersistentCheckerContext,
+    context: CheckerContext,
     protected val components: List<AbstractDiagnosticCollectorComponent>
 ) : AbstractDiagnosticCollectorVisitor(context) {
 

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

@@ -6,15 +6,10 @@
 package org.jetbrains.kotlin.fir.analysis.collectors
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.MutableCheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
-import org.jetbrains.kotlin.fir.analysis.diagnostics.impl.BaseDiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.impl.DiagnosticReporterWithSuppress
-import org.jetbrains.kotlin.fir.analysis.diagnostics.impl.SimpleDiagnosticReporter
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 
 class SimpleDiagnosticsCollector(
@@ -24,7 +19,7 @@ class SimpleDiagnosticsCollector(
 ) : AbstractDiagnosticCollector(session, scopeSession, createComponents) {
     override fun createVisitor(components: List<AbstractDiagnosticCollectorComponent>): CheckerRunningDiagnosticCollectorVisitor {
         return CheckerRunningDiagnosticCollectorVisitor(
-            PersistentCheckerContext(
+            MutableCheckerContext(
                 this,
                 ReturnTypeCalculatorForFullBodyResolve()
             ),

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/BeforeElementDiagnosticCollectionHandler.kt

@@ -8,12 +8,12 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 
 abstract class BeforeElementDiagnosticCollectionHandler: FirSessionComponent {
     open fun beforeCollectingForElement(element: FirElement) {}
-    open fun beforeGoingNestedDeclaration(declaration: FirDeclaration, context: PersistentCheckerContext) {}
+    open fun beforeGoingNestedDeclaration(declaration: FirDeclaration, context: CheckerContext) {}
 }
 
 val FirSession.beforeElementDiagnosticCollectionHandler: BeforeElementDiagnosticCollectionHandler? by FirSession.nullableSessionComponentAccessor()

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FileStructureElementDiagnosticRetriever.kt

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
@@ -42,7 +42,7 @@ internal class SingleNonLocalDeclarationDiagnosticRetriever(
 
     private class Visitor(
         private val structureElementDeclaration: FirDeclaration,
-        context: PersistentCheckerContext,
+        context: CheckerContext,
         components: List<AbstractDiagnosticCollectorComponent>
     ) : FirIdeDiagnosticVisitor(context, components) {
         private var insideAlwaysVisitableDeclarations = 0

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeDiagnosticVisitor.kt

@@ -6,14 +6,14 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.CheckerRunningDiagnosticCollectorVisitor
 import org.jetbrains.kotlin.fir.analysis.collectors.components.AbstractDiagnosticCollectorComponent
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 
 internal open class FirIdeDiagnosticVisitor(
-    context: PersistentCheckerContext,
+    context: CheckerContext,
     components: List<AbstractDiagnosticCollectorComponent>,
 ) : CheckerRunningDiagnosticCollectorVisitor(context, components) {
     private val beforeElementDiagnosticCollectionHandler = context.session.beforeElementDiagnosticCollectionHandler

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/fir/ContextCollectingDiagnosticCollectorVisitor.kt

@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorVisitor
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
@@ -24,8 +24,8 @@ private class ContextCollectingDiagnosticCollectorVisitor private constructor(
 ) : AbstractDiagnosticCollectorVisitor(
     PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(sessionHolder)
 ) {
-    private val contextCollector = object : ContextByDesignationCollector<PersistentCheckerContext>(designation) {
-        override fun getCurrentContext(): PersistentCheckerContext = context
+    private val contextCollector = object : ContextByDesignationCollector<CheckerContext>(designation) {
+        override fun getCurrentContext(): CheckerContext = context
 
         override fun goToNestedDeclaration(declaration: FirDeclaration) {
             declaration.accept(this@ContextCollectingDiagnosticCollectorVisitor, null)
@@ -43,7 +43,7 @@ private class ContextCollectingDiagnosticCollectorVisitor private constructor(
     override fun checkElement(element: FirElement) {}
 
     companion object {
-        fun collect(sessionHolder: SessionHolder, designation: FirDeclarationDesignationWithFile): PersistentCheckerContext {
+        fun collect(sessionHolder: SessionHolder, designation: FirDeclarationDesignationWithFile): CheckerContext {
             val visitor = ContextCollectingDiagnosticCollectorVisitor(sessionHolder, designation)
             designation.firFile.accept(visitor, null)
             return visitor.contextCollector.getCollectedContext()
@@ -56,7 +56,7 @@ internal object PersistenceContextCollector {
         sessionHolder: SessionHolder,
         firFile: FirFile,
         declaration: FirDeclaration,
-    ): PersistentCheckerContext {
+    ): CheckerContext {
         val isLocal = when (declaration) {
             is FirClassLikeDeclaration -> declaration.symbol.classId.isLocal
             is FirCallableDeclaration -> declaration.symbol.callableId.isLocal

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/AbstractFirContextCollectionTest.kt

@@ -5,12 +5,11 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostic
 
-import kotlinx.collections.immutable.PersistentList
 import org.jetbrains.kotlin.fir.SessionConfiguration
-import org.jetbrains.kotlin.fir.analysis.checkers.context.PersistentCheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
+import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
@@ -68,7 +67,7 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
         lateinit var elementsToCheckContext: List<FirDeclaration>
         lateinit var firFile: FirFile
 
-        override fun beforeGoingNestedDeclaration(declaration: FirDeclaration, context: PersistentCheckerContext) {
+        override fun beforeGoingNestedDeclaration(declaration: FirDeclaration, context: CheckerContext) {
             if (declaration is FirFile) {
                 return
             }
@@ -82,15 +81,15 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
             }
         }
 
-        private fun compareStructurally(expected: PersistentCheckerContext, actual: PersistentCheckerContext) {
+        private fun compareStructurally(expected: CheckerContext, actual: CheckerContext) {
             assertions.assertEquals(expected.implicitReceiverStack.asString(), actual.implicitReceiverStack.asString())
             assertions.assertEquals(expected.containingDeclarations.asString(), actual.containingDeclarations.asString())
         }
 
-        private fun PersistentImplicitReceiverStack.asString() =
+        private fun ImplicitReceiverStack.asString() =
             joinToString { it.boundSymbol.name() }
 
-        private fun PersistentList<FirDeclaration>.asString() =
+        private fun List<FirDeclaration>.asString() =
             joinToString(transform = FirDeclaration::name)
     }
 }