FIR: Support safe-calls new format in FIR builders

^KT-38444 In Progress

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

@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.builder.*
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.*
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
@@ -536,6 +536,14 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             }
         }
         require(operation == FirOperation.ASSIGN)
+
+        if (this?.elementType == SAFE_ACCESS_EXPRESSION && this != null) {
+            val safeCallNonAssignment = convert() as? FirSafeCallExpression
+            if (safeCallNonAssignment != null) {
+                return putAssignmentToSafeCall(safeCallNonAssignment, baseSource, value)
+            }
+        }
+
         return buildVariableAssignment {
             source = baseSource
             safe = false
@@ -544,6 +552,29 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
     }
 
+    // gets a?.{ $subj.x } and turns it to a?.{ $subj.x = v }
+    private fun putAssignmentToSafeCall(
+        safeCallNonAssignment: FirSafeCallExpression,
+        baseSource: FirSourceElement?,
+        value: FirExpression
+    ): FirSafeCallExpression {
+        val nestedAccess = safeCallNonAssignment.regularQualifiedAccess
+
+        val assignment = buildVariableAssignment {
+            source = baseSource
+            safe = false
+            rValue = value
+            calleeReference = nestedAccess.calleeReference
+            explicitReceiver = safeCallNonAssignment.checkedSubject.value
+        }
+
+        safeCallNonAssignment.replaceRegularQualifiedAccess(
+            assignment
+        )
+
+        return safeCallNonAssignment
+    }
+
     private fun T.generateAugmentedArraySetCall(
         baseSource: FirSourceElement?,
         operation: FirOperation,

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

@@ -19,13 +19,20 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.expressions.impl.FirModifiableQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirStubStatement
 import org.jetbrains.kotlin.fir.expressions.impl.buildSingleExpressionBlock
 import org.jetbrains.kotlin.fir.references.FirNamedReference
-import org.jetbrains.kotlin.fir.references.builder.*
+import org.jetbrains.kotlin.fir.references.builder.buildDelegateFieldReference
+import org.jetbrains.kotlin.fir.references.builder.buildImplicitThisReference
+import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
 import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType
-import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.symbols.impl.FirDelegateFieldSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertyAccessorSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.ConeStarProjection
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
@@ -479,3 +486,22 @@ private fun FirExpression.checkReceiver(name: String?): Boolean {
     val receiverName = (receiver.calleeReference as? FirNamedReference)?.name?.asString() ?: return false
     return receiverName == name
 }
+
+fun FirModifiableQualifiedAccess.wrapWithSafeCall(receiver: FirExpression): FirSafeCallExpression {
+    // TODO: Refactor tree to make FirModifiableQualifiedAccess inherit FirQualifiedAccess
+    require(this is FirQualifiedAccess) {
+        "Safe-call instances are expected to be FirQualifiedAccess, but ${this::class} was found"
+    }
+
+    val checkedSafeCallSubject = buildCheckedSafeCallSubject {
+        this.originalReceiverReference = FirSafeCallOriginalReceiverReference(receiver)
+    }
+
+    explicitReceiver = checkedSafeCallSubject
+    return buildSafeCallExpression {
+        this.receiver = receiver
+        this.checkedSubject = FirSafeCallCheckedSubjectReference(checkedSafeCallSubject)
+        this.regularQualifiedAccess = this@wrapWithSafeCall
+        this.source = this@wrapWithSafeCall.source
+    }
+}

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -483,7 +483,11 @@ class ExpressionsConverter(
         }
 
         (firSelector as? FirModifiableQualifiedAccess)?.let {
-            it.safe = isSafe
+            if (isSafe) {
+                return it.wrapWithSafeCall(firReceiver!!)
+            }
+
+            it.safe = false
             it.explicitReceiver = firReceiver
         }
         return firSelector

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1699,8 +1699,14 @@ class RawFirBuilder(
                 )
             val firSelector = selector.toFirExpression("Incorrect selector expression")
             if (firSelector is FirModifiableQualifiedAccess) {
-                firSelector.safe = expression is KtSafeQualifiedExpression
+                val receiver = expression.receiverExpression.toFirExpression("Incorrect receiver expression")
-                firSelector.explicitReceiver = expression.receiverExpression.toFirExpression("Incorrect receiver expression")
+
+                if (expression is KtSafeQualifiedExpression) {
+                    return firSelector.wrapWithSafeCall(receiver)
+                }
+
+                firSelector.safe = false
+                firSelector.explicitReceiver = receiver
             }
             return firSelector
         }

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/calls.txt

@@ -34,9 +34,9 @@ FILE: calls.kt
     }
     public? final? fun testQualified(first: My, second: My?): R|kotlin/Unit| {
         println#(first#.x#)
-        println#(second#?.x#)
+        println#(second#?.{ $subj$.x# })
         first#.foo#()
-        second#?.foo#()
+        second#?.{ $subj$.foo#() }
         first#.copy#().foo#()
         first#.x# = IntegerLiteral(42)
     }