FIR: Extract putIncrementOrDecrementStatements

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

@@ -516,38 +516,16 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             val desugaredSource = baseSource?.fakeElement(KtFakeSourceElementKind.DesugaredIncrementOrDecrement)
             source = desugaredSource
 
-            // initialValueVar is only used for postfix increment/decrement (stores the argument value before increment/decrement).
+            val convertedReceiver = unwrappedReceiver.convert()
-            val initialValueVar = generateTemporaryVariable(
+
-                baseModuleData,
+            putIncrementOrDecrementStatements(
+                convertedReceiver,
+                operationReference,
+                callName,
+                prefix,
+                unwrappedReceiver.takeIf { it.elementType == REFERENCE_EXPRESSION }?.getReferencedNameAsName(),
                 desugaredSource,
-                SpecialNames.UNARY,
+            ) { resultInitializer: FirExpression, resultVar: FirVariable ->
-                unwrappedReceiver.convert()
-            )
-
-            // resultInitializer is the expression for `argument.inc()`
-            val resultInitializer = buildFunctionCall {
-                source = desugaredSource
-                calleeReference = buildSimpleNamedReference {
-                    source = operationReference?.toFirSourceElement()
-                    name = callName
-                }
-                explicitReceiver = if (prefix) {
-                    unwrappedReceiver.convert()
-                } else {
-                    generateResolvedAccessExpression(desugaredSource, initialValueVar)
-                }
-                origin = FirFunctionCallOrigin.Operator
-            }
-
-            // resultVar is only used for prefix increment/decrement.
-            val resultVar = generateTemporaryVariable(
-                baseModuleData,
-                desugaredSource,
-                Name.special("<unary-result>"),
-                resultInitializer
-            )
-
-            fun appendAssignment() {
                 val assignment = unwrappedReceiver.generateAssignment(
                     desugaredSource,
                     null,
@@ -567,21 +545,87 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     statements += assignment
                 }
             }
+        }
+    }
 
-            if (prefix) {
+
-                if (unwrappedReceiver.elementType != REFERENCE_EXPRESSION) {
+    /**
-                    statements += resultVar
+     * given:
-                    appendAssignment()
+     * receiver++
-                    statements += generateResolvedAccessExpression(desugaredSource, resultVar)
+     *
-                } else {
+     * result:
-                    appendAssignment()
+     * {
-                    statements += generateAccessExpression(desugaredSource, desugaredSource, unwrappedReceiver.getReferencedNameAsName())
+     *     val <unary> = receiver
-                }
+     *     val resultVar = <unary>.inc()
-            } else {
+     *     appendAssignment(resultVar)
-                statements += initialValueVar
+     *     ^<unary>
-                appendAssignment()
+     * }
-                statements += generateResolvedAccessExpression(desugaredSource, initialValueVar)
+     *
+     * given:
+     * ++receiver
+     *
+     * result:
+     * {
+     *     val <unary-result> = receiver.inc()
+     *     val resultVar = <unary-result>
+     *     appendAssignment(resultVar)
+     *     ^<unary-result>
+     * }
+     *
+     */
+    private fun FirBlockBuilder.putIncrementOrDecrementStatements(
+        receiver: FirExpression,
+        operationReference: T?,
+        callName: Name, // 'inc' or 'dec'
+        prefix: Boolean,
+        nameIfSimpleReference: Name?, // 'b' if whole expression is simple `b++` or `a.b++`, but not `a[1]++`
+        desugaredSource: KtSourceElement?,
+        appendAssignment: FirBlockBuilder.(resultInitializer: FirExpression, resultVar: FirVariable) -> Unit
+    ) {
+        // initialValueVar is only used for postfix increment/decrement (stores the argument value before increment/decrement).
+        val initialValueVar = generateTemporaryVariable(
+            baseModuleData,
+            desugaredSource,
+            SpecialNames.UNARY,
+            receiver
+        )
+
+        // resultInitializer is the expression for `argument.inc()`
+        val resultInitializer = buildFunctionCall {
+            source = desugaredSource
+            calleeReference = buildSimpleNamedReference {
+                source = operationReference?.toFirSourceElement()
+                name = callName
             }
+            explicitReceiver = if (prefix) {
+                receiver
+            } else {
+                generateResolvedAccessExpression(desugaredSource, initialValueVar)
+            }
+            origin = FirFunctionCallOrigin.Operator
+        }
+
+        // resultVar is only used for prefix increment/decrement.
+        val resultVar = generateTemporaryVariable(
+            baseModuleData,
+            desugaredSource,
+            Name.special("<unary-result>"),
+            resultInitializer
+        )
+
+        if (prefix) {
+            if (nameIfSimpleReference != null) {
+                appendAssignment(resultInitializer, resultVar)
+                statements += generateAccessExpression(desugaredSource, desugaredSource, nameIfSimpleReference)
+            } else {
+                statements += resultVar
+                appendAssignment(resultInitializer, resultVar)
+                statements += generateResolvedAccessExpression(desugaredSource, resultVar)
+            }
+        } else {
+            statements += initialValueVar
+            appendAssignment(resultInitializer, resultVar)
+            statements += generateResolvedAccessExpression(desugaredSource, initialValueVar)
         }
     }
 
@@ -657,64 +701,25 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 firArgumentSelector.also { if (it is FirQualifiedAccessExpression) it.replaceExplicitReceiver(receiver) }
             }
 
-            // initialValueVar is only used for postfix increment/decrement (stores the argument value before increment/decrement).
+            putIncrementOrDecrementStatements(
-            val initialValueVar = generateTemporaryVariable(
+                firArgument, operationReference, callName, prefix,
-                baseModuleData,
+                nameIfSimpleReference = null, desugaredSource
-                desugaredSource,
+            ) { resultInitializer: FirExpression, resultVar: FirVariable ->
-                SpecialNames.UNARY,
+                if (firArgument !is FirQualifiedAccessExpression) return@putIncrementOrDecrementStatements
-                firArgument
+                statements += buildVariableAssignment {
-            )
+                    source = desugaredSource
-
+                    rValue = if (prefix) {
-            // resultInitializer is the expression for `argument.inc()`
+                        generateResolvedAccessExpression(source, resultVar)
-            val resultInitializer = buildFunctionCall {
+                    } else {
-                source = desugaredSource
+                        resultInitializer
-                calleeReference = buildSimpleNamedReference {
+                    }
-                    source = operationReference?.toFirSourceElement()
+                    explicitReceiver = generateResolvedAccessExpression(argumentReceiverVariable.source, argumentReceiverVariable)
-                    name = callName
+                    calleeReference = buildSimpleNamedReference {
-                }
+                        source = firArgument.calleeReference.source
-                explicitReceiver = if (prefix) {
+                        name = (firArgument.calleeReference as FirSimpleNamedReference).name
-                    firArgument
-                } else {
-                    generateResolvedAccessExpression(desugaredSource, initialValueVar)
-                }
-                origin = FirFunctionCallOrigin.Operator
-            }
-
-            // resultVar is only used for prefix increment/decrement.
-            val resultVar = generateTemporaryVariable(
-                baseModuleData,
-                desugaredSource,
-                Name.special("<unary-result>"),
-                resultInitializer
-            )
-
-            fun appendAssignment() {
-                if (firArgument is FirQualifiedAccessExpression) {
-                    statements += buildVariableAssignment {
-                        source = desugaredSource
-                        rValue = if (prefix) {
-                            generateResolvedAccessExpression(source, resultVar)
-                        } else {
-                            resultInitializer
-                        }
-                        explicitReceiver = generateResolvedAccessExpression(argumentReceiverVariable.source, argumentReceiverVariable)
-                        calleeReference = buildSimpleNamedReference {
-                            source = firArgument.calleeReference.source
-                            name = (firArgument.calleeReference as FirSimpleNamedReference).name
-                        }
                     }
                 }
             }
-
-            if (prefix) {
-                statements += resultVar
-                appendAssignment()
-                statements += generateResolvedAccessExpression(desugaredSource, resultVar)
-            } else {
-                statements += initialValueVar
-                appendAssignment()
-                statements += generateResolvedAccessExpression(desugaredSource, initialValueVar)
-            }
         }
     }
 
@@ -797,38 +802,10 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 origin = FirFunctionCallOrigin.Operator
             }
 
-            // initialValueVar is only used for postfix increment/decrement (stores the argument value before increment/decrement).
+            putIncrementOrDecrementStatements(
-            val initialValueVar = generateTemporaryVariable(
+                firArgument, operationReference, callName, prefix,
-                baseModuleData,
+                nameIfSimpleReference = null, desugaredSource
-                desugaredSource,
+            ) { resultInitializer: FirExpression, resultVar: FirVariable ->
-                SpecialNames.UNARY,
-                firArgument
-            )
-
-            // resultInitializer is the expression for `argument.inc()`
-            val resultInitializer = buildFunctionCall {
-                source = desugaredSource
-                calleeReference = buildSimpleNamedReference {
-                    source = operationReference?.toFirSourceElement()
-                    name = callName
-                }
-                explicitReceiver = if (prefix) {
-                    firArgument
-                } else {
-                    generateResolvedAccessExpression(desugaredSource, initialValueVar)
-                }
-                origin = FirFunctionCallOrigin.Operator
-            }
-
-            // resultVar is only used for prefix increment/decrement.
-            val resultVar = generateTemporaryVariable(
-                baseModuleData,
-                desugaredSource,
-                Name.special("<unary-result>"),
-                resultInitializer
-            )
-
-            fun appendAssignment() {
                 statements += buildFunctionCall {
                     source = desugaredSource
                     calleeReference = buildSimpleNamedReference {
@@ -849,16 +826,6 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     origin = FirFunctionCallOrigin.Operator
                 }
             }
-
-            if (prefix) {
-                statements += resultVar
-                appendAssignment()
-                statements += generateResolvedAccessExpression(desugaredSource, resultVar)
-            } else {
-                statements += initialValueVar
-                appendAssignment()
-                statements += generateResolvedAccessExpression(desugaredSource, initialValueVar)
-            }
         }
     }