[PL] Fix: Respect inlined array constructor lambdas

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/AdditionalIrUtils.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.ir.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrClassPublicSymbolImpl
@@ -23,6 +24,13 @@ import java.io.File
 
 val IrConstructor.constructedClass get() = this.parent as IrClass
 
+fun IrClassifierSymbol?.isArrayOrPrimitiveArray(builtins: IrBuiltIns): Boolean =
+    this == builtins.arrayClass || this in builtins.primitiveArraysToPrimitiveTypes
+
+// Constructors can't be marked as inline in metadata, hence this check.
+fun IrFunction.isInlineArrayConstructor(builtIns: IrBuiltIns): Boolean =
+    this is IrConstructor && valueParameters.size == 2 && constructedClass.symbol.isArrayOrPrimitiveArray(builtIns)
+
 val IrDeclarationParent.fqNameForIrSerialization: FqName
     get() = when (this) {
         is IrPackageFragment -> this.fqName

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartiallyLinkedIrTreePatcher.kt

@@ -895,11 +895,12 @@ internal class PartiallyLinkedIrTreePatcher(
             }
         ) { super.visitBlockBody(body) }
 
-        override fun visitCall(expression: IrCall) = withContext(
+        // Allows visiting any type of call: IrCall, IrConstructorCall, IrEnumConstructorCall, IrDelegatingConstructorCall.
+        override fun visitFunctionAccess(expression: IrFunctionAccessExpression) = withContext(
             { oldContext ->
                 val functionSymbol = expression.symbol
                 val function = if (functionSymbol.isBound) functionSymbol.owner else return@withContext oldContext
-                if (!function.isInline) return@withContext oldContext
+                if (!function.isInline && !function.isInlineArrayConstructor(builtIns)) return@withContext oldContext
 
                 fun IrValueParameter?.canHaveNonLocalReturns(): Boolean = this != null && !isCrossinline && !isNoinline
 
@@ -925,7 +926,7 @@ internal class PartiallyLinkedIrTreePatcher(
                     inlinedLambdaArgumentsWithPermittedNonLocalReturns = inlinedLambdaArgumentsWithPermittedNonLocalReturns.toSet()
                 )
             }
-        ) { super.visitCall(expression) }
+        ) { super.visitFunctionAccess(expression) }
 
         override fun visitReturn(expression: IrReturn) = withContext { context ->
             expression.maybeThrowLinkageError(transformer = this@NonLocalReturnsPatcher) {

compiler/testData/klibABI/functionTransformations/lib2/l2.kt

@@ -161,3 +161,21 @@ fun newNonInlineFunctionInOpenClassImpl(oci: OpenClassImpl, x: Int): String = oc
 
 fun inlineLambdaToNoinlineLambda(x: Int): String = Functions.inlineLambdaToNoinlineLambda(x) { if (it > 0) it.toString() else return "inlineLambdaToNoinlineLambda($x)" }
 fun inlineLambdaToCrossinlineLambda(x: Int): String = Functions.inlineLambdaToCrossinlineLambda(x) { if (it > 0) it.toString() else return "inlineLambdaToCrossinlineLambda($x)" }
+
+fun nonLocalReturnFromArrayConstructorLambda(expected: String, unexpected: String): String = Array(1) outer@{
+    Array(1) {
+        if ('1' in "123") { // The condition that is always true.
+            return@outer expected
+        }
+        unexpected
+    }[0]
+}[0]
+
+fun nonLocalReturnFromIntArrayConstructorLambda(expected: Int, unexpected: Int): Int = IntArray(1) outer@{
+    IntArray(1) {
+        if ('1' in "123") { // The condition that is always true.
+            return@outer expected
+        }
+        unexpected
+    }[0]
+}[0]

compiler/testData/klibABI/functionTransformations/main/m.kt

@@ -119,4 +119,7 @@ fun box() = abiTest {
     expectFailure(linkage("Illegal non-local return: The return target is function 'inlineLambdaToNoinlineLambda' while only the following return targets are allowed: lambda in function 'inlineLambdaToNoinlineLambda'")) { inlineLambdaToNoinlineLambda(-3) }
     expectSuccess("Functions.inlineLambdaToCrossinlineLambda(5) { 10 }") { inlineLambdaToCrossinlineLambda(5) }
     expectFailure(linkage("Illegal non-local return: The return target is function 'inlineLambdaToCrossinlineLambda' while only the following return targets are allowed: lambda in function 'inlineLambdaToCrossinlineLambda'")) { inlineLambdaToCrossinlineLambda(-5) }
+
+    expectSuccess("success") { nonLocalReturnFromArrayConstructorLambda("success", "failure") }
+    expectSuccess(100) { nonLocalReturnFromIntArrayConstructorLambda(100, -100) }
 }