IR: create shared variables for val-variables when needed

This is possible when a lambda's contract guarantees initialization of a
variable.

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt

@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
 import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.visitors.*
-import java.util.*
 
 val sharedVariablesPhase = makeIrFilePhase(
     ::SharedVariablesLowering,
@@ -77,7 +76,8 @@ class SharedVariablesLowering(val context: BackendContext) : FileLoweringPass {
 
         private fun collectSharedVariables() {
             irDeclaration.accept(object : IrElementVisitor<Unit, IrDeclarationParent?> {
-                val relevantVars = mutableSetOf<IrVariable>()
+                val relevantVars = HashSet<IrVariable>()
+                val relevantVals = HashSet<IrVariable>()
 
                 override fun visitElement(element: IrElement, data: IrDeclarationParent?) {
                     element.acceptChildren(this, data)
@@ -102,6 +102,12 @@ class SharedVariablesLowering(val context: BackendContext) : FileLoweringPass {
 
                     if (declaration.isVar) {
                         relevantVars.add(declaration)
+                    } else if (declaration.initializer == null) {
+                        // A val-variable can be initialized from another container (and thus can require shared variable transformation)
+                        // in case that container is a lambda with a corresponding contract, e.g. with invocation kind EXACTLY_ONCE.
+                        // Here, we collect all val-variables without immediate initializer to relevantVals, and later we copy only those
+                        // variables which are initialized in a foreign container, to sharedVariables.
+                        relevantVals.add(declaration)
                     }
                 }
 
@@ -113,6 +119,15 @@ class SharedVariablesLowering(val context: BackendContext) : FileLoweringPass {
                         sharedVariables.add(value)
                     }
                 }
+
+                override fun visitSetVariable(expression: IrSetVariable, data: IrDeclarationParent?) {
+                    super.visitSetVariable(expression, data)
+
+                    val variable = expression.symbol.owner
+                    if (variable.initializer == null && variable.parent != data && variable in relevantVals) {
+                        sharedVariables.add(variable)
+                    }
+                }
             }, irDeclaration as? IrDeclarationParent ?: irDeclaration.parent)
         }
 

compiler/testData/codegen/box/contracts/exactlyOnceNotInline.kt

@@ -1,6 +1,6 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR, NATIVE, JS_IR
+// IGNORE_BACKEND: NATIVE
 
 import kotlin.contracts.*
 

compiler/testData/codegen/boxInline/contracts/cfgDependendValInitialization.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -16,8 +16,8 @@ public inline fun <R> myrun(block: () -> R): R {
     return block()
 }
 
-
 // FILE: 2.kt
+
 import test.*
 
 fun test(b: Boolean): Int {
@@ -34,7 +34,6 @@ fun test(b: Boolean): Int {
 }
 
 fun box(): String {
-
     if (test(true) != 1) return "Fail 1"
     if (test(false) != -1) return "Fail 2"
     return "OK"

compiler/testData/codegen/boxInline/contracts/complexInitializer.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*

compiler/testData/codegen/boxInline/contracts/complexInitializerWithStackTransformation.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*

compiler/testData/codegen/boxInline/contracts/definiteLongValInitialization.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -17,6 +17,7 @@ public inline fun <R> myrun(block: () -> R): R {
 }
 
 // FILE: 2.kt
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/boxInline/contracts/definiteNestedValInitialization.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -16,8 +16,8 @@ public inline fun <R> myrun(block: () -> R): R {
     return block()
 }
 
-
 // FILE: 2.kt
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/boxInline/contracts/definiteValInitInInitializer.kt

@@ -1,10 +1,10 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
 // NO_CHECK_LAMBDA_INLINING
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -17,8 +17,8 @@ public inline fun <R> myrun(block: () -> R): R {
     return block()
 }
 
-
 // FILE: 2.kt
+
 import test.*
 
 class A {

compiler/testData/codegen/boxInline/contracts/definiteValInitialization.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -16,8 +16,8 @@ public inline fun <R> myrun(block: () -> R): R {
     return block()
 }
 
-
 // FILE: 2.kt
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/boxInline/contracts/exactlyOnceCrossinline.kt

@@ -1,8 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR, NATIVE, JS_IR
+// IGNORE_BACKEND: NATIVE
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -17,6 +18,7 @@ public inline fun myrun(crossinline block: () -> Unit): Unit {
 }
 
 // FILE: 2.kt
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/boxInline/contracts/exactlyOnceNoinline.kt

@@ -1,9 +1,10 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR, NATIVE, JS_IR
+// IGNORE_BACKEND: NATIVE
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // NO_CHECK_LAMBDA_INLINING
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -17,6 +18,7 @@ public inline fun myrun(noinline block: () -> Unit): Unit {
 }
 
 // FILE: 2.kt
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/boxInline/contracts/nonLocalReturn.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -17,6 +17,7 @@ public inline fun <R> myrun(block: () -> R): R {
 }
 
 // FILE: 2.kt
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/boxInline/contracts/nonLocalReturnWithCycle.kt

@@ -1,11 +1,11 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
 // WITH_RUNTIME
 // KJS_WITH_FULL_RUNTIME
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -19,6 +19,7 @@ public inline fun <R> myrun(block: () -> R): R {
 }
 
 // FILE: 2.kt
+
 import test.*
 
 fun test(xs: List<String>): String {
@@ -38,7 +39,6 @@ fun test(xs: List<String>): String {
     return result
 }
 
-
 fun box(): String {
     return test(listOf("O", "K", "fail"))
 }

compiler/testData/codegen/boxInline/contracts/valInitializationAndUsageInNestedLambda.kt

@@ -1,9 +1,9 @@
 // !LANGUAGE: +AllowContractsForCustomFunctions +UseCallsInPlaceEffect +ReadDeserializedContracts
 // !USE_EXPERIMENTAL: kotlin.contracts.ExperimentalContracts
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // IGNORE_BACKEND: NATIVE
+
 // FILE: 1.kt
+
 package test
 
 import kotlin.contracts.*
@@ -16,9 +16,9 @@ public inline fun <R> myrun(block: () -> R): R {
     return block()
 }
 
-
 // FILE: 2.kt
 // NO_CHECK_LAMBDA_INLINING
+
 import test.*
 
 fun box(): String {

compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedValInLambdaInitializedInside.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+// IGNORE_BACKEND: JVM_IR
+
+// In JVM IR, SharedVariablesLowering transforms `x` into a shared variable to be able to update it from a lambda,
+// which is a separate function (...$lambda-0).
+// If we keep the existing representation of lambda bodies as separate functions in JVM IR, the only viable option to fix this test
+// seems to support this case in the bytecode optimization pass CapturedVarsOptimizationMethodTransformer.
+
+fun box(): String {
+    val x: String
+    run {
+        x = "OK"
+        val y = x
+    }
+    return x
+}
+
+// 0 ObjectRef

compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedValInLambdaInitializedOutside.kt

@@ -0,0 +1,10 @@
+fun box(): String {
+    val x: String
+    x = "OK"
+    {
+        val y = x
+    }()
+    return x
+}
+
+// 0 ObjectRef

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -777,6 +777,16 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedInNoInlneInsideChainOfInlineFuns.kt");
         }
 
+        @TestMetadata("capturedValInLambdaInitializedInside.kt")
+        public void testCapturedValInLambdaInitializedInside() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedValInLambdaInitializedInside.kt");
+        }
+
+        @TestMetadata("capturedValInLambdaInitializedOutside.kt")
+        public void testCapturedValInLambdaInitializedOutside() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedValInLambdaInitializedOutside.kt");
+        }
+
         @TestMetadata("capturedVarsOfSize2.kt")
         public void testCapturedVarsOfSize2() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedVarsOfSize2.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBytecodeTextTestGenerated.java

@@ -777,6 +777,16 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedInNoInlneInsideChainOfInlineFuns.kt");
         }
 
+        @TestMetadata("capturedValInLambdaInitializedInside.kt")
+        public void testCapturedValInLambdaInitializedInside() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedValInLambdaInitializedInside.kt");
+        }
+
+        @TestMetadata("capturedValInLambdaInitializedOutside.kt")
+        public void testCapturedValInLambdaInitializedOutside() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedValInLambdaInitializedOutside.kt");
+        }
+
         @TestMetadata("capturedVarsOfSize2.kt")
         public void testCapturedVarsOfSize2() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/capturedVarsOptimization/capturedVarsOfSize2.kt");