FIR IDE: add validity token to data classes related to smartcast API

Also move KtImplicitReceiverSmartCast closer to where it's used.

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

@@ -5,10 +5,10 @@
 
 package org.jetbrains.kotlin.analysis.api.descriptors.components
 
-import org.jetbrains.kotlin.analysis.api.ImplicitReceiverSmartCast
-import org.jetbrains.kotlin.analysis.api.ImplicitReceiverSmartcastKind
+import org.jetbrains.kotlin.analysis.api.components.KtImplicitReceiverSmartCast
+import org.jetbrains.kotlin.analysis.api.components.KtImplicitReceiverSmartCastKind
+import org.jetbrains.kotlin.analysis.api.components.KtSmartCastInfo
 import org.jetbrains.kotlin.analysis.api.components.KtSmartCastProvider
-import org.jetbrains.kotlin.analysis.api.components.SmartCastInfo
 import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtType
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
@@ -24,7 +24,7 @@ internal class KtFe10SmartCastProvider(override val analysisSession: KtFe10Analy
     override val token: ValidityToken
         get() = analysisSession.token
 
-    override fun getSmartCastedInfo(expression: KtExpression): SmartCastInfo? {
+    override fun getSmartCastedInfo(expression: KtExpression): KtSmartCastInfo? {
         withValidityAssertion {
             val bindingContext = analysisSession.analyze(expression)
             val stableSmartCasts = bindingContext[BindingContext.SMARTCAST, expression]
@@ -33,11 +33,11 @@ internal class KtFe10SmartCastProvider(override val analysisSession: KtFe10Analy
             return when {
                 stableSmartCasts != null -> {
                     val type = stableSmartCasts.getKtType() ?: return null
-                    SmartCastInfo(type, true)
+                    KtSmartCastInfo(type, true, token)
                 }
                 unstableSmartCasts != null -> {
                     val type = unstableSmartCasts.getKtType() ?: return null
-                    SmartCastInfo(type, false)
+                    KtSmartCastInfo(type, false, token)
                 }
                 else -> null
             }
@@ -51,13 +51,13 @@ internal class KtFe10SmartCastProvider(override val analysisSession: KtFe10Analy
         return defaultType?.toKtType(analysisSession)
     }
 
-    override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<ImplicitReceiverSmartCast> {
+    override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<KtImplicitReceiverSmartCast> {
         withValidityAssertion {
             val bindingContext = analysisSession.analyze(expression)
             val smartCasts = bindingContext[BindingContext.IMPLICIT_RECEIVER_SMARTCAST, expression] ?: return emptyList()
             return smartCasts.receiverTypes.map { (_, type) ->
-                val kind = ImplicitReceiverSmartcastKind.DISPATCH // TODO provide precise kind
-                ImplicitReceiverSmartCast(type.toKtType(analysisSession), kind)
+                val kind = KtImplicitReceiverSmartCastKind.DISPATCH // TODO provide precise kind
+                KtImplicitReceiverSmartCast(type.toKtType(analysisSession), kind, token)
             }
         }
     }

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

@@ -5,10 +5,10 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.components
 
-import org.jetbrains.kotlin.analysis.api.ImplicitReceiverSmartCast
-import org.jetbrains.kotlin.analysis.api.ImplicitReceiverSmartcastKind
+import org.jetbrains.kotlin.analysis.api.components.KtImplicitReceiverSmartCast
+import org.jetbrains.kotlin.analysis.api.components.KtImplicitReceiverSmartCastKind
+import org.jetbrains.kotlin.analysis.api.components.KtSmartCastInfo
 import org.jetbrains.kotlin.analysis.api.components.KtSmartCastProvider
-import org.jetbrains.kotlin.analysis.api.components.SmartCastInfo
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
@@ -24,18 +24,19 @@ internal class KtFirSmartcastProvider(
     override val analysisSession: KtFirAnalysisSession,
     override val token: ValidityToken,
 ) : KtSmartCastProvider(), KtFirAnalysisSessionComponent {
-    override fun getSmartCastedInfo(expression: KtExpression): SmartCastInfo? = withValidityAssertion {
+    override fun getSmartCastedInfo(expression: KtExpression): KtSmartCastInfo? = withValidityAssertion {
         val smartCastExpression =
             expression.getOrBuildFirSafe<FirExpressionWithSmartcast>(analysisSession.firResolveState) ?: return@withValidityAssertion null
-        SmartCastInfo(
+        KtSmartCastInfo(
             smartCastExpression.smartcastType.coneTypeSafe<ConeKotlinType>()?.asKtType() ?: return@withValidityAssertion null,
-            smartCastExpression.isStable
+            smartCastExpression.isStable,
+            token,
         )
 
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<ImplicitReceiverSmartCast> = withValidityAssertion {
+    override fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<KtImplicitReceiverSmartCast> = withValidityAssertion {
         val qualifiedExpression =
             expression.getOrBuildFirSafe<FirQualifiedAccessExpression>(analysisSession.firResolveState) ?: return emptyList()
         val dispatchReceiver = qualifiedExpression.dispatchReceiver
@@ -45,15 +46,17 @@ internal class KtFirSmartcastProvider(
         ) return emptyList()
         buildList {
             dispatchReceiver.takeIf { it.isStableSmartcast() }?.let { smartCasted ->
-                ImplicitReceiverSmartCast(
+                KtImplicitReceiverSmartCast(
                     smartCasted.typeRef.coneTypeSafe<ConeKotlinType>()?.asKtType() ?: return@let null,
-                    ImplicitReceiverSmartcastKind.DISPATCH
+                    KtImplicitReceiverSmartCastKind.DISPATCH,
+                    token,
                 )
             }?.let(::add)
             extensionReceiver.takeIf { it.isStableSmartcast() }?.let { smartCasted ->
-                ImplicitReceiverSmartCast(
+                KtImplicitReceiverSmartCast(
                     smartCasted.typeRef.coneTypeSafe<ConeKotlinType>()?.asKtType() ?: return@let null,
-                    ImplicitReceiverSmartcastKind.EXTENSION
+                    KtImplicitReceiverSmartCastKind.EXTENSION,
+                    token,
                 )
             }?.let(::add)
         }

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

@@ -1,14 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.analysis.api
-
-import org.jetbrains.kotlin.analysis.api.types.KtType
-
-public data class ImplicitReceiverSmartCast(val type: KtType, val kind: ImplicitReceiverSmartcastKind)
-
-public enum class ImplicitReceiverSmartcastKind {
-    DISPATCH, EXTENSION
-}

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

@@ -5,24 +5,46 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
-import org.jetbrains.kotlin.analysis.api.ImplicitReceiverSmartCast
+import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
+import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtExpression
 
 public abstract class KtSmartCastProvider : KtAnalysisSessionComponent() {
-    public abstract fun getSmartCastedInfo(expression: KtExpression): SmartCastInfo?
-    public abstract fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<ImplicitReceiverSmartCast>
+    public abstract fun getSmartCastedInfo(expression: KtExpression): KtSmartCastInfo?
+    public abstract fun getImplicitReceiverSmartCast(expression: KtExpression): Collection<KtImplicitReceiverSmartCast>
 }
 
-public class SmartCastInfo(public val smartCastType: KtType, public val isStable: Boolean)
-
 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(): SmartCastInfo? =
+    public fun KtExpression.getSmartCastInfo(): KtSmartCastInfo? =
         analysisSession.smartCastProvider.getSmartCastedInfo(this)
 
-    public fun KtExpression.getImplicitReceiverSmartCast(): Collection<ImplicitReceiverSmartCast> =
+    public fun KtExpression.getImplicitReceiverSmartCast(): Collection<KtImplicitReceiverSmartCast> =
         analysisSession.smartCastProvider.getImplicitReceiverSmartCast(this)
-}
+}
+
+public data class KtSmartCastInfo(
+    private val _smartCastType: KtType,
+    private val _isStable: Boolean,
+    override val token: ValidityToken
+) : ValidityTokenOwner {
+    public val isStable: Boolean get() = withValidityAssertion { _isStable }
+    public val smartCastType: KtType get() = withValidityAssertion { _smartCastType }
+}
+
+public data class KtImplicitReceiverSmartCast(
+    private val _type: KtType,
+    private val _kind: KtImplicitReceiverSmartCastKind,
+    override val token: ValidityToken
+) : ValidityTokenOwner {
+    public val type: KtType get() = withValidityAssertion { _type }
+    public val kind: KtImplicitReceiverSmartCastKind get() = withValidityAssertion { _kind }
+}
+
+public enum class KtImplicitReceiverSmartCastKind {
+    DISPATCH, EXTENSION
+}