[analysis api] fix Analysis API contracts violation

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.*
 import org.jetbrains.kotlin.analysis.api.impl.base.components.KtAnalysisScopeProviderImpl
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolProvider
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.psi.KtElement
@@ -56,7 +57,8 @@ class KtFe10AnalysisSession(
     override val analysisScopeProviderImpl: KtAnalysisScopeProvider = KtAnalysisScopeProviderImpl(this, token)
     override val referenceResolveProviderImpl: KtReferenceResolveProvider = KtFe10ReferenceResolveProvider(this)
 
-    override fun createContextDependentCopy(originalKtFile: KtFile, elementToReanalyze: KtElement): KtAnalysisSession {
+    override fun createContextDependentCopy(originalKtFile: KtFile, elementToReanalyze: KtElement): KtAnalysisSession =
-        return KtFe10AnalysisSession(elementToReanalyze, token)
+        withValidityAssertion {
-    }
+            KtFe10AnalysisSession(elementToReanalyze, token)
+        }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10CompletionCandidateChecker.kt

@@ -27,7 +27,7 @@ internal class KtFe10CompletionCandidateChecker(
         originalFile: KtFile,
         nameExpression: KtSimpleNameExpression,
         possibleExplicitReceiver: KtExpression?
-    ): KtExtensionApplicabilityResult = withValidityAssertion {
+    ): KtExtensionApplicabilityResult {
         throw NotImplementedError("Method is not implemented for FE 1.0")
     }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10DiagnosticProvider.kt

@@ -31,36 +31,32 @@ internal class KtFe10DiagnosticProvider(
         get() = analysisSession.token
 
     override fun getDiagnosticsForElement(element: KtElement, filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>> {
-        withValidityAssertion {
+        val bindingContext = analysisContext.analyze(element, AnalysisMode.PARTIAL_WITH_DIAGNOSTICS)
-            val bindingContext = analysisContext.analyze(element, AnalysisMode.PARTIAL_WITH_DIAGNOSTICS)
+        val diagnostics = bindingContext.diagnostics.forElement(element)
-            val diagnostics = bindingContext.diagnostics.forElement(element)
+        return diagnostics.map { KtFe10Diagnostic(it, token) }
-            return diagnostics.map { KtFe10Diagnostic(it, token) }
-        }
     }
 
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>> {
-        withValidityAssertion {
+        val bindingContext = analysisContext.analyze(ktFile)
-            val bindingContext = analysisContext.analyze(ktFile)
+        val result = mutableListOf<KtDiagnosticWithPsi<*>>()
-            val result = mutableListOf<KtDiagnosticWithPsi<*>>()
+        for (diagnostic in bindingContext.diagnostics) {
-            for (diagnostic in bindingContext.diagnostics) {
+            if (diagnostic.psiFile == ktFile) {
-                if (diagnostic.psiFile == ktFile) {
+                result += KtFe10Diagnostic(diagnostic, token)
-                    result += KtFe10Diagnostic(diagnostic, token)
-                }
             }
-            return result
         }
+        return result
     }
 }
 
 internal class KtFe10Diagnostic(private val diagnostic: Diagnostic, override val token: KtLifetimeToken) : KtDiagnosticWithPsi<PsiElement> {
     override val severity: Severity
-        get() = diagnostic.severity
+        get() = withValidityAssertion { diagnostic.severity }
 
     override val factoryName: String
-        get() = diagnostic.factory.name
+        get() = withValidityAssertion { diagnostic.factory.name }
 
     override val defaultMessage: String
-        get() {
+        get() = withValidityAssertion {
             @Suppress("UNCHECKED_CAST")
             val factory = diagnostic.factory as DiagnosticFactory<UnboundDiagnostic>?
             return factory?.defaultRenderer?.render(diagnostic)
@@ -69,11 +65,11 @@ internal class KtFe10Diagnostic(private val diagnostic: Diagnostic, override val
         }
 
     override val psi: PsiElement
-        get() = diagnostic.psiElement
+        get() = withValidityAssertion { diagnostic.psiElement }
 
     override val textRanges: Collection<TextRange>
-        get() = diagnostic.textRanges
+        get() = withValidityAssertion { diagnostic.textRanges }
 
     override val diagnosticClass: KClass<out KtDiagnosticWithPsi<PsiElement>>
-        get() = KtFe10Diagnostic::class
+        get() = withValidityAssertion { KtFe10Diagnostic::class }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10ExpressionInfoProvider.kt

@@ -27,7 +27,7 @@ internal class KtFe10ExpressionInfoProvider(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtCallableSymbol? = withValidityAssertion {
+    override fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtCallableSymbol? {
         val bindingContext = analysisContext.analyze(returnExpression, AnalysisMode.PARTIAL)
         val targetLabel = returnExpression.getTargetLabel()
             ?: return returnExpression.parentOfType<KtNamedFunction>()
@@ -36,7 +36,7 @@ internal class KtFe10ExpressionInfoProvider(
         return with(analysisSession) { labelTarget.getSymbol() as? KtCallableSymbol }
     }
 
-    override fun getWhenMissingCases(whenExpression: KtWhenExpression): List<WhenMissingCase> = withValidityAssertion {
+    override fun getWhenMissingCases(whenExpression: KtWhenExpression): List<WhenMissingCase>  {
         val bindingContext = analysisContext.analyze(whenExpression)
         return WhenChecker.getMissingCases(whenExpression, bindingContext)
     }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10ExpressionTypeProvider.kt

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnaly
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtType
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -28,10 +27,10 @@ import org.jetbrains.kotlin.resolve.calls.util.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.util.getType
 import org.jetbrains.kotlin.resolve.sam.SamConstructorDescriptor
 import org.jetbrains.kotlin.resolve.sam.getFunctionTypeForAbstractMethod
-import org.jetbrains.kotlin.types.error.ErrorTypeKind
-import org.jetbrains.kotlin.types.error.ErrorUtils
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.checker.intersectWrappedTypes
+import org.jetbrains.kotlin.types.error.ErrorTypeKind
+import org.jetbrains.kotlin.types.error.ErrorUtils
 import org.jetbrains.kotlin.types.typeUtil.makeNullable
 
 class KtFe10ExpressionTypeProvider(
@@ -49,7 +48,7 @@ class KtFe10ExpressionTypeProvider(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun getKtExpressionType(expression: KtExpression): KtType? = withValidityAssertion {
+    override fun getKtExpressionType(expression: KtExpression): KtType? {
         // Not sure if it's safe enough. In theory, some annotations on expressions might change its type
         val unwrapped = expression.unwrapParenthesesLabelsAndAnnotations() as? KtExpression ?: return null
         if (unwrapped.getParentOfTypes(false, *NON_EXPRESSION_CONTAINERS) != null) {
@@ -66,7 +65,7 @@ class KtFe10ExpressionTypeProvider(
         return kotlinType.toKtType(analysisContext)
     }
 
-    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType = withValidityAssertion {
+    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType {
         // Handle callable declarations with explicit return type first
         if (declaration is KtCallableDeclaration) {
             val typeReference = declaration.typeReference
@@ -134,7 +133,7 @@ class KtFe10ExpressionTypeProvider(
         return analysisContext.builtIns.unitType.toKtType(analysisContext)
     }
 
-    override fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType = withValidityAssertion {
+    override fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType {
         val analysisMode = if (declaration.hasDeclaredReturnType()) AnalysisMode.PARTIAL else AnalysisMode.FULL
         val bindingContext = analysisContext.analyze(declaration, analysisMode)
         val functionDescriptor = bindingContext[BindingContext.FUNCTION, declaration]
@@ -155,7 +154,7 @@ class KtFe10ExpressionTypeProvider(
             .toKtType(analysisContext)
     }
 
-    override fun getExpectedType(expression: PsiElement): KtType? = withValidityAssertion {
+    override fun getExpectedType(expression: PsiElement): KtType? {
         val ktExpression = expression.getParentOfType<KtExpression>(false) ?: return null
         val parentExpression = ktExpression.parent
 
@@ -222,7 +221,7 @@ class KtFe10ExpressionTypeProvider(
         }
     }
 
-    override fun isDefinitelyNull(expression: KtExpression): Boolean = withValidityAssertion {
+    override fun isDefinitelyNull(expression: KtExpression): Boolean {
         val unwrapped = expression.unwrapParenthesesLabelsAndAnnotations() as? KtElement ?: return false
         val bindingContext = analysisContext.analyze(expression, AnalysisMode.PARTIAL)
 
@@ -239,7 +238,7 @@ class KtFe10ExpressionTypeProvider(
         return false
     }
 
-    override fun isDefinitelyNotNull(expression: KtExpression): Boolean = withValidityAssertion {
+    override fun isDefinitelyNotNull(expression: KtExpression): Boolean {
         val ktExpression = expression as? KtExpression ?: return false
         val bindingContext = analysisContext.analyze(ktExpression)
 

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10InheritorsProvider.kt

@@ -36,7 +36,7 @@ internal class KtFe10InheritorsProvider(
             .mapNotNull { it.toKtClassifierSymbol(analysisContext) as? KtNamedClassOrObjectSymbol }
     }
 
-    override fun getEnumEntries(classSymbol: KtNamedClassOrObjectSymbol): List<KtEnumEntrySymbol> = withValidityAssertion {
+    override fun getEnumEntries(classSymbol: KtNamedClassOrObjectSymbol): List<KtEnumEntrySymbol> {
         val enumDescriptor = getSymbolDescriptor(classSymbol) as? ClassDescriptor ?: return emptyList()
         if (enumDescriptor.kind != ClassKind.ENUM_CLASS) {
             return emptyList()

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10JvmTypeMapper.kt

@@ -24,7 +24,7 @@ internal class KtFe10JvmTypeMapper(
 
     private val typeMapper by lazy { KtFe10JvmTypeMapperContext(analysisContext.resolveSession) }
 
-    override fun mapTypeToJvmType(type: KtType, mode: TypeMappingMode): Type = withValidityAssertion {
+    override fun mapTypeToJvmType(type: KtType, mode: TypeMappingMode): Type {
         val kotlinType = (type as KtFe10Type).type
         return typeMapper.mapType(kotlinType, mode)
     }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10OverrideInfoProvider.kt

@@ -28,7 +28,7 @@ internal class KtFe10OverrideInfoProvider(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun isVisible(memberSymbol: KtCallableSymbol, classSymbol: KtClassOrObjectSymbol): Boolean = withValidityAssertion {
+    override fun isVisible(memberSymbol: KtCallableSymbol, classSymbol: KtClassOrObjectSymbol): Boolean  {
         val memberDescriptor = getSymbolDescriptor(memberSymbol) as? DeclarationDescriptorWithVisibility ?: return false
         val classDescriptor = getSymbolDescriptor(classSymbol) ?: return false
         return isVisibleWithAnyReceiver(memberDescriptor, classDescriptor, analysisSession.analysisContext.languageVersionSettings)
@@ -37,17 +37,17 @@ internal class KtFe10OverrideInfoProvider(
     override fun getImplementationStatus(
         memberSymbol: KtCallableSymbol,
         parentClassSymbol: KtClassOrObjectSymbol
-    ): ImplementationStatus? = withValidityAssertion {
+    ): ImplementationStatus?  {
         throw NotImplementedError("Method is not implemented for FE 1.0")
     }
 
-    override fun getOriginalOverriddenSymbol(symbol: KtCallableSymbol): KtCallableSymbol? = withValidityAssertion {
+    override fun getOriginalOverriddenSymbol(symbol: KtCallableSymbol): KtCallableSymbol?  {
         val callableDescriptor = getSymbolDescriptor(symbol) as? CallableMemberDescriptor ?: return null
         val originalCallableDescriptor = callableDescriptor.findOriginalTopMostOverriddenDescriptors().firstOrNull() ?: return null
         return originalCallableDescriptor.toKtCallableSymbol(analysisContext)
     }
 
-    override fun getOriginalContainingClassForOverride(symbol: KtCallableSymbol): KtClassOrObjectSymbol? = withValidityAssertion {
+    override fun getOriginalContainingClassForOverride(symbol: KtCallableSymbol): KtClassOrObjectSymbol?  {
         val callableDescriptor = getSymbolDescriptor(symbol) as? CallableMemberDescriptor ?: return null
         val originalCallableDescriptor = callableDescriptor.findOriginalTopMostOverriddenDescriptors().firstOrNull() ?: return null
         val containingClassDescriptor = originalCallableDescriptor.containingDeclaration as? ClassDescriptor ?: return null

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10PsiTypeProvider.kt

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.analysis.api.descriptors.utils.KtFe10JvmTypeMapperCo
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeMappingMode
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.load.kotlin.getOptimalModeForReturnType
@@ -41,7 +40,7 @@ internal class KtFe10PsiTypeProvider(
         useSitePosition: PsiElement,
         mode: KtTypeMappingMode,
         isAnnotationMethod: Boolean,
-    ): PsiType? = withValidityAssertion {
+    ): PsiType? {
         val kotlinType = (type as KtFe10Type).type
 
         if (kotlinType.isError || kotlinType.arguments.any { !it.isStarProjection && it.type.isError }) {

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10SamResolver.kt

@@ -25,7 +25,7 @@ internal class KtFe10SamResolver(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun getSamConstructor(ktClassLikeSymbol: KtClassLikeSymbol): KtSamConstructorSymbol? = withValidityAssertion {
+    override fun getSamConstructor(ktClassLikeSymbol: KtClassLikeSymbol): KtSamConstructorSymbol?  {
         val descriptor = getSymbolDescriptor(ktClassLikeSymbol)
         if (descriptor is ClassDescriptor && getSingleAbstractMethodOrNull(descriptor) != null) {
             val constructorDescriptor = createSamConstructorFunction(

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10ScopeProvider.kt

@@ -11,7 +11,10 @@ import org.jetbrains.kotlin.analysis.api.components.KtScopeContext
 import org.jetbrains.kotlin.analysis.api.components.KtScopeProvider
 import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
 import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnalysisSessionComponent
-import org.jetbrains.kotlin.analysis.api.descriptors.scopes.*
+import org.jetbrains.kotlin.analysis.api.descriptors.scopes.KtFe10FileScope
+import org.jetbrains.kotlin.analysis.api.descriptors.scopes.KtFe10PackageScope
+import org.jetbrains.kotlin.analysis.api.descriptors.scopes.KtFe10ScopeLexical
+import org.jetbrains.kotlin.analysis.api.descriptors.scopes.KtFe10ScopeMember
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.KtFe10FileSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.KtFe10PackageSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.base.KtFe10Symbol
@@ -21,16 +24,16 @@ import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.bas
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.psiBased.base.KtFe10PsiSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.psiBased.base.getResolutionScope
 import org.jetbrains.kotlin.analysis.api.descriptors.types.base.KtFe10Type
-import org.jetbrains.kotlin.analysis.api.impl.base.scopes.KtEmptyScope
 import org.jetbrains.kotlin.analysis.api.impl.base.scopes.KtCompositeScope
-import org.jetbrains.kotlin.analysis.api.scopes.*
+import org.jetbrains.kotlin.analysis.api.impl.base.scopes.KtEmptyScope
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.scopes.KtScope
 import org.jetbrains.kotlin.analysis.api.symbols.KtFileSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
@@ -55,21 +58,21 @@ internal class KtFe10ScopeProvider(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun getMemberScope(classSymbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getMemberScope(classSymbol: KtSymbolWithMembers): KtScope {
         val descriptor = getDescriptor<ClassDescriptor>(classSymbol)
             ?: return getEmptyScope()
 
         return KtFe10ScopeMember(descriptor.unsubstitutedMemberScope, analysisContext)
     }
 
-    override fun getDeclaredMemberScope(classSymbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getDeclaredMemberScope(classSymbol: KtSymbolWithMembers): KtScope {
         val descriptor = getDescriptor<ClassDescriptor>(classSymbol)
             ?: return getEmptyScope()
 
         return KtFe10ScopeMember(DeclaredMemberScope(descriptor), analysisContext)
     }
 
-    override fun getDelegatedMemberScope(classSymbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getDelegatedMemberScope(classSymbol: KtSymbolWithMembers): KtScope  {
         val descriptor = getDescriptor<ClassDescriptor>(classSymbol)
             ?: return getEmptyScope()
 
@@ -151,21 +154,21 @@ internal class KtFe10ScopeProvider(
 
     }
 
-    override fun getStaticMemberScope(symbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getStaticMemberScope(symbol: KtSymbolWithMembers): KtScope {
         val descriptor = getDescriptor<ClassDescriptor>(symbol) ?: return getEmptyScope()
         return KtFe10ScopeMember(descriptor.staticScope, analysisContext)
     }
 
-    override fun getEmptyScope(): KtScope = withValidityAssertion {
+    override fun getEmptyScope(): KtScope {
-        KtEmptyScope(token)
+        return KtEmptyScope(token)
     }
 
-    override fun getFileScope(fileSymbol: KtFileSymbol): KtScope = withValidityAssertion {
+    override fun getFileScope(fileSymbol: KtFileSymbol): KtScope {
         require(fileSymbol is KtFe10FileSymbol)
         return KtFe10FileScope(fileSymbol.psi, analysisContext, token)
     }
 
-    override fun getPackageScope(packageSymbol: KtPackageSymbol): KtScope = withValidityAssertion {
+    override fun getPackageScope(packageSymbol: KtPackageSymbol): KtScope {
         require(packageSymbol is KtFe10PackageSymbol)
         val packageFragments = analysisContext.resolveSession.packageFragmentProvider.packageFragments(packageSymbol.fqName)
         val scopeDescription = "Compound scope for package \"${packageSymbol.fqName}\""
@@ -173,28 +176,28 @@ internal class KtFe10ScopeProvider(
         return KtFe10PackageScope(chainedScope, packageSymbol, analysisContext)
     }
 
-    override fun getCompositeScope(subScopes: List<KtScope>): KtScope = withValidityAssertion {
+    override fun getCompositeScope(subScopes: List<KtScope>): KtScope {
         return KtCompositeScope(subScopes, token)
     }
 
-    override fun getTypeScope(type: KtType): KtScope = withValidityAssertion {
+    override fun getTypeScope(type: KtType): KtScope {
         require(type is KtFe10Type)
         return KtFe10ScopeMember(type.type.memberScope, analysisContext)
     }
 
-    override fun getScopeContextForPosition(originalFile: KtFile, positionInFakeFile: KtElement): KtScopeContext = withValidityAssertion {
+    override fun getScopeContextForPosition(originalFile: KtFile, positionInFakeFile: KtElement): KtScopeContext {
         val elementToAnalyze = positionInFakeFile.containingNonLocalDeclaration() ?: originalFile
         val bindingContext = analysisContext.analyze(elementToAnalyze)
 
         val lexicalScope = positionInFakeFile.getResolutionScope(bindingContext)
         if (lexicalScope != null) {
             val compositeScope = KtCompositeScope(listOf(KtFe10ScopeLexical(lexicalScope, analysisContext)), token)
-            return KtScopeContext(compositeScope, collectImplicitReceivers(lexicalScope))
+            return KtScopeContext(compositeScope, collectImplicitReceivers(lexicalScope), token)
         }
 
         val fileScope = analysisContext.resolveSession.fileScopeProvider.getFileResolutionScope(originalFile)
         val compositeScope = KtCompositeScope(listOf(KtFe10ScopeLexical(fileScope, analysisContext)), token)
-        return KtScopeContext(compositeScope, collectImplicitReceivers(fileScope))
+        return KtScopeContext(compositeScope, collectImplicitReceivers(fileScope), token)
     }
 
     private inline fun <reified T : DeclarationDescriptor> getDescriptor(symbol: KtSymbol): T? {

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10SmartCastProvider.kt

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnaly
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtType
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.smartcasts.ExplicitSmartCasts
@@ -28,22 +27,20 @@ internal class KtFe10SmartCastProvider(
         get() = analysisSession.token
 
     override fun getSmartCastedInfo(expression: KtExpression): KtSmartCastInfo? {
-        withValidityAssertion {
+        val bindingContext = analysisContext.analyze(expression)
-            val bindingContext = analysisContext.analyze(expression)
+        val stableSmartCasts = bindingContext[BindingContext.SMARTCAST, expression]
-            val stableSmartCasts = bindingContext[BindingContext.SMARTCAST, expression]
+        val unstableSmartCasts = bindingContext[BindingContext.UNSTABLE_SMARTCAST, expression]
-            val unstableSmartCasts = bindingContext[BindingContext.UNSTABLE_SMARTCAST, expression]
 
-            return when {
+        return when {
-                stableSmartCasts != null -> {
+            stableSmartCasts != null -> {
-                    val type = stableSmartCasts.getKtType() ?: return null
+                val type = stableSmartCasts.getKtType() ?: return null
-                    KtSmartCastInfo(type, true, token)
+                KtSmartCastInfo(type, true, token)
-                }
-                unstableSmartCasts != null -> {
-                    val type = unstableSmartCasts.getKtType() ?: return null
-                    KtSmartCastInfo(type, false, token)
-                }
-                else -> null
             }
+            unstableSmartCasts != null -> {
+                val type = unstableSmartCasts.getKtType() ?: return null
+                KtSmartCastInfo(type, false, token)
+            }
+            else -> null
         }
     }
 
@@ -55,13 +52,11 @@ internal class KtFe10SmartCastProvider(
     }
 
     override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<KtImplicitReceiverSmartCast> {
-        withValidityAssertion {
+        val bindingContext = analysisContext.analyze(expression)
-            val bindingContext = analysisContext.analyze(expression)
+        val smartCasts = bindingContext[BindingContext.IMPLICIT_RECEIVER_SMARTCAST, expression] ?: return emptyList()
-            val smartCasts = bindingContext[BindingContext.IMPLICIT_RECEIVER_SMARTCAST, expression] ?: return emptyList()
+        return smartCasts.receiverTypes.map { (_, type) ->
-            return smartCasts.receiverTypes.map { (_, type) ->
+            val kind = KtImplicitReceiverSmartCastKind.DISPATCH // TODO provide precise kind
-                val kind = KtImplicitReceiverSmartCastKind.DISPATCH // TODO provide precise kind
+            KtImplicitReceiverSmartCast(type.toKtType(analysisContext), kind, token)
-                KtImplicitReceiverSmartCast(type.toKtType(analysisContext), kind, token)
-            }
         }
     }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10SubtypingComponent.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnaly
 import org.jetbrains.kotlin.analysis.api.descriptors.types.base.KtFe10Type
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 
 internal class KtFe10SubtypingComponent(
     override val analysisSession: KtFe10AnalysisSession
@@ -19,13 +18,13 @@ internal class KtFe10SubtypingComponent(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun isEqualTo(first: KtType, second: KtType): Boolean = withValidityAssertion {
+    override fun isEqualTo(first: KtType, second: KtType): Boolean {
         require(first is KtFe10Type)
         require(second is KtFe10Type)
         return analysisContext.resolveSession.kotlinTypeCheckerOfOwnerModule.equalTypes(first.type, second.type)
     }
 
-    override fun isSubTypeOf(subType: KtType, superType: KtType): Boolean = withValidityAssertion {
+    override fun isSubTypeOf(subType: KtType, superType: KtType): Boolean {
         require(subType is KtFe10Type)
         require(superType is KtFe10Type)
         val typeChecker = analysisContext.resolveSession.kotlinTypeCheckerOfOwnerModule

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10SymbolDeclarationOverridesProvider.kt

@@ -10,11 +10,10 @@ import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
 import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnalysisSessionComponent
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.getSymbolDescriptor
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtCallableSymbol
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassOrAny
@@ -66,19 +65,19 @@ internal class KtFe10SymbolDeclarationOverridesProvider(
         return overriddenDescriptors
     }
 
-    override fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean = withValidityAssertion {
+    override fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean {
         val subClassDescriptor = getSymbolDescriptor(subClass) as? ClassDescriptor ?: return false
         val superClassDescriptor = getSymbolDescriptor(superClass) as? ClassDescriptor ?: return false
         return subClassDescriptor.isSubclassOf(superClassDescriptor)
     }
 
-    override fun isDirectSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean = withValidityAssertion {
+    override fun isDirectSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean {
         val subClassDescriptor = getSymbolDescriptor(subClass) as? ClassDescriptor ?: return false
         val superClassDescriptor = getSymbolDescriptor(superClass) as? ClassDescriptor ?: return false
         return subClassDescriptor.getSuperClassOrAny() == superClassDescriptor
     }
 
-    override fun getIntersectionOverriddenSymbols(symbol: KtCallableSymbol): Collection<KtCallableSymbol> = withValidityAssertion {
+    override fun getIntersectionOverriddenSymbols(symbol: KtCallableSymbol): Collection<KtCallableSymbol> {
         throw NotImplementedError("Method is not implemented for FE 1.0")
     }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10SymbolProvider.kt

@@ -12,9 +12,8 @@ import org.jetbrains.kotlin.analysis.api.descriptors.symbols.KtFe10PackageSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtClassSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.psiBased.*
-import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -31,11 +30,11 @@ internal class KtFe10SymbolProvider(
     override val ROOT_PACKAGE_SYMBOL: KtPackageSymbol
         get() = KtFe10PackageSymbol(FqName.ROOT, analysisContext)
 
-    override fun getFileSymbol(psi: KtFile): KtFileSymbol = withValidityAssertion {
+    override fun getFileSymbol(psi: KtFile): KtFileSymbol {
         return KtFe10FileSymbol(psi, analysisContext)
     }
 
-    override fun getParameterSymbol(psi: KtParameter): KtVariableLikeSymbol = withValidityAssertion {
+    override fun getParameterSymbol(psi: KtParameter): KtVariableLikeSymbol {
         return when {
             psi.isFunctionTypeParameter -> error("Function type parameters are not supported in getParameterSymbol()")
             psi.isLoopParameter -> KtFe10PsiLoopParameterLocalVariableSymbol(psi, analysisContext)
@@ -43,7 +42,7 @@ internal class KtFe10SymbolProvider(
         }
     }
 
-    override fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol = withValidityAssertion {
+    override fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol {
         return if (psi.hasBody() && (psi.funKeyword == null || psi.nameIdentifier == null)) {
             getAnonymousFunctionSymbol(psi)
         } else {
@@ -51,31 +50,31 @@ internal class KtFe10SymbolProvider(
         }
     }
 
-    override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol = withValidityAssertion {
+    override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol {
         return KtFe10PsiConstructorSymbol(psi, analysisContext)
     }
 
-    override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol = withValidityAssertion {
+    override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol {
         return KtFe10PsiTypeParameterSymbol(psi, analysisContext)
     }
 
-    override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol = withValidityAssertion {
+    override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol {
         return KtFe10PsiTypeAliasSymbol(psi, analysisContext)
     }
 
-    override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol = withValidityAssertion {
+    override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol {
         return KtFe10PsiEnumEntrySymbol(psi, analysisContext)
     }
 
-    override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol = withValidityAssertion {
+    override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol {
         return KtFe10PsiAnonymousFunctionSymbol(psi, analysisContext)
     }
 
-    override fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol = withValidityAssertion {
+    override fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol {
         return KtFe10PsiLiteralAnonymousFunctionSymbol(psi, analysisContext)
     }
 
-    override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol = withValidityAssertion {
+    override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol {
         return if (psi.isLocal) {
             KtFe10PsiLocalVariableSymbol(psi, analysisContext)
         } else {
@@ -83,11 +82,11 @@ internal class KtFe10SymbolProvider(
         }
     }
 
-    override fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol = withValidityAssertion {
+    override fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol {
         return KtFe10PsiAnonymousObjectSymbol(psi.objectDeclaration, analysisContext)
     }
 
-    override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol = withValidityAssertion {
+    override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol {
         return if (psi is KtObjectDeclaration && psi.isObjectLiteral()) {
             KtFe10PsiAnonymousObjectSymbol(psi, analysisContext)
         } else {
@@ -95,7 +94,7 @@ internal class KtFe10SymbolProvider(
         }
     }
 
-    override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol? = withValidityAssertion {
+    override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol? {
         if (psi is KtEnumEntry || psi.nameIdentifier == null) {
             return null
         }
@@ -103,7 +102,7 @@ internal class KtFe10SymbolProvider(
         return KtFe10PsiNamedClassOrObjectSymbol(psi, analysisContext)
     }
 
-    override fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol = withValidityAssertion {
+    override fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol {
         return if (psi.isGetter) {
             KtFe10PsiPropertyGetterSymbol(psi, analysisContext)
         } else {
@@ -111,16 +110,16 @@ internal class KtFe10SymbolProvider(
         }
     }
 
-    override fun getClassInitializerSymbol(psi: KtClassInitializer): KtClassInitializerSymbol = withValidityAssertion {
+    override fun getClassInitializerSymbol(psi: KtClassInitializer): KtClassInitializerSymbol {
         return KtFe10PsiClassInitializerSymbol(psi, analysisContext)
     }
 
-    override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? = withValidityAssertion {
+    override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? {
         val descriptor = analysisContext.resolveSession.moduleDescriptor.findClassAcrossModuleDependencies(classId) ?: return null
         return descriptor.toKtClassSymbol(analysisContext)
     }
 
-    override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): Sequence<KtSymbol> = withValidityAssertion {
+    override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): Sequence<KtSymbol> {
         val packageViewDescriptor = analysisContext.resolveSession.moduleDescriptor.getPackage(packageFqName)
         return packageViewDescriptor.memberScope.getContributedDescriptors(DescriptorKindFilter.ALL, nameFilter = { it == name })
             .asSequence()

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10TypeCreator.kt

@@ -21,12 +21,13 @@ import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtClassType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
 import org.jetbrains.kotlin.analysis.api.types.KtTypeParameterType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.types.SimpleType
-import org.jetbrains.kotlin.types.error.ErrorType
+import org.jetbrains.kotlin.types.StarProjectionImpl
+import org.jetbrains.kotlin.types.TypeProjectionImpl
+import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.error.ErrorTypeKind
 import org.jetbrains.kotlin.types.error.ErrorUtils
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
@@ -37,7 +38,7 @@ internal class KtFe10TypeCreator(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun buildClassType(builder: KtClassTypeBuilder): KtClassType = withValidityAssertion {
+    override fun buildClassType(builder: KtClassTypeBuilder): KtClassType {
         val descriptor: ClassDescriptor? = when (builder) {
             is KtClassTypeBuilder.ByClassId -> {
                 val fqName = builder.classId.asSingleFqName()
@@ -73,7 +74,7 @@ internal class KtFe10TypeCreator(
         return KtFe10UsualClassType(typeWithNullability as SimpleType, descriptor, analysisContext)
     }
 
-    override fun buildTypeParameterType(builder: KtTypeParameterTypeBuilder): KtTypeParameterType = withValidityAssertion {
+    override fun buildTypeParameterType(builder: KtTypeParameterTypeBuilder): KtTypeParameterType {
         val descriptor = when (builder) {
             is KtTypeParameterTypeBuilder.BySymbol -> {
                 getSymbolDescriptor(builder.symbol) as? TypeParameterDescriptor

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10TypeInfoProvider.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnaly
 import org.jetbrains.kotlin.analysis.api.descriptors.types.base.KtFe10Type
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.builtins.getFunctionalClassKind
 import org.jetbrains.kotlin.load.java.sam.JavaSingleAbstractMethodUtils
@@ -26,7 +25,7 @@ internal class KtFe10TypeInfoProvider(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override fun isFunctionalInterfaceType(type: KtType): Boolean = withValidityAssertion {
+    override fun isFunctionalInterfaceType(type: KtType): Boolean {
         require(type is KtFe10Type)
         return JavaSingleAbstractMethodUtils.isSamType(type.type)
     }
@@ -36,7 +35,7 @@ internal class KtFe10TypeInfoProvider(
         return type.type.constructor.declarationDescriptor?.getFunctionalClassKind()
     }
 
-    override fun canBeNull(type: KtType): Boolean = withValidityAssertion {
+    override fun canBeNull(type: KtType): Boolean {
         require(type is KtFe10Type)
         return TypeUtils.isNullableType(type.type)
     }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10TypeProvider.kt

@@ -65,39 +65,39 @@ internal class KtFe10TypeProvider(
     override val token: KtLifetimeToken
         get() = analysisSession.token
 
-    override val builtinTypes: KtBuiltinTypes by cached { KtFe10BuiltinTypes(analysisContext) }
+    override val builtinTypes: KtBuiltinTypes by lazy(LazyThreadSafetyMode.PUBLICATION) { KtFe10BuiltinTypes(analysisContext) }
 
-    override fun approximateToSuperPublicDenotableType(type: KtType): KtType? = withValidityAssertion {
+    override fun approximateToSuperPublicDenotableType(type: KtType): KtType?  {
         require(type is KtFe10Type)
         return typeApproximator.approximateToSuperType(type.type, PublicApproximatorConfiguration)?.toKtType(analysisContext)
     }
 
-    override fun buildSelfClassType(symbol: KtNamedClassOrObjectSymbol): KtType = withValidityAssertion {
+    override fun buildSelfClassType(symbol: KtNamedClassOrObjectSymbol): KtType  {
         val kotlinType = (getSymbolDescriptor(symbol) as? ClassDescriptor)?.defaultType
             ?: ErrorUtils.createErrorType(ErrorTypeKind.UNRESOLVED_CLASS_TYPE, symbol.nameOrAnonymous.toString())
 
         return kotlinType.toKtType(analysisContext)
     }
 
-    override fun commonSuperType(types: Collection<KtType>): KtType = withValidityAssertion {
+    override fun commonSuperType(types: Collection<KtType>): KtType  {
         val kotlinTypes = types.map { (it as KtFe10Type).type }
         return CommonSupertypes.commonSupertype(kotlinTypes).toKtType(analysisContext)
     }
 
-    override fun getKtType(ktTypeReference: KtTypeReference): KtType = withValidityAssertion {
+    override fun getKtType(ktTypeReference: KtTypeReference): KtType  {
         val bindingContext = analysisContext.analyze(ktTypeReference, AnalysisMode.PARTIAL)
         val kotlinType = bindingContext[BindingContext.TYPE, ktTypeReference]
             ?: ErrorUtils.createErrorType(ErrorTypeKind.UNRESOLVED_TYPE, ktTypeReference.text)
         return kotlinType.toKtType(analysisContext)
     }
 
-    override fun getReceiverTypeForDoubleColonExpression(expression: KtDoubleColonExpression): KtType? = withValidityAssertion {
+    override fun getReceiverTypeForDoubleColonExpression(expression: KtDoubleColonExpression): KtType?  {
         val bindingContext = analysisContext.analyze(expression, AnalysisMode.PARTIAL)
         val lhs = bindingContext[BindingContext.DOUBLE_COLON_LHS, expression] ?: return null
         return lhs.type.toKtType(analysisContext)
     }
 
-    override fun withNullability(type: KtType, newNullability: KtTypeNullability): KtType = withValidityAssertion {
+    override fun withNullability(type: KtType, newNullability: KtTypeNullability): KtType  {
         require(type is KtFe10Type)
         return type.type.makeNullableAsSpecified(newNullability == KtTypeNullability.NULLABLE).toKtType(analysisContext)
     }
@@ -125,7 +125,6 @@ internal class KtFe10TypeProvider(
     }
 
     override fun getDispatchReceiverType(symbol: KtCallableSymbol): KtType? {
-        assertIsValidAndAccessible()
         require(symbol is KtFe10Symbol)
 
         val descriptor = when (symbol) {

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10VisibilityChecker.kt

@@ -12,10 +12,9 @@ import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
 import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnalysisSessionComponent
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.getSymbolDescriptor
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.psiBased.base.getResolutionScope
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtFileSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithVisibility
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithVisibility
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilityUtils.isVisible
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -39,7 +38,7 @@ internal class KtFe10VisibilityChecker(
         useSiteFile: KtFileSymbol,
         position: PsiElement,
         receiverExpression: KtExpression?
-    ): Boolean = withValidityAssertion {
+    ): Boolean {
         if (candidateSymbol.visibility == Visibilities.Public) {
             return true
         }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/KtFe10FileSymbol.kt

@@ -25,7 +25,7 @@ internal class KtFe10FileSymbol(
         get() = withValidityAssertion { file }
 
     override val annotationsObject: Annotations
-        get() = analysisContext.resolveSession.getFileAnnotations(file)
+        get() = withValidityAssertion { analysisContext.resolveSession.getFileAnnotations(file) }
 
     override val origin: KtSymbolOrigin
         get() = withValidityAssertion { if (file.isCompiled) KtSymbolOrigin.LIBRARY else KtSymbolOrigin.SOURCE }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/descriptorBased/KtFe10DescFunctionSymbol.kt

@@ -58,7 +58,7 @@ internal class KtFe10DescFunctionSymbol private constructor(
         get() = withValidityAssertion { descriptor is JavaCallableMemberDescriptor && DescriptorUtils.isStaticDeclaration(descriptor) }
 
     override val isBuiltinFunctionInvoke: Boolean
-        get() = callableIdIfNonLocal in kotlinFunctionInvokeCallableIds
+        get() = withValidityAssertion { callableIdIfNonLocal in kotlinFunctionInvokeCallableIds }
 
     override val valueParameters: List<KtValueParameterSymbol>
         get() = withValidityAssertion { descriptor.valueParameters.map { KtFe10DescValueParameterSymbol(it, analysisContext) } }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/descriptorBased/KtFe10DescSyntheticFieldSymbol.kt

@@ -38,7 +38,7 @@ internal class KtFe10DescSyntheticFieldSymbol(
         }
 
     override val annotationsObject: Annotations
-        get() = descriptor.annotations
+        get() = withValidityAssertion { descriptor.annotations }
 
     override val returnType: KtType
         get() = withValidityAssertion { descriptor.propertyDescriptor.type.toKtType(analysisContext) }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/descriptorBased/KtFe10ReceiverParameterSymbol.kt

@@ -25,7 +25,7 @@ class KtFe10ReceiverParameterSymbol(
     override val descriptor: ReceiverParameterDescriptor
         get() = withValidityAssertion { _descriptor }
 
-    override fun createPointer(): KtSymbolPointer<KtSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtSymbol> = withValidityAssertion {
         TODO("Not yet implemented")
     }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/descriptorBased/base/KtFe10DescSymbol.kt

@@ -21,10 +21,10 @@ internal interface KtFe10DescSymbol<T : DeclarationDescriptor> : KtFe10Symbol, K
     val descriptor: T
 
     override val annotationsObject: Annotations
-        get() = descriptor.annotations
+        get() = withValidityAssertion { descriptor.annotations }
 
     val source: SourceElement
-        get() {
+        get() = withValidityAssertion {
             val descriptor = this.descriptor
             if (descriptor is CallableMemberDescriptor && descriptor.kind == CallableMemberDescriptor.Kind.FAKE_OVERRIDE) {
                 val firstOverridden = descriptor.overriddenDescriptors.firstOrNull()

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/psiBased/KtFe10PsiClassInitializerSymbol.kt

@@ -33,7 +33,7 @@ class KtFe10PsiClassInitializerSymbol(
     override val symbolKind: KtSymbolKind
         get() = withValidityAssertion { KtSymbolKind.CLASS_MEMBER }
 
-    override fun createPointer(): KtSymbolPointer<KtClassInitializerSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtClassInitializerSymbol> = withValidityAssertion {
         return KtPsiBasedSymbolPointer.createForSymbolFromSource(this) ?: KtFe10NeverRestoringSymbolPointer()
     }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/psiBased/KtFe10PsiFunctionSymbol.kt

@@ -64,7 +64,7 @@ internal class KtFe10PsiFunctionSymbol(
         get() = withValidityAssertion { false }
 
     override val isBuiltinFunctionInvoke: Boolean
-        get() = callableIdIfNonLocal in kotlinFunctionInvokeCallableIds
+        get() = withValidityAssertion { callableIdIfNonLocal in kotlinFunctionInvokeCallableIds }
 
     override val valueParameters: List<KtValueParameterSymbol>
         get() = withValidityAssertion { psi.valueParameters.map { KtFe10PsiValueParameterSymbol(it, analysisContext) } }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/symbols/psiBased/base/KtFe10PsiSymbol.kt

@@ -18,7 +18,7 @@ internal interface KtFe10PsiSymbol<P : KtElement, D : DeclarationDescriptor> : K
     val descriptor: D?
 
     override val annotationsObject: Annotations
-        get() = descriptor?.annotations ?: Annotations.EMPTY
+        get() = withValidityAssertion { descriptor?.annotations ?: Annotations.EMPTY }
 
     override val origin: KtSymbolOrigin
         get() = withValidityAssertion { psi.ktSymbolOrigin }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/types/KtFe10FunctionalType.kt

@@ -55,7 +55,12 @@ internal class KtFe10FunctionalType(
         get() = withValidityAssertion { type.getReturnTypeFromFunctionType().toKtType(analysisContext) }
 
     override val classId: ClassId
-        get() = ClassId(descriptor.functionKind.packageFqName, descriptor.functionKind.numberedClassName(descriptor.arity))
+        get() = withValidityAssertion {
+            ClassId(
+                descriptor.functionKind.packageFqName,
+                descriptor.functionKind.numberedClassName(descriptor.arity)
+            )
+        }
 
     override val classSymbol: KtClassLikeSymbol
         get() = withValidityAssertion { KtFe10DescNamedClassOrObjectSymbol(descriptor, analysisContext) }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSession.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbolProvider
 import org.jetbrains.kotlin.analysis.api.fir.utils.threadLocal
 import org.jetbrains.kotlin.analysis.api.impl.base.components.KtAnalysisScopeProviderImpl
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -25,7 +26,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-
+@Suppress("AnalysisApiMissingLifetimeCheck")
 internal class KtFirAnalysisSession
 private constructor(
     private val project: Project,
@@ -35,6 +36,7 @@ private constructor(
     private val mode: AnalysisSessionMode,
 ) : KtAnalysisSession(token) {
 
+    @Suppress("AnalysisApiMissingLifetimeCheck")
     override val useSiteModule: KtModule get() = firResolveSession.useSiteKtModule
 
     private enum class AnalysisSessionMode {
@@ -54,10 +56,10 @@ private constructor(
 
     override val samResolverImpl = KtFirSamResolver(this, token)
 
-    override val scopeProviderImpl by threadLocal { KtFirScopeProvider(this, firSymbolBuilder, project, firResolveSession, token) }
+    override val scopeProviderImpl by threadLocal { KtFirScopeProvider(this, firSymbolBuilder, project, firResolveSession) }
 
     override val symbolProviderImpl =
-        KtFirSymbolProvider(this, firResolveSession.useSiteFirSession.symbolProvider, firResolveSession, firSymbolBuilder, token)
+        KtFirSymbolProvider(this, firResolveSession.useSiteFirSession.symbolProvider, firResolveSession, firSymbolBuilder)
 
     override val completionCandidateCheckerImpl = KtFirCompletionCandidateChecker(this, token)
 
@@ -99,6 +101,7 @@ private constructor(
 
     override val referenceResolveProviderImpl: KtReferenceResolveProvider = KtFirReferenceResolveProvider(this)
 
+    @Suppress("AnalysisApiMissingLifetimeCheck")
     override fun createContextDependentCopy(originalKtFile: KtFile, elementToReanalyze: KtElement): KtAnalysisSession {
         check(mode == AnalysisSessionMode.REGULAR) {
             "Cannot create context-dependent copy of KtAnalysis session from a context dependent one"
@@ -120,11 +123,11 @@ private constructor(
         )
     }
 
-    val useSiteSession: FirSession get() = firResolveSession.useSiteFirSession
+    internal val useSiteSession: FirSession get() = firResolveSession.useSiteFirSession
-    val firSymbolProvider: FirSymbolProvider get() = useSiteSession.symbolProvider
+    internal val firSymbolProvider: FirSymbolProvider get() = useSiteSession.symbolProvider
-    val targetPlatform: TargetPlatform get() = useSiteSession.moduleData.platform
+    internal val targetPlatform: TargetPlatform get() = useSiteSession.moduleData.platform
 
-    fun getScopeSessionFor(session: FirSession): ScopeSession = firResolveSession.getScopeSessionFor(session)
+    fun getScopeSessionFor(session: FirSession): ScopeSession = withValidityAssertion { firResolveSession.getScopeSessionFor(session) }
 
     companion object {
         internal fun createAnalysisSessionByFirResolveSession(

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtSymbolByFirBuilder.kt

@@ -55,13 +55,13 @@ import kotlin.contracts.contract
 internal class KtSymbolByFirBuilder private constructor(
     private val project: Project,
     private val firResolveSession: LLFirResolveSession,
-    override val token: KtLifetimeToken,
+    val token: KtLifetimeToken,
     val withReadOnlyCaching: Boolean,
     private val symbolsCache: BuilderCache<FirBasedSymbol<*>, KtSymbol>,
     private val extensionReceiverSymbolsCache: BuilderCache<FirCallableSymbol<*>, KtSymbol>,
     private val filesCache: BuilderCache<FirFileSymbol, KtFileSymbol>,
     private val backingFieldCache: BuilderCache<FirBackingFieldSymbol, KtBackingFieldSymbol>,
-) : KtLifetimeOwner {
+) {
     private val firProvider get() = firResolveSession.useSiteFirSession.symbolProvider
     val rootSession: FirSession = firResolveSession.useSiteFirSession
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -16,10 +16,9 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.impl.barebone.parentOfType
 import org.jetbrains.kotlin.analysis.api.impl.base.components.AbstractKtCallResolver
-import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.AllCandidatesResolver
 import org.jetbrains.kotlin.fir.FirElement
@@ -75,7 +74,7 @@ internal class KtFirCallResolver(
         result
     }
 
-    override fun resolveCall(psi: KtElement): KtCallInfo? = withValidityAssertion {
+    override fun resolveCall(psi: KtElement): KtCallInfo? {
         val ktCallInfos = getCallInfo(psi) { psiToResolve, resolveCalleeExpressionOfFunctionCall, resolveFragmentOfCall ->
             listOfNotNull(
                 toKtCallInfo(
@@ -703,8 +702,8 @@ internal class KtFirCallResolver(
     private fun FirVariableSymbol<*>.toKtSignature(): KtVariableLikeSignature<KtVariableLikeSymbol> =
         firSymbolBuilder.variableLikeBuilder.buildVariableLikeSignature(this)
 
-    override fun collectCallCandidates(psi: KtElement): List<KtCallCandidateInfo> = withValidityAssertion {
+    override fun collectCallCandidates(psi: KtElement): List<KtCallCandidateInfo> {
-        getCallInfo(psi) { psiToResolve, resolveCalleeExpressionOfFunctionCall, resolveFragmentOfCall ->
+        return getCallInfo(psi) { psiToResolve, resolveCalleeExpressionOfFunctionCall, resolveFragmentOfCall ->
             collectCallCandidates(
                 psiToResolve,
                 resolveCalleeExpressionOfFunctionCall,

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompileTimeConstantProvider.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.api.components.KtConstantEvaluationMode
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.evaluate.FirCompileTimeConstantEvaluator
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.fir.FirElement
@@ -29,16 +28,16 @@ internal class KtFirCompileTimeConstantProvider(
     override fun evaluate(
         expression: KtExpression,
         mode: KtConstantEvaluationMode,
-    ): KtConstantValue? = withValidityAssertion {
+    ): KtConstantValue? {
-        evaluateFir(expression.getOrBuildFir(firResolveSession), expression, mode)
+        return evaluateFir(expression.getOrBuildFir(firResolveSession), expression, mode)
     }
 
     private fun evaluateFir(
         fir: FirElement?,
         sourcePsi: KtExpression,
         mode: KtConstantEvaluationMode,
-    ): KtConstantValue? = withValidityAssertion {
+    ): KtConstantValue? {
-        when (fir) {
+        return when (fir) {
             is FirPropertyAccessExpression,
             is FirExpression,
             is FirNamedReference -> {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompletionCandidateChecker.kt

@@ -9,11 +9,9 @@ import org.jetbrains.kotlin.analysis.api.components.KtCompletionCandidateChecker
 import org.jetbrains.kotlin.analysis.api.components.KtExtensionApplicabilityResult
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
-import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
-import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.ResolutionParameters
@@ -36,22 +34,19 @@ import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForReceiver
 
 internal class KtFirCompletionCandidateChecker(
-    analysisSession: KtFirAnalysisSession,
+    override val analysisSession: KtFirAnalysisSession,
     override val token: KtLifetimeToken,
 ) : KtCompletionCandidateChecker(), KtFirAnalysisSessionComponent {
-    override val analysisSession: KtFirAnalysisSession by weakRef(analysisSession)
-
     override fun checkExtensionFitsCandidate(
         firSymbolForCandidate: KtCallableSymbol,
         originalFile: KtFile,
         nameExpression: KtSimpleNameExpression,
         possibleExplicitReceiver: KtExpression?,
-    ): KtExtensionApplicabilityResult = withValidityAssertion {
+    ): KtExtensionApplicabilityResult {
         require(firSymbolForCandidate is KtFirSymbol<*>)
         firSymbolForCandidate.firSymbol.ensureResolved(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
         val declaration = firSymbolForCandidate.firSymbol.fir as FirCallableDeclaration
-        checkExtension(declaration, originalFile, nameExpression, possibleExplicitReceiver)
+        return checkExtension(declaration, originalFile, nameExpression, possibleExplicitReceiver)
-
     }
 
     private fun checkExtension(
@@ -75,15 +70,15 @@ internal class KtFirCompletionCandidateChecker(
                 val substitutor = it.createSubstitutorFromTypeArguments() ?: return@let null
                 return when {
                     candidateSymbol is FirVariable && candidateSymbol.returnTypeRef.coneType.receiverType(rootModuleSession) != null -> {
-                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall(substitutor)
+                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall(substitutor, token)
                     }
                     else -> {
-                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable(substitutor)
+                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable(substitutor, token)
                     }
                 }
             }
         }
-        return KtExtensionApplicabilityResult.NonApplicable(KtSubstitutor.Empty(token))
+        return KtExtensionApplicabilityResult.NonApplicable(KtSubstitutor.Empty(token), token)
     }
 
     private fun getImplicitReceivers(

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirDiagnosticProvider.kt

@@ -25,8 +25,8 @@ internal class KtFirDiagnosticProvider(
     override fun getDiagnosticsForElement(
         element: KtElement,
         filter: KtDiagnosticCheckerFilter
-    ): Collection<KtDiagnosticWithPsi<*>> = withValidityAssertion {
+    ): Collection<KtDiagnosticWithPsi<*>> {
-        element.getDiagnostics(firResolveSession, filter.asLLFilter()).map { it.asKtDiagnostic() }
+        return element.getDiagnostics(firResolveSession, filter.asLLFilter()).map { it.asKtDiagnostic() }
     }
 
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>> =

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirExpressionTypeProvider.kt

@@ -11,8 +11,8 @@ import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.utils.getReferencedElementType
 import org.jetbrains.kotlin.analysis.api.fir.utils.unwrap
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirSafe
@@ -37,8 +37,8 @@ internal class KtFirExpressionTypeProvider(
     override val token: KtLifetimeToken,
 ) : KtExpressionTypeProvider(), KtFirAnalysisSessionComponent {
 
-    override fun getKtExpressionType(expression: KtExpression): KtType? = withValidityAssertion {
+    override fun getKtExpressionType(expression: KtExpression): KtType? {
-        when (val fir = expression.unwrap().getOrBuildFir(firResolveSession)) {
+        return when (val fir = expression.unwrap().getOrBuildFir(firResolveSession)) {
             is FirFunctionCall -> {
                 getReturnTypeForArrayStyleAssignmentTarget(expression, fir)
                     ?: fir.typeRef.coneType.asKtType()
@@ -94,14 +94,14 @@ internal class KtFirExpressionTypeProvider(
         return setTargetArgumentParameter.returnTypeRef.coneType.asKtType()
     }
 
-    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType = withValidityAssertion {
+    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType {
         val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration>(firResolveSession)
-        firDeclaration.returnTypeRef.coneType.asKtType()
+        return firDeclaration.returnTypeRef.coneType.asKtType()
     }
 
-    override fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType = withValidityAssertion {
+    override fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType {
         val firFunction = declaration.getOrBuildFirOfType<FirFunction>(firResolveSession)
-        firFunction.constructFunctionalType(firFunction.isSuspend).asKtType()
+        return firFunction.constructFunctionalType(firFunction.isSuspend).asKtType()
     }
 
     override fun getExpectedType(expression: PsiElement): KtType? {
@@ -224,7 +224,7 @@ internal class KtFirExpressionTypeProvider(
     override fun isDefinitelyNotNull(expression: KtExpression): Boolean =
         getDefiniteNullability(expression) == DefiniteNullability.DEFINITELY_NOT_NULL
 
-    private fun getDefiniteNullability(expression: KtExpression): DefiniteNullability = withValidityAssertion {
+    private fun getDefiniteNullability(expression: KtExpression): DefiniteNullability {
         fun FirExpression.isNotNullable() = with(analysisSession.useSiteSession.typeContext) {
             !typeRef.coneType.isNullableType()
         }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirInheritorsProvider.kt

@@ -8,11 +8,10 @@ package org.jetbrains.kotlin.analysis.api.fir.components
 import org.jetbrains.kotlin.analysis.api.components.KtInheritorsProvider
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirNamedClassOrObjectSymbol
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassKind
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtNamedClassOrObjectSymbol
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.declarations.getSealedClassInheritors
 
@@ -22,20 +21,20 @@ internal class KtFirInheritorsProvider(
 ) : KtInheritorsProvider(), KtFirAnalysisSessionComponent {
     override fun getInheritorsOfSealedClass(
         classSymbol: KtNamedClassOrObjectSymbol
-    ): List<KtNamedClassOrObjectSymbol> = withValidityAssertion {
+    ): List<KtNamedClassOrObjectSymbol> {
         require(classSymbol.modality == Modality.SEALED)
         require(classSymbol is KtFirNamedClassOrObjectSymbol)
 
         val inheritorClassIds = classSymbol.firSymbol.fir.getSealedClassInheritors(analysisSession.useSiteSession)
 
-        with(analysisSession) {
+        return with(analysisSession) {
             inheritorClassIds.mapNotNull { it.getCorrespondingToplevelClassOrObjectSymbol() as? KtNamedClassOrObjectSymbol }
         }
     }
 
-    override fun getEnumEntries(classSymbol: KtNamedClassOrObjectSymbol): List<KtEnumEntrySymbol> = withValidityAssertion {
+    override fun getEnumEntries(classSymbol: KtNamedClassOrObjectSymbol): List<KtEnumEntrySymbol> {
         require(classSymbol.classKind == KtClassKind.ENUM_CLASS)
-        with(analysisSession) {
+        return with(analysisSession) {
             classSymbol.getDeclaredMemberScope().getCallableSymbols().filterIsInstance<KtEnumEntrySymbol>().toList()
         }
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirPsiTypeProvider.kt

@@ -53,13 +53,12 @@ internal class KtFirPsiTypeProvider(
         useSitePosition: PsiElement,
         mode: KtTypeMappingMode,
         isAnnotationMethod: Boolean,
-    ): PsiType? = withValidityAssertion {
+    ): PsiType? = type.coneType.asPsiType(
-        type.coneType.asPsiType(
+        rootModuleSession,
-            rootModuleSession,
+        mode.toTypeMappingMode(type, isAnnotationMethod),
-            mode.toTypeMappingMode(type, isAnnotationMethod),
+        useSitePosition
-            useSitePosition
+    )
-        )
+
-    }
 
     private fun KtTypeMappingMode.toTypeMappingMode(type: KtType, isAnnotationMethod: Boolean): TypeMappingMode {
         require(type is KtFirType)

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt

@@ -46,20 +46,14 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.util.*
 
 internal class KtFirScopeProvider(
-    analysisSession: KtFirAnalysisSession,
+    override val analysisSession: KtFirAnalysisSession,
-    builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
     private val project: Project,
-    firResolveSession: LLFirResolveSession,
+    private val firResolveSession: LLFirResolveSession,
-    override val token: KtLifetimeToken,
+) : KtScopeProvider() {
-) : KtScopeProvider(), KtLifetimeOwner {
     // KtFirScopeProvider is thread local, so it's okay to use the same session here
     private val scopeSession = analysisSession.getScopeSessionFor(analysisSession.useSiteSession)
 
-
-    override val analysisSession: KtFirAnalysisSession by weakRef(analysisSession)
-    private val builder by weakRef(builder)
-    private val firResolveSession by weakRef(firResolveSession)
-
     private val memberScopeCache = IdentityHashMap<KtSymbolWithMembers, KtScope>()
     private val declaredMemberScopeCache = IdentityHashMap<KtSymbolWithMembers, KtScope>()
     private val delegatedMemberScopeCache = IdentityHashMap<KtSymbolWithMembers, KtScope>()
@@ -67,27 +61,27 @@ internal class KtFirScopeProvider(
     private val packageMemberScopeCache = IdentityHashMap<KtPackageSymbol, KtScope>()
 
     private inline fun <T> KtSymbolWithMembers.withFirForScope(crossinline body: (FirClass) -> T): T? {
-        when (this) {
+        return when (this) {
             is KtFirNamedClassOrObjectSymbol -> {
                 firSymbol.ensureResolved(FirResolvePhase.TYPES)
-                return body(firSymbol.fir)
+                body(firSymbol.fir)
             }
             is KtFirAnonymousObjectSymbol -> {
                 firSymbol.ensureResolved(FirResolvePhase.TYPES)
-                return body(firSymbol.fir)
+                body(firSymbol.fir)
             }
             is KtFirEnumEntrySymbol -> {
                 firSymbol.ensureResolved(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
                 val initializer = firSymbol.fir.initializer
                 check(initializer is FirAnonymousObjectExpression) { "Unexpected enum entry initializer: ${initializer?.javaClass}" }
-                return body(initializer.anonymousObject)
+                body(initializer.anonymousObject)
             }
             else -> error { "Unknown KtSymbolWithDeclarations implementation ${this::class.qualifiedName}" }
         }
     }
 
-    override fun getMemberScope(classSymbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getMemberScope(classSymbol: KtSymbolWithMembers): KtScope {
-        memberScopeCache.getOrPut(classSymbol) {
+        return memberScopeCache.getOrPut(classSymbol) {
             val firScope = classSymbol.withFirForScope { fir ->
                 val firSession = analysisSession.useSiteSession
                 fir.unsubstitutedScope(
@@ -113,8 +107,8 @@ internal class KtFirScopeProvider(
         return KtFirDelegatingScope(firScope, builder, token)
     }
 
-    override fun getDeclaredMemberScope(classSymbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getDeclaredMemberScope(classSymbol: KtSymbolWithMembers): KtScope {
-        declaredMemberScopeCache.getOrPut(classSymbol) {
+        return declaredMemberScopeCache.getOrPut(classSymbol) {
             val firScope = classSymbol.withFirForScope {
                 analysisSession.useSiteSession.declaredMemberScope(it)
             } ?: return@getOrPut getEmptyScope()
@@ -123,10 +117,10 @@ internal class KtFirScopeProvider(
         }
     }
 
-    override fun getDelegatedMemberScope(classSymbol: KtSymbolWithMembers): KtScope = withValidityAssertion {
+    override fun getDelegatedMemberScope(classSymbol: KtSymbolWithMembers): KtScope {
         val declaredScope = (getDeclaredMemberScope(classSymbol) as? KtFirDelegatingScope)?.firScope
             ?: return delegatedMemberScopeCache.getOrPut(classSymbol) { getEmptyScope() }
-        delegatedMemberScopeCache.getOrPut(classSymbol) {
+        return delegatedMemberScopeCache.getOrPut(classSymbol) {
             val firScope = classSymbol.withFirForScope { fir ->
                 val delegateFields = fir.delegateFields
                 if (delegateFields.isNotEmpty()) {
@@ -145,19 +139,19 @@ internal class KtFirScopeProvider(
         }
     }
 
-    override fun getFileScope(fileSymbol: KtFileSymbol): KtScope = withValidityAssertion {
+    override fun getFileScope(fileSymbol: KtFileSymbol): KtScope {
-        fileScopeCache.getOrPut(fileSymbol) {
+        return fileScopeCache.getOrPut(fileSymbol) {
             check(fileSymbol is KtFirFileSymbol) { "KtFirScopeProvider can only work with KtFirFileSymbol, but ${fileSymbol::class} was provided" }
             KtFirFileScope(fileSymbol, token, builder)
         }
     }
 
-    override fun getEmptyScope(): KtScope = withValidityAssertion {
+    override fun getEmptyScope(): KtScope {
-        KtEmptyScope(token)
+        return KtEmptyScope(token)
     }
 
-    override fun getPackageScope(packageSymbol: KtPackageSymbol): KtScope = withValidityAssertion {
+    override fun getPackageScope(packageSymbol: KtPackageSymbol): KtScope {
-        packageMemberScopeCache.getOrPut(packageSymbol) {
+        return packageMemberScopeCache.getOrPut(packageSymbol) {
             KtFirPackageScope(
                 packageSymbol.fqName,
                 project,
@@ -170,8 +164,8 @@ internal class KtFirScopeProvider(
     }
 
 
-    override fun getCompositeScope(subScopes: List<KtScope>): KtScope = withValidityAssertion {
+    override fun getCompositeScope(subScopes: List<KtScope>): KtScope {
-        KtCompositeScope(subScopes, token)
+        return KtCompositeScope(subScopes, token)
     }
 
     override fun getTypeScope(type: KtType): KtScope? {
@@ -196,7 +190,7 @@ internal class KtFirScopeProvider(
     override fun getScopeContextForPosition(
         originalFile: KtFile,
         positionInFakeFile: KtElement
-    ): KtScopeContext = withValidityAssertion {
+    ): KtScopeContext {
         val towerDataContext =
             analysisSession.firResolveSession.getTowerContextProvider(originalFile).getClosestAvailableParentContext(positionInFakeFile)
                 ?: error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
@@ -221,9 +215,10 @@ internal class KtFirScopeProvider(
             firLocalScopes.mapTo(this, ::convertToKtScope)
         }
 
-        KtScopeContext(
+        return KtScopeContext(
             getCompositeScope(allKtScopes.asReversed()),
-            implicitKtReceivers.asReversed()
+            implicitKtReceivers.asReversed(),
+            token
         )
     }
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSmartcastProvider.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.api.components.KtSmartCastInfo
 import org.jetbrains.kotlin.analysis.api.components.KtSmartCastProvider
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.FirImplicitInvokeCall
@@ -59,9 +58,9 @@ internal class KtFirSmartcastProvider(
         }
     }
 
-    override fun getSmartCastedInfo(expression: KtExpression): KtSmartCastInfo? = withValidityAssertion {
+    override fun getSmartCastedInfo(expression: KtExpression): KtSmartCastInfo? {
         val firSmartCastExpression = getMatchingFirExpressionWithSmartCast(expression) ?: return null
-        getSmartCastedInfo(firSmartCastExpression)
+        return getSmartCastedInfo(firSmartCastExpression)
     }
 
     private fun getSmartCastedInfo(expression: FirExpressionWithSmartcast): KtSmartCastInfo? {
@@ -87,10 +86,10 @@ internal class KtFirSmartcastProvider(
         }
     }
 
-    override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<KtImplicitReceiverSmartCast> = withValidityAssertion {
+    override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<KtImplicitReceiverSmartCast> {
         val firQualifiedExpression = getMatchingFirQualifiedAccessExpression(expression) ?: return emptyList()
 
-        listOfNotNull(
+        return listOfNotNull(
             smartCastedImplicitReceiver(firQualifiedExpression, KtImplicitReceiverSmartCastKind.DISPATCH),
             smartCastedImplicitReceiver(firQualifiedExpression, KtImplicitReceiverSmartCastKind.EXTENSION),
         )

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSubtypingComponent.kt

@@ -5,20 +5,19 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.components
 
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.components.KtSubtypingComponent
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.types.KtFirType
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.assertIsValidAndAccessible
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 
 internal class KtFirSubtypingComponent(
     override val analysisSession: KtFirAnalysisSession,
     override val token: KtLifetimeToken,
 ) : KtSubtypingComponent(), KtFirAnalysisSessionComponent {
-    override fun isEqualTo(first: KtType, second: KtType): Boolean = withValidityAssertion {
+    override fun isEqualTo(first: KtType, second: KtType): Boolean {
         second.assertIsValidAndAccessible()
         check(first is KtFirType)
         check(second is KtFirType)
@@ -29,7 +28,7 @@ internal class KtFirSubtypingComponent(
         )
     }
 
-    override fun isSubTypeOf(subType: KtType, superType: KtType): Boolean = withValidityAssertion {
+    override fun isSubTypeOf(subType: KtType, superType: KtType): Boolean {
         superType.assertIsValidAndAccessible()
         check(subType is KtFirType)
         check(superType is KtFirType)

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt

@@ -26,12 +26,9 @@ import org.jetbrains.kotlin.psi.*
 
 internal class KtFirSymbolContainingDeclarationProvider(
     override val analysisSession: KtFirAnalysisSession,
-    override val token: KtLifetimeToken,
+    override val token: KtLifetimeToken
 ) : KtSymbolContainingDeclarationProvider(), KtFirAnalysisSessionComponent {
-    override fun getContainingDeclaration(symbol: KtSymbol): KtSymbolWithKind? {
+    override fun getContainingDeclaration(symbol: KtSymbol): KtSymbolWithKind? { if (symbol is KtReceiverParameterSymbol) {
-        assertIsValidAndAccessible()
-
-        if (symbol is KtReceiverParameterSymbol) {
             return firSymbolBuilder.buildSymbol((symbol as KtFirReceiverParameterSymbol).firSymbol) as KtSymbolWithKind
         }
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolDeclarationOverridesProvider.kt

@@ -30,7 +30,7 @@ import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
 
 internal class KtFirSymbolDeclarationOverridesProvider(
     override val analysisSession: KtFirAnalysisSession,
-    override val token: KtLifetimeToken,
+    override val token: KtLifetimeToken
 ) : KtSymbolDeclarationOverridesProvider(), KtFirAnalysisSessionComponent {
 
     override fun <T : KtSymbol> getAllOverriddenSymbols(

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolDeclarationRendererProvider.kt

@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.fir.symbols.ensureResolved
 
 internal class KtFirSymbolDeclarationRendererProvider(
     override val analysisSession: KtFirAnalysisSession,
-    override val token: KtLifetimeToken,
+    override val token: KtLifetimeToken
 ) : KtSymbolDeclarationRendererProvider() {
 
     override fun render(type: KtType, options: KtTypeRendererOptions): String {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolInfoProvider.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.resolve.deprecation.DeprecationInfo
 
 internal class KtFirSymbolInfoProvider(
     override val analysisSession: KtFirAnalysisSession,
-    override val token: KtLifetimeToken,
+    override val token: KtLifetimeToken
 ) : KtSymbolInfoProvider(), KtFirAnalysisSessionComponent {
     override fun getDeprecation(symbol: KtSymbol): DeprecationInfo? {
         if (symbol is KtFirBackingFieldSymbol || symbol is KtFirPackageSymbol) return null

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeCreator.kt

@@ -29,7 +29,7 @@ internal class KtFirTypeCreator(
     override val token: KtLifetimeToken
 ) : KtTypeCreator(), KtFirAnalysisSessionComponent {
 
-    override fun buildClassType(builder: KtClassTypeBuilder): KtClassType = withValidityAssertion {
+    override fun buildClassType(builder: KtClassTypeBuilder): KtClassType {
         val lookupTag = when (builder) {
             is KtClassTypeBuilder.ByClassId -> {
                 val classSymbol = rootModuleSession.symbolProvider.getClassLikeSymbolByClassId(builder.classId)
@@ -53,7 +53,7 @@ internal class KtFirTypeCreator(
         return coneType.asKtType() as KtClassType
     }
 
-    override fun buildTypeParameterType(builder: KtTypeParameterTypeBuilder): KtTypeParameterType = withValidityAssertion {
+    override fun buildTypeParameterType(builder: KtTypeParameterTypeBuilder): KtTypeParameterType  {
         val coneType = when (builder) {
             is KtTypeParameterTypeBuilder.BySymbol -> {
                 val symbol = builder.symbol

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeInfoProvider.kt

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.fir.types.typeApproximator
 
 internal class KtFirTypeInfoProvider(
     override val analysisSession: KtFirAnalysisSession,
-    override val token: KtLifetimeToken,
+    override val token: KtLifetimeToken
 ) : KtTypeInfoProvider(), KtFirAnalysisSessionComponent {
 
     override fun isFunctionalInterfaceType(type: KtType): Boolean {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirTypeProvider.kt

@@ -50,7 +50,7 @@ import org.jetbrains.kotlin.util.bfs
 
 internal class KtFirTypeProvider(
     override val analysisSession: KtFirAnalysisSession,
-    override val token: KtLifetimeToken,
+    override val token: KtLifetimeToken
 ) : KtTypeProvider(), KtFirAnalysisSessionComponent {
     override val builtinTypes: KtBuiltinTypes = KtFirBuiltInTypes(rootModuleSession.builtinTypes, firSymbolBuilder, token)
 
@@ -85,16 +85,16 @@ internal class KtFirTypeProvider(
             ?.asKtType()
     }
 
-    override fun getKtType(ktTypeReference: KtTypeReference): KtType = withValidityAssertion {
+    override fun getKtType(ktTypeReference: KtTypeReference): KtType {
-        when (val fir = ktTypeReference.getOrBuildFir(firResolveSession)) {
+        return when (val fir = ktTypeReference.getOrBuildFir(firResolveSession)) {
             is FirResolvedTypeRef -> fir.coneType.asKtType()
             is FirDelegatedConstructorCall -> fir.constructedTypeRef.coneType.asKtType()
             else -> throwUnexpectedFirElementError(fir, ktTypeReference)
         }
     }
 
-    override fun getReceiverTypeForDoubleColonExpression(expression: KtDoubleColonExpression): KtType? = withValidityAssertion {
+    override fun getReceiverTypeForDoubleColonExpression(expression: KtDoubleColonExpression): KtType? {
-        when (val fir = expression.getOrBuildFir(firResolveSession)) {
+        return when (val fir = expression.getOrBuildFir(firResolveSession)) {
             is FirGetClassCall ->
                 fir.typeRef.coneType.getReceiverOfReflectionType()?.asKtType()
             is FirCallableReferenceAccess ->

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirDelegatedMemberScope.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analysis.api.scopes.KtScopeNameFilter
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 
 internal class KtFirDelegatedMemberScope(
@@ -18,7 +19,7 @@ internal class KtFirDelegatedMemberScope(
     builder: KtSymbolByFirBuilder
 ) : KtFirDelegatingScope(firScope, builder, token) {
 
-    override fun getCallableSymbols(nameFilter: KtScopeNameFilter): Sequence<KtCallableSymbol> {
+    override fun getCallableSymbols(nameFilter: KtScopeNameFilter): Sequence<KtCallableSymbol> = withValidityAssertion {
         return super.getCallableSymbols(nameFilter).filter { it.origin == KtSymbolOrigin.DELEGATED }
     }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirDelegatingScope.kt

@@ -31,7 +31,7 @@ internal open class KtFirDelegatingScope(
         getPossibleCallableNames() + getPossibleClassifierNames()
     }
 
-    override fun getAllPossibleNames(): Set<Name> = allNamesCached
+    override fun getAllPossibleNames(): Set<Name> = withValidityAssertion { withValidityAssertion { allNamesCached } }
 
     override fun getPossibleCallableNames(): Set<Name> = withValidityAssertion {
         firScope.getCallableNames()

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirFileScope.kt

@@ -32,7 +32,7 @@ internal class KtFirFileScope(
         _callableNames + _classifierNames
     }
 
-    override fun getAllPossibleNames(): Set<Name> = allNamesCached
+    override fun getAllPossibleNames(): Set<Name> = withValidityAssertion { allNamesCached }
 
     private val _callableNames: Set<Name> by cached {
         val result = mutableSetOf<Name>()
@@ -48,7 +48,7 @@ internal class KtFirFileScope(
         result
     }
 
-    override fun getPossibleCallableNames(): Set<Name> = _callableNames
+    override fun getPossibleCallableNames(): Set<Name> = withValidityAssertion { _callableNames }
 
     private val _classifierNames: Set<Name> by cached {
         val result = mutableSetOf<Name>()
@@ -60,7 +60,7 @@ internal class KtFirFileScope(
         result
     }
 
-    override fun getPossibleClassifierNames(): Set<Name> = _classifierNames
+    override fun getPossibleClassifierNames(): Set<Name> = withValidityAssertion { _classifierNames }
 
     override fun getCallableSymbols(nameFilter: KtScopeNameFilter): Sequence<KtCallableSymbol> = withValidityAssertion {
         sequence {
@@ -94,7 +94,7 @@ internal class KtFirFileScope(
     }
 
 
-    override fun getConstructors(): Sequence<KtConstructorSymbol> = emptySequence()
+    override fun getConstructors(): Sequence<KtConstructorSymbol> = withValidityAssertion { emptySequence() }
 
     override fun getPackageSymbols(nameFilter: KtScopeNameFilter): Sequence<KtPackageSymbol> = withValidityAssertion {
         emptySequence()

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirNonStarImportingScope.kt

@@ -57,7 +57,7 @@ internal class KtFirNonStarImportingScope(
         emptySequence()
     }
 
-    override fun getConstructors(): Sequence<KtConstructorSymbol> = emptySequence()
+    override fun getConstructors(): Sequence<KtConstructorSymbol> = withValidityAssertion { emptySequence() }
 
     override fun getPossibleCallableNames(): Set<Name> = withValidityAssertion {
         imports.mapNotNullTo(hashSetOf()) { it.callableName }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirStarImportingScope.kt

@@ -48,7 +48,7 @@ internal class KtFirStarImportingScope(
         firScope.getClassifierSymbols(getPossibleClassifierNames().filter(nameFilter), builder)
     }
 
-    override fun getConstructors(): Sequence<KtConstructorSymbol> = emptySequence()
+    override fun getConstructors(): Sequence<KtConstructorSymbol> = withValidityAssertion { emptySequence() }
 
     // todo cache?
     @OptIn(ExperimentalStdlibApi::class)

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirAnonymousFunctionSymbol.kt

@@ -11,13 +11,13 @@ import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.analysis.api.fir.annotations.KtFirAnnotationListForDeclaration
 import org.jetbrains.kotlin.analysis.api.fir.findPsi
 import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtAnonymousFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtValueParameterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtPsiBasedSymbolPointer
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.fir.resolve.getHasStableParameterNames
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
@@ -49,7 +49,7 @@ internal class KtFirAnonymousFunctionSymbol(
     override val isExtension: Boolean get() = withValidityAssertion { firSymbol.isExtension }
 
 
-    override fun createPointer(): KtSymbolPointer<KtAnonymousFunctionSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtAnonymousFunctionSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         error("Could not create a pointer for anonymous function from library")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirAnonymousObjectSymbol.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.pointers.CanNotCreateSymbolPoin
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtPsiBasedSymbolPointer
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousObjectSymbol
@@ -32,8 +33,12 @@ internal class KtFirAnonymousObjectSymbol(
     override val superTypes: List<KtType> by cached { firSymbol.superTypesList(builder) }
 
     override fun createPointer(): KtSymbolPointer<KtAnonymousObjectSymbol> =
-        KtPsiBasedSymbolPointer.createForSymbolFromSource(this)
+        withValidityAssertion {
-            ?: throw CanNotCreateSymbolPointerForLocalLibraryDeclarationException("Cannot create pointer for KtFirAnonymousObjectSymbol")
+            withValidityAssertion {
+                KtPsiBasedSymbolPointer.createForSymbolFromSource(this)
+                    ?: throw CanNotCreateSymbolPointerForLocalLibraryDeclarationException("Cannot create pointer for KtFirAnonymousObjectSymbol")
+            }
+        }
 
     override fun equals(other: Any?): Boolean = symbolEquals(other)
     override fun hashCode(): Int = symbolHashCode()

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirClassInitializerSymbol.kt

@@ -8,11 +8,12 @@ package org.jetbrains.kotlin.analysis.api.fir.symbols
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.fir.findPsi
 import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassInitializerSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousInitializerSymbol
 
@@ -23,9 +24,9 @@ internal class KtFirClassInitializerSymbol(
 ) : KtClassInitializerSymbol(), KtFirSymbol<FirAnonymousInitializerSymbol> {
     override val psi: PsiElement? by cached { firSymbol.findPsi() }
 
-    override fun createPointer(): KtSymbolPointer<KtSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtSymbol> = withValidityAssertion {
         TODO("Figure out how to create such a pointer. Should we give an index to class initializers?")
     }
 
-    override val symbolKind: KtSymbolKind get() = KtSymbolKind.CLASS_MEMBER
+    override val symbolKind: KtSymbolKind get() = withValidityAssertion { KtSymbolKind.CLASS_MEMBER }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirEnumEntrySymbol.kt

@@ -40,11 +40,11 @@ internal class KtFirEnumEntrySymbol(
 
     override val name: Name get() = withValidityAssertion { firSymbol.name }
     override val returnType: KtType get() = withValidityAssertion { firSymbol.returnType(builder) }
-    override val containingEnumClassIdIfNonLocal: ClassId? get() = callableIdIfNonLocal?.classId
+    override val containingEnumClassIdIfNonLocal: ClassId? get() = withValidityAssertion { callableIdIfNonLocal?.classId }
 
     override val callableIdIfNonLocal: CallableId? get() = withValidityAssertion { firSymbol.getCallableIdIfNonLocal() }
 
-    override fun createPointer(): KtSymbolPointer<KtEnumEntrySymbol> {
+    override fun createPointer(): KtSymbolPointer<KtEnumEntrySymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         return KtFirEnumEntrySymbolPointer(
             firSymbol.containingClass()?.classId ?: error("Containing class should present for enum entry"),

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirFileSymbol.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithDeclaration
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtPsiBasedSymbolPointer
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.fir.symbols.impl.FirFileSymbol
 
@@ -24,7 +25,7 @@ internal class KtFirFileSymbol(
 ) : KtFileSymbol(), KtSymbolWithDeclarations, KtFirSymbol<FirFileSymbol> {
     override val psi: PsiElement? by cached { firSymbol.findPsi() }
 
-    override fun createPointer(): KtSymbolPointer<KtFileSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtFileSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         TODO("Creating pointers for files from library is not supported yet")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirFunctionSymbol.kt

@@ -83,7 +83,7 @@ internal class KtFirFunctionSymbol(
     override val modality: Modality get() = withValidityAssertion { firSymbol.modalityOrFinal }
     override val visibility: Visibility get() = withValidityAssertion { firSymbol.visibility }
 
-    override fun createPointer(): KtSymbolPointer<KtFunctionSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtFunctionSymbol> = withValidityAssertion {
         if (firSymbol.fir.origin != FirDeclarationOrigin.SubstitutionOverride) {
             KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirJavaFieldSymbol.kt

@@ -52,7 +52,7 @@ internal class KtFirJavaFieldSymbol(
 
     override val isStatic: Boolean get() = withValidityAssertion { firSymbol.isStatic }
 
-    override fun createPointer(): KtSymbolPointer<KtJavaFieldSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtJavaFieldSymbol> = withValidityAssertion {
         val containingClassId = firSymbol.containingClass()?.classId
             ?: error("Cannot find parent class for java field $callableIdIfNonLocal")
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirKotlinPropertySymbol.kt

@@ -116,7 +116,7 @@ internal class KtFirKotlinPropertySymbol(
     override val hasGetter: Boolean get() = withValidityAssertion { firSymbol.getterSymbol != null }
     override val hasSetter: Boolean get() = withValidityAssertion { firSymbol.setterSymbol != null }
 
-    override fun createPointer(): KtSymbolPointer<KtKotlinPropertySymbol> {
+    override fun createPointer(): KtSymbolPointer<KtKotlinPropertySymbol> = withValidityAssertion {
         if (firSymbol.fir.origin != FirDeclarationOrigin.SubstitutionOverride) {
             KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirLocalVariableSymbol.kt

@@ -45,9 +45,9 @@ internal class KtFirLocalVariableSymbol(
     override val name: Name get() = withValidityAssertion { firSymbol.name }
     override val returnType: KtType get() = withValidityAssertion { firSymbol.returnType(builder) }
 
-    override val symbolKind: KtSymbolKind get() = KtSymbolKind.LOCAL
+    override val symbolKind: KtSymbolKind get() = withValidityAssertion { KtSymbolKind.LOCAL }
 
-    override fun createPointer(): KtSymbolPointer<KtLocalVariableSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtLocalVariableSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         throw CanNotCreateSymbolPointerForLocalLibraryDeclarationException(name.asString())
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirNamedClassOrObjectSymbol.kt

@@ -60,9 +60,11 @@ internal class KtFirNamedClassOrObjectSymbol(
 
     /* FirRegularClass visibility is not modified by STATUS only for Unknown, so it can be taken from RAW */
     override val visibility: Visibility
-        get() = when (val possiblyRawVisibility = firSymbol.fir.visibility) {
+        get() = withValidityAssertion {
-            Visibilities.Unknown -> if (firSymbol.fir.isLocal) Visibilities.Local else Visibilities.Public
+            when (val possiblyRawVisibility = firSymbol.fir.visibility) {
-            else -> possiblyRawVisibility
+                Visibilities.Unknown -> if (firSymbol.fir.isLocal) Visibilities.Local else Visibilities.Public
+                else -> possiblyRawVisibility
+            }
         }
 
     override val annotationsList by cached { KtFirAnnotationListForDeclaration.create(firSymbol, firResolveSession.useSiteFirSession, token) }
@@ -112,7 +114,7 @@ internal class KtFirNamedClassOrObjectSymbol(
         }
 
 
-    override fun createPointer(): KtSymbolPointer<KtNamedClassOrObjectSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtNamedClassOrObjectSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         if (symbolKind == KtSymbolKind.LOCAL) {
             throw CanNotCreateSymbolPointerForLocalLibraryDeclarationException(classIdIfNonLocal?.asString().orEmpty())

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirPackageSymbol.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.symbolPointer
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.providers.createPackageProvider
 import org.jetbrains.kotlin.name.FqName
 
@@ -33,11 +34,13 @@ class KtFirPackageSymbol(
     }
 
     override val origin: KtSymbolOrigin
-        get() = KtSymbolOrigin.SOURCE // TODO
+        get() = withValidityAssertion { KtSymbolOrigin.SOURCE } // TODO
 
-    override fun createPointer(): KtSymbolPointer<KtPackageSymbol> = symbolPointer { session ->
+    override fun createPointer(): KtSymbolPointer<KtPackageSymbol> = withValidityAssertion {
-        check(session is KtFirAnalysisSession)
+        symbolPointer { session ->
-        session.firSymbolBuilder.createPackageSymbolIfOneExists(fqName)
+            check(session is KtFirAnalysisSession)
+            session.firSymbolBuilder.createPackageSymbolIfOneExists(fqName)
+        }
     }
 
     override fun equals(other: Any?): Boolean {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirPropertyGetterSymbol.kt

@@ -72,7 +72,7 @@ internal class KtFirPropertyGetterSymbol(
     override val hasStableParameterNames: Boolean
         get() = withValidityAssertion { firSymbol.fir.getHasStableParameterNames(firSymbol.moduleData.session) }
 
-    override fun createPointer(): KtSymbolPointer<KtPropertyGetterSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtPropertyGetterSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         TODO("Creating pointers for getters from library is not supported yet")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirPropertySetterSymbol.kt

@@ -77,7 +77,7 @@ internal class KtFirPropertySetterSymbol(
     override val hasStableParameterNames: Boolean
         get() = withValidityAssertion { firSymbol.fir.getHasStableParameterNames(firSymbol.moduleData.session) }
 
-    override fun createPointer(): KtSymbolPointer<KtPropertySetterSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtPropertySetterSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         TODO("Creating pointers for setters from library is not supported yet")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirReceiverParameterSymbol.kt

@@ -41,7 +41,7 @@ internal class KtFirReceiverParameterSymbol(
 
     override val origin: KtSymbolOrigin = withValidityAssertion { firSymbol.fir.ktSymbolOrigin() }
 
-    override fun createPointer(): KtSymbolPointer<KtFirReceiverParameterSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtFirReceiverParameterSymbol> = withValidityAssertion {
         TODO("Probably just create a pointer based on the underlying declaration")
     }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSamConstructorSymbol.kt

@@ -63,7 +63,7 @@ internal class KtFirSamConstructorSymbol(
     override val typeParameters: List<KtTypeParameterSymbol>
         get() = withValidityAssertion { firSymbol.createKtTypeParameters(builder) }
 
-    override fun createPointer(): KtSymbolPointer<KtSamConstructorSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtSamConstructorSymbol> = withValidityAssertion {
         val callableId = firSymbol.callableId
         return KtFirSamConstructorSymbolPointer(ClassId(callableId.packageName, callableId.callableName))
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbol.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.fir.symbols
 
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
@@ -27,7 +28,7 @@ internal interface KtFirSymbol<out S : FirBasedSymbol<*>> : KtSymbol, KtLifetime
 
     val firResolveSession: LLFirResolveSession
 
-    override val origin: KtSymbolOrigin get() = firSymbol.fir.ktSymbolOrigin()
+    override val origin: KtSymbolOrigin get() = withValidityAssertion { firSymbol.fir.ktSymbolOrigin() }
 }
 
 internal fun KtFirSymbol<*>.symbolEquals(other: Any?): Boolean {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt

@@ -6,12 +6,8 @@
 package org.jetbrains.kotlin.analysis.api.fir.symbols
 
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
 import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
-import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
 import org.jetbrains.kotlin.analysis.api.symbols.*
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
-import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbolOfType
@@ -26,21 +22,23 @@ import org.jetbrains.kotlin.psi.*
 
 internal class KtFirSymbolProvider(
     override val analysisSession: KtAnalysisSession,
-    firSymbolProvider: FirSymbolProvider,
+    private val firSymbolProvider: FirSymbolProvider,
     private val firResolveSession: LLFirResolveSession,
     private val firSymbolBuilder: KtSymbolByFirBuilder,
-    override val token: KtLifetimeToken,
+) : KtSymbolProvider() {
-) : KtSymbolProvider(), KtLifetimeOwner {
-    private val firSymbolProvider by weakRef(firSymbolProvider)
 
-    override fun getParameterSymbol(psi: KtParameter): KtVariableLikeSymbol = withValidityAssertion {
+    override fun getParameterSymbol(psi: KtParameter): KtVariableLikeSymbol {
-        when {
+        return when {
             psi.isFunctionTypeParameter -> error(
                 "Creating KtValueParameterSymbol for function type parameter is not possible. Please see the KDoc of getParameterSymbol"
             )
 
             psi.isLoopParameter -> {
-                firSymbolBuilder.variableLikeBuilder.buildLocalVariableSymbol(psi.resolveToFirSymbolOfType<FirPropertySymbol>(firResolveSession))
+                firSymbolBuilder.variableLikeBuilder.buildLocalVariableSymbol(
+                    psi.resolveToFirSymbolOfType<FirPropertySymbol>(
+                        firResolveSession
+                    )
+                )
             }
 
             else -> {
@@ -51,12 +49,12 @@ internal class KtFirSymbolProvider(
         }
     }
 
-    override fun getFileSymbol(psi: KtFile): KtFileSymbol = withValidityAssertion {
+    override fun getFileSymbol(psi: KtFile): KtFileSymbol {
-        firSymbolBuilder.buildFileSymbol(psi.getOrBuildFirFile(firResolveSession).symbol)
+        return firSymbolBuilder.buildFileSymbol(psi.getOrBuildFirFile(firResolveSession).symbol)
     }
 
-    override fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol = withValidityAssertion {
+    override fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol {
-        when (val firSymbol = psi.resolveToFirSymbolOfType<FirFunctionSymbol<*>>(firResolveSession)) {
+        return when (val firSymbol = psi.resolveToFirSymbolOfType<FirFunctionSymbol<*>>(firResolveSession)) {
             is FirNamedFunctionSymbol -> {
                 if (firSymbol.origin == FirDeclarationOrigin.SamConstructor) {
                     firSymbolBuilder.functionLikeBuilder.buildSamConstructorSymbol(firSymbol)
@@ -69,70 +67,80 @@ internal class KtFirSymbolProvider(
         }
     }
 
-    override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol = withValidityAssertion {
+    override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol {
-
+        return firSymbolBuilder.functionLikeBuilder.buildConstructorSymbol(
-        firSymbolBuilder.functionLikeBuilder.buildConstructorSymbol(psi.resolveToFirSymbolOfType<FirConstructorSymbol>(firResolveSession))
+            psi.resolveToFirSymbolOfType<FirConstructorSymbol>(
-
+                firResolveSession
+            )
+        )
     }
 
-    override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol = withValidityAssertion {
+    override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol {
-        firSymbolBuilder.classifierBuilder.buildTypeParameterSymbol(psi.resolveToFirSymbolOfType<FirTypeParameterSymbol>(firResolveSession))
+        return firSymbolBuilder.classifierBuilder.buildTypeParameterSymbol(
+            psi.resolveToFirSymbolOfType<FirTypeParameterSymbol>(
+                firResolveSession
+            )
+        )
     }
 
-    override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol = withValidityAssertion {
+    override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol {
-        firSymbolBuilder.classifierBuilder.buildTypeAliasSymbol(psi.resolveToFirSymbolOfType<FirTypeAliasSymbol>(firResolveSession))
+        return firSymbolBuilder.classifierBuilder.buildTypeAliasSymbol(psi.resolveToFirSymbolOfType<FirTypeAliasSymbol>(firResolveSession))
     }
 
-    override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol = withValidityAssertion {
+    override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol {
-        firSymbolBuilder.buildEnumEntrySymbol(psi.resolveToFirSymbolOfType<FirEnumEntrySymbol>(firResolveSession))
+        return firSymbolBuilder.buildEnumEntrySymbol(psi.resolveToFirSymbolOfType<FirEnumEntrySymbol>(firResolveSession))
     }
 
-    override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol = withValidityAssertion {
+    override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol {
-        firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(
+        return firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(
             psi.resolveToFirSymbolOfType<FirAnonymousFunctionSymbol>(firResolveSession)
         )
     }
 
-    override fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol = withValidityAssertion {
+    override fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol {
-        firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(
+        return firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(
             psi.resolveToFirSymbolOfType<FirAnonymousFunctionSymbol>(firResolveSession)
         )
     }
 
-    override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol = withValidityAssertion {
+    override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol {
-        firSymbolBuilder.variableLikeBuilder.buildVariableSymbol(psi.resolveToFirSymbolOfType<FirPropertySymbol>(firResolveSession))
+        return firSymbolBuilder.variableLikeBuilder.buildVariableSymbol(psi.resolveToFirSymbolOfType<FirPropertySymbol>(firResolveSession))
     }
 
-    override fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol = withValidityAssertion {
+    override fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol {
-        firSymbolBuilder.classifierBuilder.buildAnonymousObjectSymbol(
+        return firSymbolBuilder.classifierBuilder.buildAnonymousObjectSymbol(
             psi.objectDeclaration.resolveToFirSymbolOfType<FirAnonymousObjectSymbol>(firResolveSession)
         )
     }
 
-    override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol = withValidityAssertion {
+    override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol {
-        firSymbolBuilder.classifierBuilder.buildClassOrObjectSymbol(psi.resolveToFirSymbolOfType<FirClassSymbol<*>>(firResolveSession))
+        return firSymbolBuilder.classifierBuilder.buildClassOrObjectSymbol(psi.resolveToFirSymbolOfType<FirClassSymbol<*>>(firResolveSession))
     }
 
-    override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol? = withValidityAssertion {
+    override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol? {
         require(psi !is KtObjectDeclaration || psi.parent !is KtObjectLiteralExpression)
         // A KtClassOrObject may also map to an FirEnumEntry. Hence, we need to return null in this case.
         if (psi is KtEnumEntry) return null
-        firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(psi.resolveToFirSymbolOfType<FirRegularClassSymbol>(firResolveSession))
+        return firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(
+            psi.resolveToFirSymbolOfType<FirRegularClassSymbol>(firResolveSession)
+        )
     }
 
-    override fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol = withValidityAssertion {
+    override fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol {
-        firSymbolBuilder.callableBuilder.buildPropertyAccessorSymbol(psi.resolveToFirSymbolOfType<FirPropertyAccessorSymbol>(firResolveSession))
+        return firSymbolBuilder.callableBuilder.buildPropertyAccessorSymbol(
+            psi.resolveToFirSymbolOfType<FirPropertyAccessorSymbol>(firResolveSession)
+        )
     }
 
-    override fun getClassInitializerSymbol(psi: KtClassInitializer): KtClassInitializerSymbol = withValidityAssertion {
+    override fun getClassInitializerSymbol(psi: KtClassInitializer): KtClassInitializerSymbol {
-        firSymbolBuilder.anonymousInitializerBuilder.buildClassInitializer(
+        return firSymbolBuilder.anonymousInitializerBuilder.buildClassInitializer(
             psi.resolveToFirSymbolOfType<FirAnonymousInitializerSymbol>(firResolveSession)
         )
     }
 
-    override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? = withValidityAssertion {
+    override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? {
         val symbol = firSymbolProvider.getClassLikeSymbolByClassId(classId) as? FirRegularClassSymbol ?: return null
-        firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(symbol)
+        return firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(symbol)
     }
 
     override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): Sequence<KtSymbol> {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSyntheticJavaPropertySymbol.kt

@@ -65,12 +65,12 @@ internal class KtFirSyntheticJavaPropertySymbol(
             builder.callableBuilder.buildGetterSymbol(firSymbol.getterSymbol!!)
         }
     override val javaGetterSymbol: KtFunctionSymbol
-        get() {
+        get() = withValidityAssertion {
             val fir = firSymbol.fir as FirSyntheticProperty
             return builder.functionLikeBuilder.buildFunctionSymbol(fir.getter.delegate.symbol)
         }
     override val javaSetterSymbol: KtFunctionSymbol?
-        get() {
+        get() = withValidityAssertion {
             val fir = firSymbol.fir as FirSyntheticProperty
             return fir.setter?.delegate?.let { builder.functionLikeBuilder.buildFunctionSymbol(it.symbol) }
         }
@@ -88,7 +88,7 @@ internal class KtFirSyntheticJavaPropertySymbol(
 
     override val origin: KtSymbolOrigin get() = withValidityAssertion { KtSymbolOrigin.JAVA_SYNTHETIC_PROPERTY }
 
-    override fun createPointer(): KtSymbolPointer<KtSyntheticJavaPropertySymbol> {
+    override fun createPointer(): KtSymbolPointer<KtSyntheticJavaPropertySymbol> = withValidityAssertion {
         val containingClassId = firSymbol.containingClass()?.classId
             ?: error("Cannot find parent class for synthetic java property $callableIdIfNonLocal")
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirTypeAliasSymbol.kt

@@ -52,7 +52,7 @@ internal class KtFirTypeAliasSymbol(
         KtFirAnnotationListForDeclaration.create(firSymbol, firResolveSession.useSiteFirSession, token)
     }
 
-    override fun createPointer(): KtSymbolPointer<KtTypeAliasSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtTypeAliasSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         TODO("Creating symbols for library typealiases is not supported yet")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirTypeParameterSymbol.kt

@@ -44,7 +44,7 @@ internal class KtFirTypeParameterSymbol(
     override val variance: Variance get() = withValidityAssertion { firSymbol.variance }
     override val isReified: Boolean get() = withValidityAssertion { firSymbol.isReified }
 
-    override fun createPointer(): KtSymbolPointer<KtTypeParameterSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtTypeParameterSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         TODO("Creating symbols for library type parameters is not supported yet")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirValueParameterSymbol.kt

@@ -66,7 +66,7 @@ internal class KtFirValueParameterSymbol(
         ktPropertySymbol
     }
 
-    override fun createPointer(): KtSymbolPointer<KtValueParameterSymbol> {
+    override fun createPointer(): KtSymbolPointer<KtValueParameterSymbol> = withValidityAssertion {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         TODO("Creating pointers for functions parameters from library is not supported yet")
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/firSymbolUtils.kt

@@ -5,9 +5,13 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.symbols
 
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.KtInitializerValue
 import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.analysis.api.fir.utils.asKtInitializerValue
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtValueParameterSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getElementTextInContext
@@ -69,4 +73,4 @@ internal fun FirVariableSymbol<*>.getKtConstantInitializer(): KtInitializerValue
     ensureResolved(FirResolvePhase.BODY_RESOLVE)
     val firInitializer = fir.initializer ?: return null
     return firInitializer.asKtInitializerValue()
-}
+}

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/types/KtFirSubstitutor.kt

@@ -41,7 +41,7 @@ internal class KtFirMapBackedSubstitutor(
     builder: KtSymbolByFirBuilder,
     token: KtLifetimeToken
 ) : AbstractKtFirSubstitutor<ConeSubstitutorByMap>(_substitutor, builder, token), KtMapBackedSubstitutor {
-    override fun getAsMap(): Map<KtTypeParameterSymbol, KtType> {
+    override fun getAsMap(): Map<KtTypeParameterSymbol, KtType> = withValidityAssertion {
         val result = mutableMapOf<KtTypeParameterSymbol, KtType>()
         for ((typeParameter, type) in substitutor.substitution) {
             val typeParameterSymbol = builderRef.classifierBuilder.buildTypeParameterSymbolByLookupTag(typeParameter.toLookupTag())

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/symbolDeclarationOverridesProvider/AbstractOverriddenDeclarationProviderTest.kt

@@ -58,7 +58,8 @@ abstract class AbstractOverriddenDeclarationProviderTest : AbstractAnalysisApiBa
         append(symbol.returnType.render(KtTypeRendererOptions.SHORT_NAMES))
     }
 
-    private fun getPath(symbol: KtCallableSymbol): String = when (symbol) {
+    @Suppress("unused")
+    private fun KtAnalysisSession.getPath(symbol: KtCallableSymbol): String = when (symbol) {
         is KtSyntheticJavaPropertySymbol -> symbol.callableIdIfNonLocal?.toString()!!
         else -> {
             val ktDeclaration = symbol.psi as KtDeclaration

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/symbols/AbstractSymbolTest.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.impl.base.test.cases.symbols
 
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.components.KtDeclarationRendererOptions
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiSingleFileTest
 import org.jetbrains.kotlin.analysis.api.symbols.DebugSymbolRenderer
 import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/KtAnalysisSession.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.psi.KtFile
  *
  * To create analysis session consider using [analyse]
  */
+@Suppress("AnalysisApiMissingLifetimeCheck")
 public abstract class KtAnalysisSession(final override val token: KtLifetimeToken) : KtLifetimeOwner,
     KtSmartCastProviderMixIn,
     KtCallResolverMixIn,

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/KtTypeArgument.kt

@@ -20,8 +20,10 @@ public class KtStarProjectionTypeArgument(override val token: KtLifetimeToken) :
 }
 
 public class KtTypeArgumentWithVariance(
-    override val type: KtType,
+    private val _type: KtType,
     public val variance: Variance,
     override val token: KtLifetimeToken,
-) : KtTypeArgument()
+) : KtTypeArgument() {
+    override val type: KtType get() = withValidityAssertion { _type }
+}
 

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtAnalysisScopeProvider.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.api.components
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 
 public abstract class KtAnalysisScopeProvider : KtAnalysisSessionComponent() {
     public abstract fun getAnalysisScope(): GlobalSearchScope
@@ -21,7 +22,7 @@ public interface KtAnalysisScopeProviderMixIn : KtAnalysisSessionMixIn {
      * That means [org.jetbrains.kotlin.analysis.api.symbols.KtSymbol] can be built for the declarations from this scope.
      */
     public val analysisScope: GlobalSearchScope
-        get() = analysisSession.analysisScopeProvider.getAnalysisScope()
+        get() = withValidityAssertion { analysisSession.analysisScopeProvider.getAnalysisScope() }
 
 
     /**
@@ -31,5 +32,5 @@ public interface KtAnalysisScopeProviderMixIn : KtAnalysisSessionMixIn {
      * @see analysisScope
      */
     public fun PsiElement.canBeAnalysed(): Boolean =
-        analysisSession.analysisScopeProvider.canBeAnalysed(this)
+        withValidityAssertion { analysisSession.analysisScopeProvider.canBeAnalysed(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtAnalysisSessionComponent.kt

@@ -7,8 +7,10 @@ package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 
-public abstract class KtAnalysisSessionComponent : KtLifetimeOwner {
+public abstract class KtAnalysisSessionComponent {
     protected abstract val analysisSession: KtAnalysisSession
+    protected open val token: KtLifetimeToken get() = analysisSession.token
 }
 

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtAnalysisSessionMixIn.kt

@@ -6,7 +6,8 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
 
-public interface KtAnalysisSessionMixIn {
+public interface KtAnalysisSessionMixIn : KtLifetimeOwner {
     public val analysisSession: KtAnalysisSession
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtCallResolver.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.calls.KtCallCandidateInfo
 import org.jetbrains.kotlin.analysis.api.calls.KtCallInfo
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtArrayAccessExpression
 import org.jetbrains.kotlin.psi.KtCallElement
 import org.jetbrains.kotlin.psi.KtElement
@@ -20,17 +21,23 @@ public abstract class KtCallResolver : KtAnalysisSessionComponent() {
 public interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
 
     public fun KtElement.resolveCall(): KtCallInfo? =
-        analysisSession.callResolver.resolveCall(this)
+        withValidityAssertion { withValidityAssertion { analysisSession.callResolver.resolveCall(this) } }
 
     public fun KtCallElement.resolveCall(): KtCallInfo =
-        analysisSession.callResolver.resolveCall(this)
+        withValidityAssertion {
-            ?: error("KtCallElement should always resolve to a KtCallInfo")
+            analysisSession.callResolver.resolveCall(this)
+                ?: error("KtCallElement should always resolve to a KtCallInfo")
+        }
 
     public fun KtUnaryExpression.resolveCall(): KtCallInfo =
-        analysisSession.callResolver.resolveCall(this) ?: error("KtUnaryExpression should always resolve to a KtCallInfo")
+        withValidityAssertion {
+            analysisSession.callResolver.resolveCall(this) ?: error("KtUnaryExpression should always resolve to a KtCallInfo")
+        }
 
     public fun KtArrayAccessExpression.resolveCall(): KtCallInfo =
-        analysisSession.callResolver.resolveCall(this) ?: error("KtArrayAccessExpression should always resolve to a KtCallInfo")
+        withValidityAssertion {
+            analysisSession.callResolver.resolveCall(this) ?: error("KtArrayAccessExpression should always resolve to a KtCallInfo")
+        }
 
     /**
      * Returns all the candidates considered during [overload resolution](https://kotlinlang.org/spec/overload-resolution.html) for the call
@@ -40,5 +47,5 @@ public interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
      * applicability and choosing the most specific candidate.
      */
     public fun KtElement.collectCallCandidates(): List<KtCallCandidateInfo> =
-        analysisSession.callResolver.collectCallCandidates(this)
+        withValidityAssertion { analysisSession.callResolver.collectCallCandidates(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtCompileTimeConstantProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.base.KtConstantValue
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtExpression
 
 public enum class KtConstantEvaluationMode {
@@ -35,5 +36,5 @@ public abstract class KtCompileTimeConstantProvider : KtAnalysisSessionComponent
 
 public interface KtCompileTimeConstantProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtExpression.evaluate(mode: KtConstantEvaluationMode): KtConstantValue? =
-        analysisSession.compileTimeConstantProvider.evaluate(this, mode)
+        withValidityAssertion { analysisSession.compileTimeConstantProvider.evaluate(this, mode) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtCompletionCandidateChecker.kt

@@ -5,6 +5,9 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
 import org.jetbrains.kotlin.psi.KtExpression
@@ -20,20 +23,34 @@ public abstract class KtCompletionCandidateChecker : KtAnalysisSessionComponent(
     ): KtExtensionApplicabilityResult
 }
 
-public sealed class KtExtensionApplicabilityResult {
+
+public sealed class KtExtensionApplicabilityResult : KtLifetimeOwner {
     public abstract val isApplicable: Boolean
     public abstract val substitutor: KtSubstitutor
 
-    public class ApplicableAsExtensionCallable(override val substitutor: KtSubstitutor) : KtExtensionApplicabilityResult() {
+    public class ApplicableAsExtensionCallable(
-        override val isApplicable: Boolean get() = true
+        private val _substitutor: KtSubstitutor,
+        override val token: KtLifetimeToken
+    ) : KtExtensionApplicabilityResult() {
+        override val substitutor: KtSubstitutor = withValidityAssertion { _substitutor }
+        override val isApplicable: Boolean get() = withValidityAssertion { true }
     }
 
-    public class ApplicableAsFunctionalVariableCall(override val substitutor: KtSubstitutor) : KtExtensionApplicabilityResult() {
+    public class ApplicableAsFunctionalVariableCall(
-        override val isApplicable: Boolean get() = true
+        private val _substitutor: KtSubstitutor,
+        override val token: KtLifetimeToken
+    ) : KtExtensionApplicabilityResult() {
+        override val substitutor: KtSubstitutor get() = withValidityAssertion { _substitutor }
+        override val isApplicable: Boolean get() = withValidityAssertion { true }
     }
 
-    public class NonApplicable(override val substitutor: KtSubstitutor) : KtExtensionApplicabilityResult() {
+
-        override val isApplicable: Boolean = false
+    public class NonApplicable(
+        private val _substitutor: KtSubstitutor,
+        override val token: KtLifetimeToken
+    ) : KtExtensionApplicabilityResult() {
+        override val substitutor: KtSubstitutor = withValidityAssertion { _substitutor }
+        override val isApplicable: Boolean get() = withValidityAssertion { false }
     }
 }
 
@@ -42,11 +59,12 @@ public interface KtCompletionCandidateCheckerMixIn : KtAnalysisSessionMixIn {
         originalPsiFile: KtFile,
         psiFakeCompletionExpression: KtSimpleNameExpression,
         psiReceiverExpression: KtExpression?,
-    ): KtExtensionApplicabilityResult =
+    ): KtExtensionApplicabilityResult = withValidityAssertion {
         analysisSession.completionCandidateChecker.checkExtensionFitsCandidate(
             this,
             originalPsiFile,
             psiFakeCompletionExpression,
             psiReceiverExpression
         )
+    }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtDiagnosticProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnosticWithPsi
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -16,10 +17,10 @@ public abstract class KtDiagnosticProvider : KtAnalysisSessionComponent() {
 
 public interface KtDiagnosticProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtElement.getDiagnostics(filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>> =
-        analysisSession.diagnosticProvider.getDiagnosticsForElement(this, filter)
+        withValidityAssertion { analysisSession.diagnosticProvider.getDiagnosticsForElement(this, filter) }
 
     public fun KtFile.collectDiagnosticsForFile(filter: KtDiagnosticCheckerFilter): Collection<KtDiagnosticWithPsi<*>> =
-        analysisSession.diagnosticProvider.collectDiagnosticsForFile(this, filter)
+        withValidityAssertion { analysisSession.diagnosticProvider.collectDiagnosticsForFile(this, filter) }
 }
 
 public enum class KtDiagnosticCheckerFilter {

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtExpressionInfoProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.diagnostics.WhenMissingCase
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.psi.KtReturnExpression
@@ -17,7 +18,8 @@ public abstract class KtExpressionInfoProvider : KtAnalysisSessionComponent() {
 
 public interface KtExpressionInfoProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtReturnExpression.getReturnTargetSymbol(): KtCallableSymbol? =
-        analysisSession.expressionInfoProvider.getReturnExpressionTargetSymbol(this)
+        withValidityAssertion { analysisSession.expressionInfoProvider.getReturnExpressionTargetSymbol(this) }
 
-    public fun KtWhenExpression.getMissingCases(): List<WhenMissingCase> = analysisSession.expressionInfoProvider.getWhenMissingCases(this)
+    public fun KtWhenExpression.getMissingCases(): List<WhenMissingCase> =
+        withValidityAssertion { analysisSession.expressionInfoProvider.getWhenMissingCases(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtExpressionTypeProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
@@ -31,13 +32,13 @@ public interface KtExpressionTypeProviderMixIn : KtAnalysisSessionMixIn {
      * - `Unit` type for statements;
      */
     public fun KtExpression.getKtType(): KtType? =
-        analysisSession.expressionTypeProvider.getKtExpressionType(this)
+        withValidityAssertion { analysisSession.expressionTypeProvider.getKtExpressionType(this) }
 
     /**
      * Returns the return type of the given [KtDeclaration] as [KtType]
      */
     public fun KtDeclaration.getReturnKtType(): KtType =
-        analysisSession.expressionTypeProvider.getReturnTypeForKtDeclaration(this)
+        withValidityAssertion { analysisSession.expressionTypeProvider.getReturnTypeForKtDeclaration(this) }
 
     /**
      * Returns the functional type of the given [KtFunction].
@@ -50,14 +51,14 @@ public interface KtExpressionTypeProviderMixIn : KtAnalysisSessionMixIn {
      * such as `SuspendFunction`, `KFunction`, or `KSuspendFunction`, will be constructed.
      */
     public fun KtFunction.getFunctionalType(): KtType =
-        analysisSession.expressionTypeProvider.getFunctionalTypeForKtFunction(this)
+        withValidityAssertion { analysisSession.expressionTypeProvider.getFunctionalTypeForKtFunction(this) }
 
     /**
      * Returns the expected [KtType] of this [PsiElement] if it is an expression. The returned value should not be a
      * [org.jetbrains.kotlin.analysis.api.types.KtClassErrorType].
      */
     public fun PsiElement.getExpectedType(): KtType? =
-        analysisSession.expressionTypeProvider.getExpectedType(this)
+        withValidityAssertion { analysisSession.expressionTypeProvider.getExpectedType(this) }
 
     /**
      * Returns `true` if this expression is definitely null, based on declared nullability and smart cast types derived from
@@ -82,11 +83,11 @@ public interface KtExpressionTypeProviderMixIn : KtAnalysisSessionMixIn {
      * [spec](https://kotlinlang.org/spec/type-inference.html#smart-cast-sink-stability) provides an explanation on smart cast stability.
      */
     public fun KtExpression.isDefinitelyNull(): Boolean =
-        analysisSession.expressionTypeProvider.isDefinitelyNull(this)
+        withValidityAssertion { analysisSession.expressionTypeProvider.isDefinitelyNull(this) }
 
     /**
      * Returns `true` if this expression is definitely not null. See [isDefinitelyNull] for examples.
      */
     public fun KtExpression.isDefinitelyNotNull(): Boolean =
-        analysisSession.expressionTypeProvider.isDefinitelyNotNull(this)
+        withValidityAssertion { analysisSession.expressionTypeProvider.isDefinitelyNotNull(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtImportOptimizer.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtImportDirective
 
@@ -21,7 +22,7 @@ public interface KtImportOptimizerMixIn : KtAnalysisSessionMixIn {
      *
      * Does **not** change the file.
      */
-    public fun analyseImports(file: KtFile): KtImportOptimizerResult {
+    public fun analyseImports(file: KtFile): KtImportOptimizerResult = withValidityAssertion {
         return analysisSession.importOptimizer.analyseImports(file)
     }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtInheritorsProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtNamedClassOrObjectSymbol
 
@@ -15,8 +16,8 @@ public abstract class KtInheritorsProvider : KtAnalysisSessionComponent() {
 
 public interface KtInheritorsProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtNamedClassOrObjectSymbol.getSealedClassInheritors(): List<KtNamedClassOrObjectSymbol> =
-        analysisSession.inheritorsProvider.getInheritorsOfSealedClass(this)
+        withValidityAssertion { analysisSession.inheritorsProvider.getInheritorsOfSealedClass(this) }
 
     public fun KtNamedClassOrObjectSymbol.getEnumEntries(): List<KtEnumEntrySymbol> =
-        analysisSession.inheritorsProvider.getEnumEntries(this)
+        withValidityAssertion { analysisSession.inheritorsProvider.getEnumEntries(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtJvmTypeMapper.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.org.objectweb.asm.Type
@@ -20,5 +21,5 @@ public interface KtJvmTypeMapperMixIn : KtAnalysisSessionMixIn {
      * @see TypeMappingMode
      */
     public fun KtType.mapTypeToJvmType(mode: TypeMappingMode = TypeMappingMode.DEFAULT): Type =
-        analysisSession.jvmTypeMapper.mapTypeToJvmType(this, mode)
+        withValidityAssertion { analysisSession.jvmTypeMapper.mapTypeToJvmType(this, mode) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtOverrideInfoProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.util.ImplementationStatus
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
@@ -24,14 +25,14 @@ public interface KtMemberSymbolProviderMixin : KtAnalysisSessionMixIn {
 
     /** Checks if the given symbol (possibly a symbol inherited from a super class) is visible in the given class. */
     public fun KtCallableSymbol.isVisibleInClass(classSymbol: KtClassOrObjectSymbol): Boolean =
-        analysisSession.overrideInfoProvider.isVisible(this, classSymbol)
+        withValidityAssertion { analysisSession.overrideInfoProvider.isVisible(this, classSymbol) }
 
     /**
      * Gets the [ImplementationStatus] of the [this] member symbol in the given [parentClassSymbol]. Or null if this symbol is not a
      * member.
      */
     public fun KtCallableSymbol.getImplementationStatus(parentClassSymbol: KtClassOrObjectSymbol): ImplementationStatus? =
-        analysisSession.overrideInfoProvider.getImplementationStatus(this, parentClassSymbol)
+        withValidityAssertion { analysisSession.overrideInfoProvider.getImplementationStatus(this, parentClassSymbol) }
 
     /**
      * Gets the original symbol for the given callable symbol. In a class scope, a symbol may be derived from symbols declared in super
@@ -54,11 +55,11 @@ public interface KtMemberSymbolProviderMixin : KtAnalysisSessionMixIn {
      * after specialization) and delegation.
      */
     public val KtCallableSymbol.originalOverriddenSymbol: KtCallableSymbol?
-        get() = analysisSession.overrideInfoProvider.getOriginalOverriddenSymbol(this)
+        get() = withValidityAssertion { analysisSession.overrideInfoProvider.getOriginalOverriddenSymbol(this) }
 
     /**
      * Gets the class symbol where the given callable symbol is declared. See [originalOverriddenSymbol] for more details.
      */
     public val KtCallableSymbol.originalContainingClassForOverride: KtClassOrObjectSymbol?
-        get() = analysisSession.overrideInfoProvider.getOriginalContainingClassForOverride(this)
+        get() = withValidityAssertion { analysisSession.overrideInfoProvider.getOriginalContainingClassForOverride(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtPsiTypeProvider.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.components
 
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeMappingMode
 
@@ -41,6 +42,6 @@ public interface KtPsiTypeProviderMixIn : KtAnalysisSessionMixIn {
         mode: KtTypeMappingMode = KtTypeMappingMode.DEFAULT,
         isAnnotationMethod: Boolean = false,
     ): PsiType? =
-        analysisSession.psiTypeProvider.asPsiType(this, useSitePosition, mode, isAnnotationMethod)
+        withValidityAssertion { analysisSession.psiTypeProvider.asPsiType(this, useSitePosition, mode, isAnnotationMethod) }
 
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtReferenceResolveMixIn.kt

@@ -5,15 +5,16 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.references.KtReference
 
 public interface KtReferenceResolveMixIn : KtAnalysisSessionMixIn {
-    public fun KtReference.resolveToSymbols(): Collection<KtSymbol> {
+    public fun KtReference.resolveToSymbols(): Collection<KtSymbol> = withValidityAssertion {
         return analysisSession.referenceResolveProvider.resolveToSymbols(this)
     }
 
-    public fun KtReference.resolveToSymbol(): KtSymbol? {
+    public fun KtReference.resolveToSymbol(): KtSymbol? = withValidityAssertion {
         return resolveToSymbols().singleOrNull()
     }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtReferenceShortener.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.api.components
 import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.analysis.api.components.ShortenOption.Companion.defaultCallableShortenOption
 import org.jetbrains.kotlin.analysis.api.components.ShortenOption.Companion.defaultClassShortenOption
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
@@ -64,19 +65,28 @@ public interface KtReferenceShortenerMixIn : KtAnalysisSessionMixIn {
         classShortenOption: (KtClassLikeSymbol) -> ShortenOption = defaultClassShortenOption,
         callableShortenOption: (KtCallableSymbol) -> ShortenOption = defaultCallableShortenOption
     ): ShortenCommand =
-        analysisSession.referenceShortener.collectShortenings(file, selection, classShortenOption, callableShortenOption)
+        withValidityAssertion {
+            analysisSession.referenceShortener.collectShortenings(
+                file,
+                selection,
+                classShortenOption,
+                callableShortenOption
+            )
+        }
 
     public fun collectPossibleReferenceShorteningsInElement(
         element: KtElement,
         classShortenOption: (KtClassLikeSymbol) -> ShortenOption = defaultClassShortenOption,
         callableShortenOption: (KtCallableSymbol) -> ShortenOption = defaultCallableShortenOption
     ): ShortenCommand =
-        analysisSession.referenceShortener.collectShortenings(
+        withValidityAssertion {
-            element.containingKtFile,
+            analysisSession.referenceShortener.collectShortenings(
-            element.textRange,
+                element.containingKtFile,
-            classShortenOption,
+                element.textRange,
-            callableShortenOption
+                classShortenOption,
-        )
+                callableShortenOption
+            )
+        }
 }
 
 public interface ShortenCommand {

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSamResolver.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSamConstructorSymbol
 import org.jetbrains.kotlin.name.ClassId
@@ -18,5 +19,5 @@ public interface KtSamResolverMixIn : KtAnalysisSessionMixIn {
      * Returns [KtSamConstructorSymbol] if the given [KtClassLikeSymbol] is a functional interface type, a.k.a. SAM.
      */
     public fun KtClassLikeSymbol.getSamConstructor(): KtSamConstructorSymbol? =
-        analysisSession.samResolver.getSamConstructor(this)
+        withValidityAssertion { analysisSession.samResolver.getSamConstructor(this) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtScopeProvider.kt

@@ -6,12 +6,14 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
-import org.jetbrains.kotlin.analysis.api.scopes.*
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.scopes.KtScope
+import org.jetbrains.kotlin.analysis.api.scopes.KtTypeScope
 import org.jetbrains.kotlin.analysis.api.symbols.KtFileSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
-import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -33,7 +35,7 @@ public abstract class KtScopeProvider : KtAnalysisSessionComponent() {
 
     public abstract fun getCompositeScope(subScopes: List<KtScope>): KtScope
 
-    public abstract fun getTypeScope(type: KtType): KtScope?
+    public abstract fun getTypeScope(type: KtType): KtTypeScope?
 
     public abstract fun getScopeContextForPosition(
         originalFile: KtFile,
@@ -43,40 +45,50 @@ public abstract class KtScopeProvider : KtAnalysisSessionComponent() {
 
 public interface KtScopeProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtSymbolWithMembers.getMemberScope(): KtScope =
-        analysisSession.scopeProvider.getMemberScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getMemberScope(this) }
 
     public fun KtSymbolWithMembers.getDeclaredMemberScope(): KtScope =
-        analysisSession.scopeProvider.getDeclaredMemberScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getDeclaredMemberScope(this) }
 
     public fun KtSymbolWithMembers.getDelegatedMemberScope(): KtScope =
-        analysisSession.scopeProvider.getDelegatedMemberScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getDelegatedMemberScope(this) }
 
     public fun KtSymbolWithMembers.getStaticMemberScope(): KtScope =
-        analysisSession.scopeProvider.getStaticMemberScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getStaticMemberScope(this) }
 
     public fun KtFileSymbol.getFileScope(): KtScope =
-        analysisSession.scopeProvider.getFileScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getFileScope(this) }
 
     public fun KtPackageSymbol.getPackageScope(): KtScope =
-        analysisSession.scopeProvider.getPackageScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getPackageScope(this) }
 
     public fun List<KtScope>.asCompositeScope(): KtScope =
-        analysisSession.scopeProvider.getCompositeScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getCompositeScope(this) }
 
-    public fun KtType.getTypeScope(): KtScope? =
+    public fun KtType.getTypeScope(): KtTypeScope? =
-        analysisSession.scopeProvider.getTypeScope(this)
+        withValidityAssertion { analysisSession.scopeProvider.getTypeScope(this) }
 
     public fun KtFile.getScopeContextForPosition(positionInFakeFile: KtElement): KtScopeContext =
-        analysisSession.scopeProvider.getScopeContextForPosition(this, positionInFakeFile)
+        withValidityAssertion { analysisSession.scopeProvider.getScopeContextForPosition(this, positionInFakeFile) }
 
     public fun KtFile.getScopeContextForFile(): KtScopeContext =
-        analysisSession.scopeProvider.getScopeContextForPosition(this, this)
+        withValidityAssertion { analysisSession.scopeProvider.getScopeContextForPosition(this, this) }
 }
 
-public data class KtScopeContext(val scopes: KtScope, val implicitReceivers: List<KtImplicitReceiver>)
+public class KtScopeContext(
+    private val _scopes: KtScope,
+    private val _implicitReceivers: List<KtImplicitReceiver>,
+    override val token: KtLifetimeToken
+) : KtLifetimeOwner {
+    public val implicitReceivers: List<KtImplicitReceiver> get() = withValidityAssertion { _implicitReceivers }
+    public val scopes: KtScope get() = withValidityAssertion { _scopes }
+}
 
 public class KtImplicitReceiver(
     override val token: KtLifetimeToken,
-    public val type: KtType,
+    private val _type: KtType,
-    public val ownerSymbol: KtSymbol
+    private val _ownerSymbol: KtSymbol
-) : KtLifetimeOwner
+) : KtLifetimeOwner {
+    public val ownerSymbol: KtSymbol get() = withValidityAssertion { _ownerSymbol }
+    public val type: KtType get() = withValidityAssertion { _type }
+}

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSmartCastProvider.kt

@@ -21,7 +21,7 @@ public interface KtSmartCastProviderMixIn : KtAnalysisSessionMixIn {
      * Gets the smart-cast information of the given expression or null if the expression is not smart casted.
      */
     public fun KtExpression.getSmartCastInfo(): KtSmartCastInfo? =
-        analysisSession.smartCastProvider.getSmartCastedInfo(this)
+        withValidityAssertion { analysisSession.smartCastProvider.getSmartCastedInfo(this) }
 
     /**
      * Returns the list of implicit smart-casts which are required for the expression to be called. Includes only implicit
@@ -36,7 +36,7 @@ public interface KtSmartCastProviderMixIn : KtAnalysisSessionMixIn {
      * ```
      */
     public fun KtExpression.getImplicitReceiverSmartCast(): Collection<KtImplicitReceiverSmartCast> =
-        analysisSession.smartCastProvider.getImplicitReceiverSmartCast(this)
+        withValidityAssertion { analysisSession.smartCastProvider.getImplicitReceiverSmartCast(this) }
 }
 
 public data class KtSmartCastInfo(

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSubtypingComponent.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.types.KtType
 
 public abstract class KtSubtypingComponent : KtAnalysisSessionComponent() {
@@ -14,11 +15,11 @@ public abstract class KtSubtypingComponent : KtAnalysisSessionComponent() {
 
 public interface KtSubtypingComponentMixIn : KtAnalysisSessionMixIn {
     infix public fun KtType.isEqualTo(other: KtType): Boolean =
-        analysisSession.subtypingComponent.isEqualTo(this, other)
+        withValidityAssertion { analysisSession.subtypingComponent.isEqualTo(this, other) }
 
     infix public fun KtType.isSubTypeOf(superType: KtType): Boolean =
-        analysisSession.subtypingComponent.isSubTypeOf(this, superType)
+        withValidityAssertion { analysisSession.subtypingComponent.isSubTypeOf(this, superType) }
 
     infix public fun KtType.isNotSubTypeOf(superType: KtType): Boolean =
-        !analysisSession.subtypingComponent.isSubTypeOf(this, superType)
+        withValidityAssertion { !analysisSession.subtypingComponent.isSubTypeOf(this, superType) }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolContainingDeclarationProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -23,8 +24,8 @@ public interface KtSymbolContainingDeclarationProviderMixIn : KtAnalysisSessionM
      *   for local declaration returns declaration it was declared it
      */
     public fun KtSymbol.getContainingSymbol(): KtSymbolWithKind? =
-        analysisSession.containingDeclarationProvider.getContainingDeclaration(this)
+        withValidityAssertion { analysisSession.containingDeclarationProvider.getContainingDeclaration(this) }
 
     public fun KtSymbol.getContainingModule(): KtModule =
-        analysisSession.containingDeclarationProvider.getContainingModule(this)
+        withValidityAssertion { analysisSession.containingDeclarationProvider.getContainingModule(this) }
 }