FIR: Allow selector of safe calls to be a FirStatement

It's necessary to allow a?.b += v be interpreted as a?.(b += v)
But currently FirAssignmentOperatorStatement is not FirQualifiedAccess

^KT-41034 In Progress

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

@@ -155,7 +155,7 @@ internal class KtFirCallResolver(
             )
             // FIR does not resolve to a symbol for equality calls.
             is FirEqualityOperatorCall -> toKtCallInfo(psi)
-            is FirSafeCallExpression -> regularQualifiedAccess.toKtCallInfo(
+            is FirSafeCallExpression -> selector.toKtCallInfo(
                 psi,
                 resolveCalleeExpressionOfFunctionCall,
                 resolveFragmentOfCall

compiler/fir/dump/src/org/jetbrains/kotlin/fir/dump/HtmlFirDump.kt

@@ -1595,21 +1595,21 @@ class HtmlFirDump internal constructor(private var linkResolver: FirLinkResolver
 
         +"?."
 
-        val nestedQualifier = safeCallExpression.regularQualifiedAccess
+        val selector = safeCallExpression.selector
-        if (nestedQualifier.explicitReceiver == safeCallExpression.checkedSubjectRef.value) {
+        if (selector is FirQualifiedAccess && selector.explicitReceiver == safeCallExpression.checkedSubjectRef.value) {
-            when (nestedQualifier) {
+            when (selector) {
                 is FirFunctionCall -> {
-                    return generate(nestedQualifier, skipReceiver = true)
+                    return generate(selector, skipReceiver = true)
                 }
                 is FirQualifiedAccessExpression -> {
-                    return generate(nestedQualifier, skipReceiver = true)
+                    return generate(selector, skipReceiver = true)
                 }
             }
         }
 
         +"{ "
 
-        generate(nestedQualifier)
+        generate(selector)
 
         +" }"
     }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -391,7 +391,7 @@ class Fir2IrVisitor(
         )
 
         return conversionScope.withSafeCallSubject(receiverVariable) {
-            val afterNotNullCheck = safeCallExpression.regularQualifiedAccess.accept(this, data) as IrExpression
+            val afterNotNullCheck = safeCallExpression.selector.accept(this, data) as IrExpression
             components.createSafeCallConstruction(receiverVariable, variableSymbol, afterNotNullCheck)
         }
     }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -966,7 +966,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         value: FirExpression,
         annotations: List<FirAnnotation>
     ): FirSafeCallExpression {
-        val nestedAccess = safeCallNonAssignment.regularQualifiedAccess
+        val nestedAccess = safeCallNonAssignment.selector as FirQualifiedAccess
 
         val assignment = buildVariableAssignment {
             source = baseSource
@@ -976,7 +976,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             this.annotations += annotations
         }
 
-        safeCallNonAssignment.replaceRegularQualifiedAccess(
+        safeCallNonAssignment.replaceSelector(
             assignment
         )
 

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -43,7 +43,6 @@ import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
-import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -583,7 +582,7 @@ fun FirQualifiedAccess.wrapWithSafeCall(receiver: FirExpression, source: KtSourc
         this.checkedSubjectRef = FirExpressionRef<FirCheckedSafeCallSubject>().apply {
             bind(checkedSafeCallSubject)
         }
-        this.regularQualifiedAccess = this@wrapWithSafeCall
+        this.selector = this@wrapWithSafeCall
         this.source = source
     }
 }

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

@@ -30,11 +30,9 @@ import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirPropertyWithExplicitBackingFieldResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
-import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.dfa.PropertyStability
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
-import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
 import org.jetbrains.kotlin.fir.scopes.impl.importedFromObjectData
@@ -50,7 +48,6 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
 import org.jetbrains.kotlin.types.SmartcastStability
-import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 fun List<FirQualifierPart>.toTypeProjections(): Array<ConeTypeProjection> =
@@ -395,7 +392,7 @@ fun FirSafeCallExpression.propagateTypeFromQualifiedAccessAfterNullCheck(
     session: FirSession,
 ) {
     val receiverType = nullableReceiverExpression.typeRef.coneTypeSafe<ConeKotlinType>()
-    val typeAfterNullCheck = regularQualifiedAccess.expressionTypeOrUnitForAssignment() ?: return
+    val typeAfterNullCheck = (selector as? FirQualifiedAccess)?.expressionTypeOrUnitForAssignment() ?: return
     val isReceiverActuallyNullable = if (session.languageVersionSettings.supportsFeature(LanguageFeature.SafeCallsAreAlwaysNullable)) {
         true
     } else {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -15,7 +15,8 @@ import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
-import org.jetbrains.kotlin.fir.resolve.inference.*
+import org.jetbrains.kotlin.fir.resolve.inference.preprocessCallableReference
+import org.jetbrains.kotlin.fir.resolve.inference.preprocessLambdaArgument
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.transformers.ensureResolvedTypeDeclaration
 import org.jetbrains.kotlin.fir.returnExpressions
@@ -63,8 +64,8 @@ fun Candidate.resolveArgumentExpression(
         // and then add constraint: typeOf(`$not-null-receiver$.bar()`).makeNullable() <: EXPECTED_TYPE
         // NB: argument.regularQualifiedAccess is either a call or a qualified access
         is FirSafeCallExpression -> {
-            val nestedQualifier = argument.regularQualifiedAccess
+            val nestedQualifier = argument.selector
-            if (nestedQualifier is FirExpression) {
+            if (nestedQualifier is FirQualifiedAccessExpression) {
                 resolveSubCallArgument(
                     csBuilder,
                     nestedQualifier,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CandidateFactory.kt

@@ -151,7 +151,7 @@ fun PostponedArgumentsAnalyzerContext.addSubsystemFromExpression(statement: FirS
         is FirElvisExpression
         -> (statement as FirResolvable).candidate()?.let { addOtherSystem(it.system.asReadOnlyStorage()) }
 
-        is FirSafeCallExpression -> addSubsystemFromExpression(statement.regularQualifiedAccess)
+        is FirSafeCallExpression -> addSubsystemFromExpression(statement.selector)
         is FirWrappedArgumentExpression -> addSubsystemFromExpression(statement.expression)
         is FirBlock -> statement.returnExpressions().forEach { addSubsystemFromExpression(it) }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/ConstraintSystemCompleter.kt

@@ -464,7 +464,7 @@ fun FirStatement.processAllContainingCallCandidates(processBlocks: Boolean, proc
         }
 
         is FirSafeCallExpression -> {
-            this.regularQualifiedAccess.processAllContainingCallCandidates(processBlocks, processor)
+            this.selector.processAllContainingCallCandidates(processBlocks, processor)
         }
 
         is FirWhenExpression -> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -354,7 +354,7 @@ class FirCallCompletionResultsWriterTransformer(
         safeCallExpression: FirSafeCallExpression,
         data: ExpectedArgumentType?
     ): FirStatement {
-        safeCallExpression.transformRegularQualifiedAccess(
+        safeCallExpression.transformSelector(
             this,
             data?.getExpectedType(
                 safeCallExpression

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

@@ -119,7 +119,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                 transformSuperReceiver(
                     callee,
                     qualifiedAccessExpression,
-                    containingSafeCallExpression?.takeIf { qualifiedAccessExpression == it.receiver }?.regularQualifiedAccess
+                    containingSafeCallExpression?.takeIf { qualifiedAccessExpression == it.receiver }?.selector as? FirQualifiedAccess
                 )
             }
             is FirDelegateFieldReference -> {
@@ -312,7 +312,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
 
             safeCallExpression.apply {
                 checkedSubjectRef.value.propagateTypeFromOriginalReceiver(receiver, components.session)
-                transformRegularQualifiedAccess(this@FirExpressionsResolveTransformer, data)
+                transformSelector(this@FirExpressionsResolveTransformer, data)
                 propagateTypeFromQualifiedAccessAfterNullCheck(receiver, session)
             }
 
@@ -685,7 +685,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             // AS operator doesn't add expected type to any other expressions
             // See https://kotlinlang.org/docs/whatsnew12.html#support-for-foo-as-a-shorthand-for-this-foo
             // And limitations at org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleterKt.isFunctionForExpectTypeFromCastFeature(org.jetbrains.kotlin.fir.declarations.FirFunction<?>)
-            if (argument is FirFunctionCall || (argument is FirSafeCallExpression && argument.regularQualifiedAccess is FirFunctionCall)) {
+            if (argument is FirFunctionCall || (argument is FirSafeCallExpression && argument.selector is FirFunctionCall)) {
                 val expectedType = conversionTypeRef.coneTypeSafe<ConeKotlinType>()?.takeIf {
                     // is not bare type
                     it !is ConeClassLikeType ||

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorageImpl.kt

@@ -62,7 +62,7 @@ class VariableStorageImpl(private val session: FirSession) : VariableStorage() {
     private fun FirElement.unwrapElement(): FirElement = when (this) {
         is FirWhenSubjectExpression -> whenRef.value.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this
         is FirExpressionWithSmartcast -> originalExpression.unwrapElement()
-        is FirSafeCallExpression -> regularQualifiedAccess.unwrapElement()
+        is FirSafeCallExpression -> selector.unwrapElement()
         is FirCheckedSafeCallSubject -> originalReceiverRef.value.unwrapElement()
         is FirCheckNotNullCall -> argument.unwrapElement()
         else -> this

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt

@@ -98,7 +98,7 @@ val FirElement.symbol: FirBasedSymbol<*>?
         is FirResolvable -> symbol
         is FirDeclaration -> symbol
         is FirWhenSubjectExpression -> whenRef.value.subject?.symbol
-        is FirSafeCallExpression -> regularQualifiedAccess.symbol
+        is FirSafeCallExpression -> selector.symbol
         else -> null
     }?.takeIf {
         (this as? FirExpression)?.unwrapSmartcastExpression() is FirThisReceiverExpression ||

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirSafeCallExpression.kt

@@ -22,7 +22,7 @@ abstract class FirSafeCallExpression : FirExpression() {
     abstract override val annotations: List<FirAnnotation>
     abstract val receiver: FirExpression
     abstract val checkedSubjectRef: FirExpressionRef<FirCheckedSafeCallSubject>
-    abstract val regularQualifiedAccess: FirQualifiedAccess
+    abstract val selector: FirStatement
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitSafeCallExpression(this, data)
 
@@ -32,11 +32,11 @@ abstract class FirSafeCallExpression : FirExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
-    abstract fun replaceRegularQualifiedAccess(newRegularQualifiedAccess: FirQualifiedAccess)
+    abstract fun replaceSelector(newSelector: FirStatement)
 
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirSafeCallExpression
 
     abstract fun <D> transformReceiver(transformer: FirTransformer<D>, data: D): FirSafeCallExpression
 
-    abstract fun <D> transformRegularQualifiedAccess(transformer: FirTransformer<D>, data: D): FirSafeCallExpression
+    abstract fun <D> transformSelector(transformer: FirTransformer<D>, data: D): FirSafeCallExpression
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirSafeCallExpressionBuilder.kt

@@ -13,8 +13,8 @@ import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirSafeCallExpressionImpl
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -33,7 +33,7 @@ class FirSafeCallExpressionBuilder : FirAnnotationContainerBuilder, FirExpressio
     override val annotations: MutableList<FirAnnotation> = mutableListOf()
     lateinit var receiver: FirExpression
     lateinit var checkedSubjectRef: FirExpressionRef<FirCheckedSafeCallSubject>
-    lateinit var regularQualifiedAccess: FirQualifiedAccess
+    lateinit var selector: FirStatement
 
     override fun build(): FirSafeCallExpression {
         return FirSafeCallExpressionImpl(
@@ -42,7 +42,7 @@ class FirSafeCallExpressionBuilder : FirAnnotationContainerBuilder, FirExpressio
             annotations,
             receiver,
             checkedSubjectRef,
-            regularQualifiedAccess,
+            selector,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirSafeCallExpressionImpl.kt

@@ -10,8 +10,8 @@ import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 
@@ -26,20 +26,20 @@ internal class FirSafeCallExpressionImpl(
     override val annotations: MutableList<FirAnnotation>,
     override var receiver: FirExpression,
     override val checkedSubjectRef: FirExpressionRef<FirCheckedSafeCallSubject>,
-    override var regularQualifiedAccess: FirQualifiedAccess,
+    override var selector: FirStatement,
 ) : FirSafeCallExpression() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         typeRef.accept(visitor, data)
         annotations.forEach { it.accept(visitor, data) }
         receiver.accept(visitor, data)
-        regularQualifiedAccess.accept(visitor, data)
+        selector.accept(visitor, data)
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirSafeCallExpressionImpl {
         typeRef = typeRef.transform(transformer, data)
         transformAnnotations(transformer, data)
         transformReceiver(transformer, data)
-        transformRegularQualifiedAccess(transformer, data)
+        transformSelector(transformer, data)
         return this
     }
 
@@ -53,8 +53,8 @@ internal class FirSafeCallExpressionImpl(
         return this
     }
 
-    override fun <D> transformRegularQualifiedAccess(transformer: FirTransformer<D>, data: D): FirSafeCallExpressionImpl {
+    override fun <D> transformSelector(transformer: FirTransformer<D>, data: D): FirSafeCallExpressionImpl {
-        regularQualifiedAccess = regularQualifiedAccess.transform(transformer, data)
+        selector = selector.transform(transformer, data)
         return this
     }
 
@@ -62,7 +62,7 @@ internal class FirSafeCallExpressionImpl(
         typeRef = newTypeRef
     }
 
-    override fun replaceRegularQualifiedAccess(newRegularQualifiedAccess: FirQualifiedAccess) {
+    override fun replaceSelector(newSelector: FirStatement) {
-        regularQualifiedAccess = newRegularQualifiedAccess
+        selector = newSelector
     }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -707,7 +707,7 @@ open class FirRenderer(builder: StringBuilder, protected val mode: RenderMode =
     override fun visitSafeCallExpression(safeCallExpression: FirSafeCallExpression) {
         safeCallExpression.receiver.accept(this)
         print("?.{ ")
-        safeCallExpression.regularQualifiedAccess.accept(this)
+        safeCallExpression.selector.accept(this)
         print(" }")
     }
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -506,7 +506,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             // Special node that might be used as a reference to receiver of a safe call after null check
             +field("checkedSubjectRef", safeCallCheckedSubjectReferenceType)
             // One that uses checkedReceiver as a receiver
-            +field("regularQualifiedAccess", qualifiedAccess, withReplace = true).withTransform()
+            +field("selector", statement, withReplace = true).withTransform()
         }
 
         checkedSafeCallSubject.configure {

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt

@@ -12,14 +12,11 @@ import org.jetbrains.kotlin.checkers.utils.TypeOfCall
 import org.jetbrains.kotlin.diagnostics.*
 import org.jetbrains.kotlin.diagnostics.rendering.Renderers
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.analysis.diagnostics.*
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
-import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
@@ -176,13 +173,16 @@ class FirDiagnosticsHandler(testServices: TestServices) : FirAnalysisHandler(tes
             }
 
             override fun visitSafeCallExpression(safeCallExpression: FirSafeCallExpression) {
-                result.addIfNotNull(
+                val selector = safeCallExpression.selector
-                    createCallDiagnosticIfExpected(
+                if (selector is FirQualifiedAccess) {
-                        safeCallExpression,
+                    result.addIfNotNull(
-                        safeCallExpression.regularQualifiedAccess.calleeReference as FirNamedReference,
+                        createCallDiagnosticIfExpected(
-                        diagnosedRangesToDiagnosticNames
+                            safeCallExpression,
+                            selector.calleeReference as FirNamedReference,
+                            diagnosedRangesToDiagnosticNames
+                        )
                     )
-                )
+                }
 
                 super.visitSafeCallExpression(safeCallExpression)
             }

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -742,7 +742,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         override fun visitSafeCallExpression(safeCallExpression: FirSafeCallExpression, data: StringBuilder) {
             safeCallExpression.receiver.accept(this, data)
             data.append("?.{ ")
-            safeCallExpression.regularQualifiedAccess.accept(this, data)
+            safeCallExpression.selector.accept(this, data)
             data.append(" }")
         }