Make interpreted functions implicitly return Unit as result

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

@@ -192,11 +192,18 @@ private fun unfoldBlock(block: IrBlock, callStack: CallStack) {
 }
 
 private fun unfoldStatements(statements: List<IrStatement>, callStack: CallStack) {
+    fun dropUnitResult() {
+        if (callStack.peekState().isUnit()) callStack.popState()
+    }
+
     for (i in statements.indices.reversed()) {
         when (val statement = statements[i]) {
             is IrClass -> if (!statement.isLocal) TODO("Only local classes are supported")
             is IrFunction -> if (!statement.isLocal) TODO("Only local functions are supported")
-            else -> callStack.addInstruction(CompoundInstruction(statement))
+            else -> {
+                callStack.addInstruction(CustomInstruction(::dropUnitResult))
+                callStack.addInstruction(CompoundInstruction(statement))
+            }
         }
     }
 }

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

@@ -57,9 +57,7 @@ class IrInterpreter private constructor(
         if (commandCount >= IrInterpreterEnvironment.MAX_COMMANDS) InterpreterTimeOutError().handleUserException(environment)
     }
 
-    private fun getUnitState(): State {
+    private fun getUnitState(): State = environment.mapOfObjects[irBuiltIns.unitClass]!!
-        return environment.mapOfObjects.getOrPut(irBuiltIns.unitClass) { Common(irBuiltIns.unitClass.owner) }
-    }
 
     private fun Instruction.handle() {
         when (this) {
@@ -91,7 +89,7 @@ class IrInterpreter private constructor(
                 incrementAndCheckCommands()
             }
 
-            if (callStack.peekState() == null) callStack.pushState(getUnitState()) // TODO maybe move this logic to body/block
+            if (callStack.peekState() == null) callStack.pushState(getUnitState())
             callStack.popState().apply { callStack.dropFrame() }
         }
     }
@@ -141,9 +139,8 @@ class IrInterpreter private constructor(
 
     private fun interpretFunction(function: IrSimpleFunction) {
         function.tryResetFunctionBody()
-        if (function.checkCast(environment)) {
+        if (callStack.peekState() == null) callStack.pushState(getUnitState()) // implicit Unit return
-            callStack.dropFrameAndCopyResult()
+        if (function.checkCast(environment)) callStack.dropFrameAndCopyResult()
-        }
     }
 
     private fun interpretValueParameter(valueParameter: IrValueParameter) {
@@ -251,7 +248,7 @@ class IrInterpreter private constructor(
         if (irClass.isLocal) callStack.storeUpValues(objectVar.state as StateWithClosure)
         val outerClass = if (irClass.isInner) callStack.popState() else null
 
-        callStack.dropSubFrame() // TODO check that data stack is empty
+        callStack.dropSubFrame()
         callStack.newFrame(constructor, listOf(SimpleInstruction(constructor)))
         callStack.addVariable(objectVar)
         constructor.valueParameters.forEachIndexed { i, param -> callStack.addVariable(Variable(param.symbol, valueArguments[i])) }

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.interpreter.stack.CallStack
+import org.jetbrains.kotlin.ir.interpreter.state.Common
 import org.jetbrains.kotlin.ir.interpreter.state.Complex
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.isSubclassOf
@@ -18,6 +19,10 @@ internal class IrInterpreterEnvironment(val irBuiltIns: IrBuiltIns, val callStac
     var mapOfEnums = mutableMapOf<IrSymbol, Complex>()
     var mapOfObjects = mutableMapOf<IrSymbol, Complex>()
 
+    init {
+        mapOfObjects[irBuiltIns.unitClass] = Common(irBuiltIns.unitClass.owner)
+    }
+
     private constructor(environment: IrInterpreterEnvironment) : this(environment.irBuiltIns, CallStack()) {
         irExceptions.addAll(environment.irExceptions)
         mapOfEnums = environment.mapOfEnums

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/stack/Frame.kt

@@ -66,7 +66,7 @@ internal class Frame(subFrame: SubFrame, val irFile: IrFile? = null) {
     }
 
     fun getVariable(symbol: IrSymbol): Variable {
-        return innerStack.firstNotNullResult { it.getVariable(symbol) }
+        return innerStack.reversed().firstNotNullResult { it.getVariable(symbol) }
             ?: throw InterpreterError("$symbol not found") // TODO better message
     }
 

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/state/State.kt

@@ -34,6 +34,7 @@ internal interface State {
 }
 
 internal fun State.isNull() = this is Primitive<*> && this.value == null
+internal fun State?.isUnit() = this is Common && this.irClassFqName() == "kotlin.Unit"
 
 internal fun State.asInt() = (this as Primitive<*>).value as Int
 internal fun State.asBoolean() = (this as Primitive<*>).value as Boolean