[IR] Rewrite how interpreter works with the Unit result

This way interpreter works in more correct way. Every expression
produces some result, and if we don't need it, the interpreter will
discard it.

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/InstructionsUnfolder.kt

@@ -18,15 +18,12 @@ import org.jetbrains.kotlin.ir.interpreter.exceptions.verify
 import org.jetbrains.kotlin.ir.interpreter.stack.CallStack
 import org.jetbrains.kotlin.ir.interpreter.state.*
 import org.jetbrains.kotlin.ir.types.classOrNull
-import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.types.makeNullable
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
 
-internal fun IrExpression.handleAndDropResult(callStack: CallStack, dropOnlyUnit: Boolean = false) {
-    val dropResult = fun() {
-        if (!dropOnlyUnit && !this.type.isUnit() || callStack.peekState().isUnit()) callStack.popState()
-    }
+internal fun IrExpression.handleAndDropResult(callStack: CallStack) {
+    val dropResult = fun() { callStack.popState() }
     callStack.pushInstruction(CustomInstruction(dropResult))
     callStack.pushCompoundInstruction(this)
 }
@@ -212,6 +209,7 @@ private fun unfoldInstanceInitializerCall(instanceInitializerCall: IrInstanceIni
     val toInitialize = irClass.declarations.filter { it is IrProperty || it is IrAnonymousInitializer }
     val state = irClass.thisReceiver?.symbol?.let { callStack.loadState(it) } // try to avoid recalculation of properties
 
+    callStack.pushSimpleInstruction(instanceInitializerCall)
     toInitialize.reversed().forEach {
         when {
             it is IrAnonymousInitializer -> callStack.pushCompoundInstruction(it.body)
@@ -258,7 +256,7 @@ private fun unfoldStatements(statements: List<IrStatement>, callStack: CallStack
             is IrExpression ->
                 when {
                     i.isLastIndex() -> callStack.pushCompoundInstruction(statement)
-                    else -> statement.handleAndDropResult(callStack, dropOnlyUnit = true)
+                    else -> statement.handleAndDropResult(callStack)
                 }
             else -> callStack.pushCompoundInstruction(statement)
         }
@@ -435,7 +433,7 @@ private fun unfoldStringConcatenation(expression: IrStringConcatenation, environ
 private fun unfoldComposite(element: IrComposite, callStack: CallStack) {
     when (element.origin) {
         IrStatementOrigin.DESTRUCTURING_DECLARATION, IrStatementOrigin.DO_WHILE_LOOP, null -> // is null for body of do while loop
-            element.statements.reversed().forEach { callStack.pushCompoundInstruction(it) }
+            unfoldStatements(element.statements, callStack)
         else -> TODO("${element.origin} not implemented")
     }
 }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/IrInterpreter.kt

@@ -63,7 +63,7 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
             callStack.popInstruction().handle()
         }
 
-        return environment.stateToIrExpression(callStack.popState(), expression).apply { callStack.dropFrame() }
+        return environment.stateToIrExpression(extractResultAndAssertThatStackIsEmpty(), expression)
     }
 
     internal fun withNewCallStack(call: IrCall, init: IrInterpreter.() -> Any?): State {
@@ -75,10 +75,17 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
                 callStack.popInstruction().handle()
             }
 
-            callStack.popState().apply { callStack.dropFrame() }
+            extractResultAndAssertThatStackIsEmpty()
         }
     }
 
+    private fun extractResultAndAssertThatStackIsEmpty(): State {
+        val result = callStack.popState()
+        assert(callStack.peekState() == null)
+        callStack.dropFrame()
+        return result
+    }
+
     private fun interpret(element: IrElement) {
         when (element) {
             is IrSimpleFunction -> interpretFunction(element)
@@ -87,6 +94,7 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
             is IrConstructorCall -> interpretConstructorCall(element)
             is IrEnumConstructorCall -> interpretEnumConstructorCall(element)
             is IrDelegatingConstructorCall -> interpretDelegatingConstructorCall(element)
+            is IrInstanceInitializerCall -> callStack.pushState(getUnitState())
             is IrValueParameter -> interpretValueParameter(element)
             is IrField -> interpretField(element)
             is IrBody -> interpretBody(element)
@@ -98,13 +106,13 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
             is IrGetEnumValue -> interpretGetEnumValue(element)
             is IrEnumEntry -> interpretEnumEntry(element)
             is IrConst<*> -> interpretConst(element)
-            is IrVariable -> callStack.storeState(element.symbol, callStack.popState())
-            is IrSetValue -> callStack.rewriteState(element.symbol, callStack.popState())
+            is IrVariable -> interpretVariable(element)
+            is IrSetValue -> interpretSetValue(element)
             is IrTypeOperatorCall -> interpretTypeOperatorCall(element)
             is IrBranch -> interpretBranch(element)
             is IrWhileLoop -> interpretWhile(element)
             is IrDoWhileLoop -> interpretDoWhile(element)
-            is IrWhen -> callStack.dropSubFrame()
+            is IrWhen -> interpretWhen(element)
             is IrVararg -> interpretVararg(element)
             is IrTry -> interpretTry(element)
             is IrThrow -> interpretThrow(element)
@@ -198,15 +206,13 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
         receiverState.setField(field.correspondingPropertySymbol!!, callStack.popState())
     }
 
-    @Suppress("UNUSED_PARAMETER")
     private fun interpretBody(body: IrBody) {
-        if (callStack.peekState() == null) callStack.pushState(getUnitState()) // implicit Unit result
+        if (body.statements.isEmpty()) callStack.pushState(getUnitState()) // implicit Unit result
     }
 
-    @Suppress("UNUSED_PARAMETER")
     private fun interpretBlock(block: IrBlock) {
         callStack.dropSubFrame()
-        if (callStack.peekState() == null) callStack.pushState(getUnitState()) // implicit Unit result
+        if (block.statements.isEmpty()) callStack.pushState(getUnitState()) // implicit Unit result
     }
 
     private fun interpretConstructor(constructor: IrConstructor) {
@@ -267,6 +273,8 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
         callInterceptor.interceptConstructor(constructorCall, valueArguments) {
             callStack.pushCompoundInstruction(constructor)
         }
+
+        callStack.pushState(getUnitState())
     }
 
     private fun interpretDelegatingConstructorCall(constructorCall: IrDelegatingConstructorCall) {
@@ -300,6 +308,16 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
         callStack.pushState(expression.toPrimitive())
     }
 
+    private fun interpretVariable(variable: IrVariable) {
+        callStack.storeState(variable.symbol, callStack.popState())
+        callStack.pushState(getUnitState())
+    }
+
+    private fun interpretSetValue(expression: IrSetValue) {
+        callStack.rewriteState(expression.symbol, callStack.popState())
+        callStack.pushState(getUnitState())
+    }
+
     private fun interpretReturn(expression: IrReturn) {
         callStack.returnFromFrameWithResult(expression)
     }
@@ -326,10 +344,18 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
         }
     }
 
+    @Suppress("UNUSED_PARAMETER")
+    private fun interpretWhen(whenExpression: IrWhen) {
+        // This method is reachable only if none of the branches were selected.
+        // In that case, we just return `Unit` result that will be dropped later.
+        callStack.dropSubFrame()
+        callStack.pushState(getUnitState())
+    }
+
     private fun interpretBranch(branch: IrBranch) {
         val result = callStack.popState().asBoolean()
         if (result) {
-            callStack.dropSubFrame()
+            callStack.dropSubFrame() // drop entire `when` expression from frame
             callStack.pushCompoundInstruction(branch.result)
         }
     }
@@ -338,6 +364,7 @@ class IrInterpreter(internal val environment: IrInterpreterEnvironment, internal
         val receiver = (expression.receiver as IrDeclarationReference).symbol
         val propertySymbol = expression.symbol.owner.correspondingPropertySymbol!!
         callStack.loadState(receiver).apply { this.setField(propertySymbol, callStack.popState()) }
+        callStack.pushState(getUnitState())
     }
 
     private fun interpretGetField(expression: IrGetField) {