From 7c63105bb76839e2a78851e70ab0b1230fd0024c Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Tue, 9 Mar 2021 21:55:01 +0100
Subject: [PATCH] Mark script result expression as used to avoid coercion to
 unit

---
 .../kotlin/cfg/ControlFlowProcessor.kt        | 30 +++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/compiler/frontend/cfg/src/org/jetbrains/kotlin/cfg/ControlFlowProcessor.kt b/compiler/frontend/cfg/src/org/jetbrains/kotlin/cfg/ControlFlowProcessor.kt
index f1e784c78c1..bf130f8ce39 100644
--- a/compiler/frontend/cfg/src/org/jetbrains/kotlin/cfg/ControlFlowProcessor.kt
+++ b/compiler/frontend/cfg/src/org/jetbrains/kotlin/cfg/ControlFlowProcessor.kt
@@ -43,10 +43,9 @@ import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.isAncestor
+import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.BindingContext.USED_AS_EXPRESSION
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.CompileTimeConstantUtils
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getEnclosingFunctionDescriptor
@@ -61,6 +60,8 @@ import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluat
 import org.jetbrains.kotlin.resolve.scopes.receivers.*
 import org.jetbrains.kotlin.types.expressions.DoubleColonLHS
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
+import org.jetbrains.kotlin.types.typeUtil.isNothing
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.util.slicedMap.ReadOnlySlice
 import java.util.*
 
@@ -925,7 +926,7 @@ class ControlFlowProcessor(
                 if (builder.getLoopExitPoint(jumpTarget) != null ||
                     // Local class secondary constructors are handled differently
                     // They are the only local class element NOT included in owner pseudocode
-                    // See generateInitializersForScriptClassOrObject && generateDeclarationForLocalClassOrObjectIfNeeded
+                    // See generateInitializersForClassOrObject && generateDeclarationForLocalClassOrObjectIfNeeded
                     labelExprEnclosingFunc is ConstructorDescriptor && !labelExprEnclosingFunc.isPrimary
                 ) {
                     trace.report(BREAK_OR_CONTINUE_JUMPS_ACROSS_FUNCTION_BOUNDARY.on(jumpExpression))
@@ -1409,8 +1410,27 @@ class ControlFlowProcessor(
         }
 
         override fun visitScript(script: KtScript) {
+
+            // the same logic is implemented in the LazyScriptDescriptor
+            // TODO: consider extracting last expression type logic to a common place
+            val lastInitializer = script
+                .getChildOfType<KtBlockExpression>()
+                ?.getChildrenOfType<KtScriptInitializer>()?.lastOrNull()
+            val resultExpression = lastInitializer?.getChildOfType<KtExpression>()
+
+            val resultType = resultExpression?.let {
+                trace.bindingContext.getType(it)
+            }
+
+            val hasResultField = resultType != null && !resultType.isUnit() && !resultType.isNothing()
+
             for (declaration in script.declarations) {
-                if (declaration is KtProperty || declaration is KtAnonymousInitializer || declaration is KtDestructuringDeclaration) {
+                if (declaration is KtAnonymousInitializer) {
+                    generateInstructions(declaration)
+                    if (hasResultField && declaration == lastInitializer) {
+                        trace.record(USED_AS_EXPRESSION, resultExpression, true)
+                    }
+                } else if (declaration is KtProperty || declaration is KtDestructuringDeclaration) {
                     generateInstructions(declaration)
                 }
             }