[IR] Support if and when for MFVC, fix several bugs, refactor

#KT-1179
2022-10-03 08:49:52 +02:00
parent e76041bb3b
commit ca789607ba
11 changed files with 347 additions and 124 deletions
@@ -50195,6 +50195,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
            runTest("compiler/testData/codegen/box/valueClasses/complex.kt", TransformersFunctions.getReplaceOptionalJvmInlineAnnotationWithReal());
        }

+        @Test
+        @TestMetadata("conditionalExpressions.kt")
+        public void testConditionalExpressions() throws Exception {
+            runTest("compiler/testData/codegen/box/valueClasses/conditionalExpressions.kt", TransformersFunctions.getReplaceOptionalJvmInlineAnnotationWithReal());
+        }
+
        @Test
        @TestMetadata("equality.kt")
        public void testEquality() throws Exception {
@@ -5789,6 +5789,12 @@ public class FirBytecodeTextTestGenerated extends AbstractFirBytecodeTextTest {
            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/bytecodeText/valueClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
        }

+        @Test
+        @TestMetadata("conditionalExpressions.kt")
+        public void testConditionalExpressions() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/valueClasses/conditionalExpressions.kt", TransformersFunctions.getReplaceOptionalJvmInlineAnnotationWithReal());
+        }
+
        @Test
        @TestMetadata("equalsBoxTest.kt")
        public void testEqualsBoxTest() throws Exception {
@@ -23,6 +23,8 @@ import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
@@ -93,14 +95,14 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
            return expressions.subList(0, expressions.size - repeatable.size) to repeatable
        }

-        fun makeReplacement(scope: IrBuilderWithScope, expression: IrSetValue, safe: Boolean): IrExpression? = with(scope) {
+        fun addReplacement(scope: IrBlockBuilder, expression: IrSetValue, safe: Boolean): IrExpression? = with(scope) {
            oldValueSymbol2NewValueSymbol[expression.symbol]?.let { return irSet(it.owner, expression.value) }
            val instance = oldSymbol2MfvcNodeInstance[expression.symbol] ?: return@with null
-            var values: List<IrExpression>? = null
-            return irBlock {
-                values = makeFlattenedExpressionsWithGivenSafety(instance.node, safe, expression.value)
-                instance.addSetterStatements(this, values!!)
-            }.also { expression2MfvcNodeInstanceAccessor[it] = MfvcNodeInstanceAccessor.Setter(instance, values!!) }
+            val values: List<IrExpression> = makeFlattenedExpressionsWithGivenSafety(instance.node, safe, expression.value)
+            val setterExpressions = instance.makeSetterExpressions(values)
+            expression2MfvcNodeInstanceAccessor[setterExpressions] = MfvcNodeInstanceAccessor.Setter(instance, values)
+            +setterExpressions
+            return setterExpressions
        }

        /**
@@ -124,50 +126,52 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
        private val IrFieldAccessExpression.field: IrField
            get() = this.symbol.owner

-        fun makeReplacement(scope: IrBuilderWithScope, expression: IrGetField): IrExpression? {
+        fun addReplacement(scope: IrBlockBuilder, expression: IrGetField): IrExpression? {
            val property = expression.field.property ?: return null
-            expression.receiver?.get(property.name)?.let { return it }
+            expression.receiver?.get(property.name)?.let { scope.run { +it }; return it }
            val node = replacements.getMfvcPropertyNode(property) ?: return null
            val typeArguments = makeTypeArgumentsFromField(expression)
-            var instance: ReceiverBasedMfvcNodeInstance? = null
-            return scope.irBlock {
-                instance = node.createInstanceFromBox(this, typeArguments, expression.receiver, AccessType.AlwaysPrivate, ::variablesSaver)
-                +instance!!.makeGetterExpression()
-            }.also { expression2MfvcNodeInstanceAccessor[it] = MfvcNodeInstanceAccessor.Getter(instance!!) }
+            val instance: ReceiverBasedMfvcNodeInstance =
+                node.createInstanceFromBox(scope, typeArguments, expression.receiver, AccessType.AlwaysPrivate, ::variablesSaver)
+            val getterExpression = instance.makeGetterExpression()
+            expression2MfvcNodeInstanceAccessor[getterExpression] = MfvcNodeInstanceAccessor.Getter(instance)
+            scope.run { +getterExpression }
+            return getterExpression
        }

-        fun makeReplacement(scope: IrBuilderWithScope, expression: IrSetField, safe: Boolean): IrExpression? {
+        fun addReplacement(scope: IrBlockBuilder, expression: IrSetField, safe: Boolean): IrExpression? {
            val property = expression.field.property ?: return null
-            expression.receiver?.get(property.name)?.let { return it }
+            expression.receiver?.get(property.name)?.let { scope.run { +it }; return it }
            val node = replacements.getMfvcPropertyNode(property) ?: return null
            val typeArguments = makeTypeArgumentsFromField(expression)
-            var instance: ReceiverBasedMfvcNodeInstance? = null
-            var values: List<IrExpression>? = null
-            return scope.irBlock {
-                instance = node.createInstanceFromBox(this, typeArguments, expression.receiver, AccessType.AlwaysPrivate, ::variablesSaver)
-                values = makeFlattenedExpressionsWithGivenSafety(node, safe, expression.value)
-                instance!!.addSetterStatements(this, values!!)
-            }.also { expression2MfvcNodeInstanceAccessor[it] = MfvcNodeInstanceAccessor.Setter(instance!!, values!!) }
+            val instance: ReceiverBasedMfvcNodeInstance =
+                node.createInstanceFromBox(scope, typeArguments, expression.receiver, AccessType.AlwaysPrivate, ::variablesSaver)
+            val values: List<IrExpression> = scope.makeFlattenedExpressionsWithGivenSafety(node, safe, expression.value)
+            val setterExpressions = instance.makeSetterExpressions(values)
+            expression2MfvcNodeInstanceAccessor[setterExpressions] = MfvcNodeInstanceAccessor.Setter(instance, values)
+            scope.run { +setterExpressions }
+            return setterExpressions
        }

-        fun makeReplacement(scope: IrBuilderWithScope, expression: IrCall): IrExpression? {
+        fun addReplacement(scope: IrBlockBuilder, expression: IrCall): IrExpression? {
            val function = expression.symbol.owner
            val property = function.property?.takeIf { function.isGetter } ?: return null
            val dispatchReceiver = expression.dispatchReceiver
-            dispatchReceiver?.get(property.name)?.let { return it }
+            dispatchReceiver?.get(property.name)?.let { scope.run { +it }; return it }
            val node = replacements.getMfvcPropertyNode(property) ?: return null
            val typeArguments = makeTypeArgumentsFromFunction(expression)
-            var instance: ReceiverBasedMfvcNodeInstance? = null
-            return scope.irBlock {
-                // Optimization: pure function access to leaf can be replaced with field access if the field itself is accessible
-                val accessType = when {
-                    !node.hasPureUnboxMethod -> AccessType.AlwaysPublic
-                    dispatchReceiver == null -> AccessType.PrivateWhenNoBox
-                    else -> getOptimizedPublicAccess(dispatchReceiver.type.erasedUpperBound)
-                }
-                instance = node.createInstanceFromBox(this, typeArguments, dispatchReceiver, accessType, ::variablesSaver)
-                +instance!!.makeGetterExpression()
-            }.also { expression2MfvcNodeInstanceAccessor[it] = MfvcNodeInstanceAccessor.Getter(instance!!) }
+            // Optimization: pure function access to leaf can be replaced with field access if the field itself is accessible
+            val accessType = when {
+                !node.hasPureUnboxMethod -> AccessType.AlwaysPublic
+                dispatchReceiver == null -> AccessType.PrivateWhenNoBox
+                else -> getOptimizedPublicAccess(dispatchReceiver.type.erasedUpperBound)
+            }
+            val instance: ReceiverBasedMfvcNodeInstance =
+                node.createInstanceFromBox(scope, typeArguments, dispatchReceiver, accessType, ::variablesSaver)
+            val getterExpression = instance.makeGetterExpression()
+            expression2MfvcNodeInstanceAccessor[getterExpression] = MfvcNodeInstanceAccessor.Getter(instance)
+            scope.run { +getterExpression }
+            return getterExpression
        }

        private val IrField.property
@@ -187,25 +191,51 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
            expression.dispatchReceiver?.type?.let { putAll(makeTypeArgumentsFromType(it as IrSimpleType)) }
        }

-        operator fun IrExpression.get(name: Name): IrExpression? {
-            val accessor = getMfvcNodeInstanceAccessor(this) ?: return null
-            val newAccessor = accessor[name] ?: return null
+        private fun handleSavedExpression(
+            expression: IrExpression, handler: IrBlockBuilder.(accessor: MfvcNodeInstanceAccessor) -> Unit?
+        ): IrExpression? {
+            val accessor = expression2MfvcNodeInstanceAccessor[expression]
+            return when {
+                accessor != null -> accessor.instance.scope.irBlock { handler(accessor) ?: return null }
+                expression !is IrContainerExpression -> null
+                else -> when (val lastExpression = expression.statements.lastOrNull()) {
+                    is IrExpression -> {
+                        val inner = handleSavedExpression(lastExpression, handler) ?: return null
+                        if (expression.isTransparentScope) IrCompositeImpl(
+                            startOffset = expression.startOffset,
+                            endOffset = expression.endOffset,
+                            type = inner.type,
+                            origin = expression.origin,
+                            statements = expression.statements.dropLast(1) + inner,
+                        ) else IrBlockImpl(
+                            startOffset = expression.startOffset,
+                            endOffset = expression.endOffset,
+                            type = inner.type,
+                            origin = expression.origin,
+                            statements = expression.statements.dropLast(1) + inner,
+                        )
+                    }
+
+                    else -> null
+                }
+            }
+        }
+
+        operator fun IrExpression.get(name: Name): IrExpression? = handleSavedExpression(this) { accessor ->
+            val newAccessor = accessor[name] ?: return@handleSavedExpression null
            val expression = when (newAccessor) {
                is MfvcNodeInstanceAccessor.Getter -> newAccessor.instance.makeGetterExpression()
                is MfvcNodeInstanceAccessor.Setter -> newAccessor.instance.makeSetterExpressions(newAccessor.values)
            }
            expression2MfvcNodeInstanceAccessor[expression] = newAccessor
-            return expression
+            +expression
        }

-        fun getMfvcNodeInstanceAccessor(expression: IrExpression): MfvcNodeInstanceAccessor? =
-            expression2MfvcNodeInstanceAccessor[expression]
-
-        fun getMfvcNodeInstanceIfIsGetter(expression: IrExpression): MfvcNodeInstance? =
-            getMfvcNodeInstanceAccessor(expression)?.takeIf { it is MfvcNodeInstanceAccessor.Getter }?.instance
-
-        fun getMfvcNodeInstanceIfIsSetter(expression: IrExpression): MfvcNodeInstance? =
-            getMfvcNodeInstanceAccessor(expression)?.takeIf { it is MfvcNodeInstanceAccessor.Setter }?.instance
+        fun handleFlattenedGetterExpressions(
+            expression: IrExpression,
+            handler: IrBlockBuilder.(values: List<IrExpression>) -> Unit
+        ): IrExpression? =
+            handleSavedExpression(expression) { handler(it.instance.makeFlattenedGetterExpressions()) }

        fun registerReplacement(expression: IrExpression, instance: MfvcNodeInstance) {
            expression2MfvcNodeInstanceAccessor[expression] = MfvcNodeInstanceAccessor.Getter(instance)
@@ -375,7 +405,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
    private object UNSAFE_MFVC_SET_ORIGIN : IrStatementOrigin

    private fun RootMfvcNode.replaceFields() {
-        mfvc.declarations.removeIf { it is IrField }
+        mfvc.declarations.removeIf { it is IrField && (!it.isStatic || it.type.needsMfvcFlattening()) }
        mfvc.declarations += fields
    }

@@ -579,7 +609,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
            when (val structure = parametersStructure[i]) {
                is MultiFieldValueClassMapping -> {
                    val mfvcNodeInstance = structure.rootMfvcNode.createInstanceFromValueDeclarationsAndBoxType(
-                        scope, structure.boxedType, newParamList, listOf()
+                        scope, structure.boxedType, newParamList
                    )
                    valueDeclarationsRemapper.registerReplacement(param, mfvcNodeInstance)
                }
@@ -598,7 +628,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
        val typeArguments = makeTypeParameterSubstitutionMap(mfvc, primaryConstructorImpl)
        primaryConstructorImpl.body = context.createIrBuilder(primaryConstructorImpl.symbol).irBlockBody {
            val mfvcNodeInstance =
-                ValueDeclarationMfvcNodeInstance(this, rootMfvcNode, typeArguments, primaryConstructorImpl.valueParameters, listOf())
+                ValueDeclarationMfvcNodeInstance(this, rootMfvcNode, typeArguments, primaryConstructorImpl.valueParameters)
            valueDeclarationsRemapper.registerReplacement(
                oldPrimaryConstructor.constructedClass.thisReceiver!!,
                mfvcNodeInstance
@@ -626,15 +656,17 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
        return when {
            function is IrConstructor && function.isPrimary && function.constructedClass.isMultiFieldValueClass &&
                    currentScope.origin != JvmLoweredDeclarationOrigin.SYNTHETIC_MULTI_FIELD_VALUE_CLASS_MEMBER -> {
-                var instance: MfvcNodeInstance? = null
                context.createIrBuilder(currentScope.symbol).irBlock {
+                    val instance: MfvcNodeInstance
                    val rootNode = replacements.getRootMfvcNode(function.constructedClass)!!
                    instance = rootNode.createInstanceFromValueDeclarationsAndBoxType(
                        this, function.constructedClassType as IrSimpleType, Name.identifier("constructor_tmp"), ::variablesSaver
                    )
-                    flattenExpressionTo(expression, instance!!)
-                    +instance!!.makeGetterExpression()
-                }.also { valueDeclarationsRemapper.registerReplacement(it, instance!!) }
+                    flattenExpressionTo(expression, instance)
+                    val getterExpression = instance.makeGetterExpression()
+                    valueDeclarationsRemapper.registerReplacement(getterExpression, instance)
+                    +getterExpression
+                }
            }

            replacement != null -> context.createIrBuilder(currentScope.symbol).irBlock {
@@ -663,7 +695,9 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
        ) {
            require(callee.valueParameters.isEmpty()) { "Unexpected getter:\n${callee.dump()}" }
            expression.dispatchReceiver = expression.dispatchReceiver?.transform(this, null)
-            return valueDeclarationsRemapper.makeReplacement(context.createIrBuilder(getCurrentScopeSymbol()), expression) ?: expression
+            return context.createIrBuilder(getCurrentScopeSymbol()).irBlock {
+                valueDeclarationsRemapper.addReplacement(this, expression) ?: return expression
+            }
        }
        if (expression.isSpecializedMFVCEqEq) {
            return context.createIrBuilder(getCurrentScopeSymbol()).irBlock {
@@ -826,31 +860,49 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
        val forbiddenVariables = mutableSetOf<IrVariable>()
        val variablesSet = variables.toSet()
        val standaloneExpressions = mutableListOf<IrExpression>()
-        for ((expression, variable) in block.statements.asReversed() zip variables.asReversed()) {
-            when {
-                expression !is IrSetValue -> break
-                expression.symbol.owner != variable -> break
-                expression.symbol.owner in forbiddenVariables -> break
-                else -> {
-                    standaloneExpressions.add(expression.value)
-                    expression.value.acceptVoid(object : IrElementVisitorVoid {
-                        override fun visitExpression(expression: IrExpression) {
-                            expression.acceptChildrenVoid(this)
-                        }
+        val resultVariables = variables.toMutableList()

-                        override fun visitValueAccess(expression: IrValueAccessExpression) {
-                            val valueDeclaration = expression.symbol.owner
-                            if (valueDeclaration is IrVariable && valueDeclaration in variablesSet) {
-                                forbiddenVariables.add(valueDeclaration)
-                            }
-                            super.visitValueAccess(expression)
+        fun recur(block: IrContainerExpression): Boolean /* stop optimization */ {
+            while (block.statements.isNotEmpty() && resultVariables.isNotEmpty()) {
+                val statement = block.statements.last()
+                //also stop
+                when {
+                    statement is IrContainerExpression -> if (recur(statement)) {
+                        if (statement.statements.isEmpty()) {
+                            block.statements.removeLast()
                        }
-                    })
+                        return true
+                    } else {
+                        require(statement.statements.isEmpty() || resultVariables.isEmpty()) { "Not all statements removed" }
+                    }
+
+                    statement !is IrSetValue -> return true
+                    statement.symbol.owner != resultVariables.last() -> return true
+                    statement.symbol.owner in forbiddenVariables -> return true
+                    else -> {
+                        standaloneExpressions.add(statement.value)
+                        resultVariables.removeLast()
+                        block.statements.removeLast()
+                        statement.value.acceptVoid(object : IrElementVisitorVoid {
+                            override fun visitElement(element: IrElement) {
+                                element.acceptChildrenVoid(this)
+                            }
+
+                            override fun visitValueAccess(expression: IrValueAccessExpression) {
+                                val valueDeclaration = expression.symbol.owner
+                                if (valueDeclaration is IrVariable && valueDeclaration in variablesSet) {
+                                    forbiddenVariables.add(valueDeclaration)
+                                }
+                                super.visitValueAccess(expression)
+                            }
+                        })
+                    }
                }
            }
+            return false
        }
-        repeat(standaloneExpressions.size) { block.statements.removeLast() }
-        return variables.dropLast(standaloneExpressions.size).map { irGet(it) } + standaloneExpressions.asReversed()
+        recur(block)
+        return resultVariables.map { irGet(it) } + standaloneExpressions.asReversed()
    }

    // Note that reference equality (x === y) is not allowed on values of MFVC class type,
@@ -862,27 +914,26 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV

    override fun visitGetField(expression: IrGetField): IrExpression {
        expression.receiver = expression.receiver?.transform(this, null)
-        return valueDeclarationsRemapper.makeReplacement(context.createIrBuilder(expression.symbol), expression)
-            ?: expression
+        return context.createIrBuilder(expression.symbol).irBlock {
+            valueDeclarationsRemapper.addReplacement(this, expression) ?: return expression
+        }
    }

    override fun visitSetField(expression: IrSetField): IrExpression {
        expression.receiver = expression.receiver?.transform(this, null)
-        return valueDeclarationsRemapper.makeReplacement(
-            context.createIrBuilder(getCurrentScopeSymbol()), expression, safe = expression.origin != UNSAFE_MFVC_SET_ORIGIN
-        ) ?: expression.also { it.value = it.value.transform(this, null) }
-    }
-
-    override fun visitGetValue(expression: IrGetValue): IrExpression = with(context.createIrBuilder(getCurrentScopeSymbol())) b@{
-        with(valueDeclarationsRemapper) {
-            this.makeReplacement(this@b, expression) ?: super.visitGetValue(expression)
+        return context.createIrBuilder(getCurrentScopeSymbol()).irBlock {
+            valueDeclarationsRemapper.addReplacement(this, expression, safe = expression.origin != UNSAFE_MFVC_SET_ORIGIN)
+                ?: return expression.also { it.value = it.value.transform(this@JvmMultiFieldValueClassLowering, null) }
        }
    }

-    override fun visitSetValue(expression: IrSetValue): IrExpression = with(context.createIrBuilder(getCurrentScopeSymbol())) b@{
-        with(valueDeclarationsRemapper) {
-            this.makeReplacement(this@b, expression, safe = expression.origin != UNSAFE_MFVC_SET_ORIGIN)
-        } ?: super.visitSetValue(expression)
+    override fun visitGetValue(expression: IrGetValue): IrExpression = with(context.createIrBuilder(getCurrentScopeSymbol())) b@{
+        valueDeclarationsRemapper.makeReplacement(this@b, expression) ?: super.visitGetValue(expression)
+    }
+
+    override fun visitSetValue(expression: IrSetValue): IrExpression = context.createIrBuilder(getCurrentScopeSymbol()).irBlock b@{
+        valueDeclarationsRemapper.addReplacement(this@b, expression, safe = expression.origin != UNSAFE_MFVC_SET_ORIGIN)
+            ?: return super.visitSetValue(expression)
    }

    override fun visitVariable(declaration: IrVariable): IrStatement {
@@ -921,7 +972,7 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                saveVariable = ::variablesSaver
            )
        }
-        val instance = ValueDeclarationMfvcNodeInstance(this, rootMfvcNode, typeArguments, variables, listOf())
+        val instance = ValueDeclarationMfvcNodeInstance(this, rootMfvcNode, typeArguments, variables)
        val block = irBlock {
            flattenExpressionTo(expression, instance)
        }
@@ -949,6 +1000,16 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
        require(rootNode.leavesCount == instance.size) {
            "Required ${rootNode.leavesCount} variable/field to store regular value but got ${instance.size}"
        }
+        if (expression is IrWhen) {
+            for (branch in expression.branches) {
+                branch.condition = branch.condition.transform(this@JvmMultiFieldValueClassLowering, null)
+                branch.result = irBlock {
+                    flattenExpressionTo(branch.result, instance)
+                }
+            }
+            +expression
+            return
+        }
        if (expression is IrConstructorCall) {
            val constructor = expression.symbol.owner
            if (constructor.isPrimary && constructor.constructedClass.isMultiFieldValueClass &&
@@ -971,24 +1032,15 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
            }
        }
        val transformedExpression = expression.transform(this@JvmMultiFieldValueClassLowering, null)
-        val flattenedExpressionInstance: MfvcNodeInstance? = valueDeclarationsRemapper.getMfvcNodeInstanceIfIsGetter(transformedExpression)
-            ?.also { flattenedInstance -> flattenedInstance.initializationStatements.forEach { +it } }
-        if (flattenedExpressionInstance == null && transformedExpression is IrContainerExpression && transformedExpression.statements.isNotEmpty()) {
-            val last = transformedExpression.statements.last()
-            if (last is IrExpression) {
-                fun irContainer(builder: IrBlockBuilder.() -> Unit) =
-                    if (transformedExpression.isTransparentScope) irComposite { builder() } else irBlock { builder() }
-
-                +irContainer {
-                    for (statement in transformedExpression.statements.dropLast(1)) {
-                        +statement
-                    }
-                    flattenExpressionTo(last, instance)
-                }
-                return
-            }
+        val addedSettersToFlattened = valueDeclarationsRemapper.handleFlattenedGetterExpressions(transformedExpression) {
+            require(it.size == instance.size) { "Incompatible assignment sizes: ${it.size}, ${instance.size}" }
+            instance.addSetterStatements(this, it)
        }
-        val expressionInstance = flattenedExpressionInstance ?: rootNode.createInstanceFromBox(
+        if (addedSettersToFlattened != null) {
+            +addedSettersToFlattened
+            return
+        }
+        val expressionInstance = rootNode.createInstanceFromBox(
            this, transformedExpression, getOptimizedPublicAccess(rootNode.mfvc), ::variablesSaver,
        )
        require(expressionInstance.size == instance.size) { "Incompatible assignment sizes: ${expressionInstance.size}, ${instance.size}" }
@@ -1028,6 +1080,15 @@ private class JvmMultiFieldValueClassLowering(context: JvmBackendContext) : JvmV
                handleStatementContainer(expression, data)
            }

+            override fun visitWhen(expression: IrWhen, data: Boolean) {
+                expression.acceptChildren(this, data) // when's are transparent
+            }
+
+            override fun visitBranch(branch: IrBranch, data: Boolean) {
+                branch.condition.accept(this, true)
+                branch.result.accept(this, data)
+            }
+
            override fun visitBlockBody(body: IrBlockBody, data: Boolean) {
                handleStatementContainer(body, data)
            }
@@ -1110,9 +1171,9 @@ private fun findNearestBlocksForVariables(variables: Set<IrVariable>, body: IrBo
 private fun IrStatement.containsUsagesOf(variablesSet: Set<IrVariable>): Boolean {
    var used = false
    acceptVoid(object : IrElementVisitorVoid {
-        override fun visitExpression(expression: IrExpression) {
+        override fun visitElement(element: IrElement) {
            if (!used) {
-                expression.acceptChildrenVoid(this)
+                element.acceptChildrenVoid(this)
            }
        }

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.backend.jvm
 import org.jetbrains.kotlin.backend.jvm.ir.erasedUpperBound
 import org.jetbrains.kotlin.backend.jvm.ir.isMultiFieldValueClassType
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
-import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.IrBlockBuilder
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irCall
@@ -61,13 +60,13 @@ fun MfvcNode.createInstanceFromValueDeclarations(
            saveVariable = saveVariable
        )
    }
-    return ValueDeclarationMfvcNodeInstance(scope, this, typeArguments, valueDeclarations, listOf())
+    return ValueDeclarationMfvcNodeInstance(scope, this, typeArguments, valueDeclarations)
 }

 fun MfvcNode.createInstanceFromValueDeclarationsAndBoxType(
-    scope: IrBuilderWithScope, type: IrSimpleType, fieldValues: List<IrValueDeclaration>, initializationStatements: List<IrStatement>
+    scope: IrBuilderWithScope, type: IrSimpleType, fieldValues: List<IrValueDeclaration>
 ): ValueDeclarationMfvcNodeInstance =
-    ValueDeclarationMfvcNodeInstance(scope, this, makeTypeArgumentsFromType(type), fieldValues, initializationStatements)
+    ValueDeclarationMfvcNodeInstance(scope, this, makeTypeArgumentsFromType(type), fieldValues)

 fun makeTypeArgumentsFromType(type: IrSimpleType): TypeArguments {
    if (type.classifierOrNull !is IrClassSymbol) return mapOf()
@@ -26,7 +26,6 @@ interface MfvcNodeInstance {
    val node: MfvcNode
    val typeArguments: TypeArguments
    val type: IrSimpleType
-    val initializationStatements: List<IrStatement>

    fun makeFlattenedGetterExpressions(): List<IrExpression>
    fun makeGetterExpression(): IrExpression
@@ -56,7 +55,6 @@ class ValueDeclarationMfvcNodeInstance(
    override val node: MfvcNode,
    override val typeArguments: TypeArguments,
    val valueDeclarations: List<IrValueDeclaration>,
-    override val initializationStatements: List<IrStatement>,
 ) : MfvcNodeInstance {
    init {
        require(valueDeclarations.size == size) { "Expected value declarations list of size $size but got of size ${valueDeclarations.size}" }
@@ -73,7 +71,7 @@ class ValueDeclarationMfvcNodeInstance(

    override fun get(name: Name): ValueDeclarationMfvcNodeInstance? {
        val (newNode, indices) = node.getSubnodeAndIndices(name) ?: return null
-        return ValueDeclarationMfvcNodeInstance(scope, newNode, typeArguments, valueDeclarations.slice(indices), initializationStatements)
+        return ValueDeclarationMfvcNodeInstance(scope, newNode, typeArguments, valueDeclarations.slice(indices))
    }

    override fun makeStatements(values: List<IrExpression>): List<IrStatement> {
@@ -86,7 +84,6 @@ internal class ExpressionCopierImpl(
    expression: IrExpression?,
    private val scope: IrBlockBuilder,
    private val saveVariable: (IrVariable) -> Unit,
-    private val saveReceiverSetter: (IrSetValue) -> Unit
 ) {
    private sealed interface CopyableExpression {
        fun makeExpression(scope: IrBuilderWithScope): IrExpression
@@ -109,7 +106,6 @@ internal class ExpressionCopierImpl(
                origin = IrDeclarationOrigin.TEMPORARY_MULTI_FIELD_VALUE_CLASS_VARIABLE,
                saveVariable = saveVariable,
                isTemporary = true,
-                saveSetter = saveReceiverSetter,
            )
        )

@@ -156,10 +152,7 @@ class ReceiverBasedMfvcNodeInstance(
 ) : MfvcNodeInstance {
    override val type: IrSimpleType = makeTypeFromMfvcNodeAndTypeArguments(node, typeArguments)

-    override val initializationStatements: List<IrStatement> = mutableListOf()
-    private val receiverCopier = ExpressionCopierImpl(
-        receiver, scope, saveVariable, saveReceiverSetter = { (initializationStatements as MutableList).add(it) }
-    )
+    private val receiverCopier = ExpressionCopierImpl(receiver, scope, saveVariable)

    private fun makeReceiverCopy() = receiverCopier.makeCopy()

@@ -280,8 +273,7 @@ fun <T : IrElement> IrStatementsBuilder<T>.savableStandaloneVariableWithSetter(
    isMutable: Boolean = false,
    origin: IrDeclarationOrigin,
    isTemporary: Boolean = origin == IrDeclarationOrigin.IR_TEMPORARY_VARIABLE,
-    saveSetter: (IrSetValue) -> Unit = {},
    saveVariable: (IrVariable) -> Unit,
 ) = savableStandaloneVariable(expression.type, name, isMutable, origin, isTemporary, saveVariable).also {
-    +irSet(it, expression).also(saveSetter)
+    +irSet(it, expression)
 }
@@ -0,0 +1,53 @@
+// CHECK_BYTECODE_LISTING
+// WITH_STDLIB
+// TARGET_BACKEND: JVM_IR
+// WORKS_WHEN_VALUE_CLASS
+// LANGUAGE: +ValueClasses
+
+@JvmInline
+value class DPoint(val x: Double, val y: Double) {
+    init {
+        counter++
+    }
+    companion object {
+        @JvmStatic
+        var counter: Int = 0
+    }
+}
+
+fun ifExpr() = if (DPoint(0.0, 1.0).x > 0.0) DPoint(2.0, 3.0) else DPoint(4.0, 5.0)
+
+fun whenExpr() = when {
+    DPoint(6.0, 7.0).x > 0.0 -> DPoint(8.0, 9.0)
+    DPoint(10.0, 11.0).x > 0.0 -> DPoint(12.0, 13.0)
+    else -> DPoint(14.0, 15.0)
+}
+
+fun ifBody() {
+    if (DPoint(0.0, 1.0).x > 0.0) DPoint(2.0, 3.0) else DPoint(4.0, 5.0)
+    val x = if (DPoint(0.0, 1.0).x > 0.0) DPoint(2.0, 3.0) else DPoint(4.0, 5.0)
+    require(x == DPoint(4.0, 5.0))
+}
+
+fun whenBody() {
+    when {
+        DPoint(6.0, 7.0).x > 0.0 -> DPoint(8.0, 9.0)
+        DPoint(10.0, 11.0).x > 0.0 -> DPoint(12.0, 13.0)
+        else -> DPoint(14.0, 15.0)
+    }
+    val x = when {
+        DPoint(6.0, 7.0).x > 0.0 -> DPoint(8.0, 9.0)
+        DPoint(10.0, 11.0).x > 0.0 -> DPoint(12.0, 13.0)
+        else -> DPoint(14.0, 15.0)
+    }
+    require(x == DPoint(8.0, 9.0))
+}
+
+fun box(): String {
+    ifBody()
+    whenBody()
+    require(ifExpr() == DPoint(4.0, 5.0))
+    require(whenExpr() == DPoint(8.0, 9.0))
+    require(DPoint.counter == 5 + 5 + 2 + 1 + 2 + 1)
+    return "OK"
+}
@@ -0,0 +1,51 @@
+@kotlin.Metadata
+public final class ConditionalExpressionsKt {
+    // source: 'conditionalExpressions.kt'
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method ifBody(): void
+    public final static @org.jetbrains.annotations.NotNull method ifExpr(): DPoint
+    public final static method whenBody(): void
+    public final static @org.jetbrains.annotations.NotNull method whenExpr(): DPoint
+    public final inner class DPoint$Companion
+}
+
+@kotlin.Metadata
+public final class DPoint$Companion {
+    // source: 'conditionalExpressions.kt'
+    private method <init>(): void
+    public synthetic method <init>(p0: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public synthetic deprecated static @kotlin.jvm.JvmStatic method getCounter$annotations(): void
+    public final method getCounter(): int
+    public final method setCounter(p0: int): void
+    public final inner class DPoint$Companion
+}
+
+@kotlin.jvm.JvmInline
+@kotlin.Metadata
+public final class DPoint {
+    // source: 'conditionalExpressions.kt'
+    public final static @org.jetbrains.annotations.NotNull field Companion: DPoint$Companion
+    private static field counter: int
+    private final field field-0: double
+    private final field field-1: double
+    static method <clinit>(): void
+    private synthetic method <init>(p0: double, p1: double): void
+    public synthetic final static method access$getCounter$cp(): int
+    public synthetic final static method access$setCounter$cp(p0: int): void
+    public synthetic final static method box-impl(p0: double, p1: double): DPoint
+    public final static method constructor-impl(p0: double, p1: double): void
+    public method equals(@org.jetbrains.annotations.Nullable p0: java.lang.Object): boolean
+    public static method equals-impl(p0: double, p1: double, p2: java.lang.Object): boolean
+    public final static method equals-impl0(p0: double, p1: double, p2: double, p3: double): boolean
+    public final static method getCounter(): int
+    public final static method getX-impl(p0: double, p1: double): double
+    public final static method getY-impl(p0: double, p1: double): double
+    public method hashCode(): int
+    public static method hashCode-impl(p0: double, p1: double): int
+    public final static method setCounter(p0: int): void
+    public @org.jetbrains.annotations.NotNull method toString(): java.lang.String
+    public static method toString-impl(p0: double, p1: double): java.lang.String
+    public synthetic final method unbox-impl-0(): double
+    public synthetic final method unbox-impl-1(): double
+    public final inner class DPoint$Companion
+}
@@ -26,6 +26,7 @@ public final class Internal {
    public final static @org.jetbrains.annotations.NotNull field Companion: Internal$Companion
    private final field field-0: int
    private final field field-1: int
+    private static field x: int
    static method <clinit>(): void
    private synthetic method <init>(p0: int, p1: int): void
    public synthetic final static method access$getX$cp(): int
@@ -228,6 +229,7 @@ public final class Public {
    public final static @org.jetbrains.annotations.NotNull field Companion: Public$Companion
    private final field field-0: int
    private final field field-1: int
+    private static field x: int
    static method <clinit>(): void
    private synthetic method <init>(p0: int, p1: int): void
    public synthetic final static method access$getX$cp(): int
@@ -0,0 +1,41 @@
+// CHECK_BYTECODE_LISTING
+// WITH_STDLIB
+// TARGET_BACKEND: JVM_IR
+// WORKS_WHEN_VALUE_CLASS
+// LANGUAGE: +ValueClasses
+
+@JvmInline
+value class DPoint(val x: Double, val y: Double)
+
+fun ifExpr() = if (DPoint(0.0, 1.0).x > 0.0) DPoint(2.0, 3.0) else DPoint(4.0, 5.0)
+
+fun whenExpr() = when {
+    DPoint(6.0, 7.0).x > 0.0 -> DPoint(8.0, 9.0)
+    DPoint(10.0, 11.0).x > 0.0 -> DPoint(12.0, 13.0)
+    else -> DPoint(14.0, 15.0)
+}
+
+fun ifBody() {
+    if (DPoint(0.0, 1.0).x > 0.0) DPoint(2.0, 3.0) else DPoint(4.0, 5.0)
+    val x = if (DPoint(0.0, 1.0).x > 0.0) DPoint(2.0, 3.0) else DPoint(4.0, 5.0)
+}
+
+fun whenBody() {
+    when {
+        DPoint(6.0, 7.0).x > 0.0 -> DPoint(8.0, 9.0)
+        DPoint(10.0, 11.0).x > 0.0 -> DPoint(12.0, 13.0)
+        else -> DPoint(14.0, 15.0)
+    }
+    val x = when {
+        DPoint(6.0, 7.0).x > 0.0 -> DPoint(8.0, 9.0)
+        DPoint(10.0, 11.0).x > 0.0 -> DPoint(12.0, 13.0)
+        else -> DPoint(14.0, 15.0)
+    }
+}
+
+// 1 ifExpr.*(\n  .*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)
+// 0 ifExpr.*(\n  .*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)
+// 1 whenExpr.*(\n  .*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)
+// 0 whenExpr.*(\n  .*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)(\n   .*)*(\n   .*box-impl.*)
+// 0 ifBody.*(\n   .*)*(\n   .*box-impl.*)
+// 0 whenBody.*(\n   .*)*(\n   .*box-impl.*)
@@ -50195,6 +50195,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
            runTest("compiler/testData/codegen/box/valueClasses/complex.kt", TransformersFunctions.getReplaceOptionalJvmInlineAnnotationWithReal());
        }

+        @Test
+        @TestMetadata("conditionalExpressions.kt")
+        public void testConditionalExpressions() throws Exception {
+            runTest("compiler/testData/codegen/box/valueClasses/conditionalExpressions.kt", TransformersFunctions.getReplaceOptionalJvmInlineAnnotationWithReal());
+        }
+
        @Test
        @TestMetadata("equality.kt")
        public void testEquality() throws Exception {
@@ -5789,6 +5789,12 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/bytecodeText/valueClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
        }

+        @Test
+        @TestMetadata("conditionalExpressions.kt")
+        public void testConditionalExpressions() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/valueClasses/conditionalExpressions.kt", TransformersFunctions.getReplaceOptionalJvmInlineAnnotationWithReal());
+        }
+
        @Test
        @TestMetadata("equalsBoxTest.kt")
        public void testEqualsBoxTest() throws Exception {