Normalize local returns before other inlining transformations

Local returns normalization can generate POP instructions.
These POP instructions can drop functional parameters, as in KT-17590,
and should be processed in markPlacesForInlineAndRemoveInlinable just
as other POP instructions.

KT-17590 conditional return in inline function parameter argument causes compilation exception

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/MethodInliner.java

@@ -421,8 +421,9 @@ public class MethodInliner {
     ) {
         node = prepareNode(node, finallyDeepShift);
 
-        Frame<SourceValue>[] sources = analyzeMethodNodeBeforeInline(node);
+        normalizeLocalReturns(node, labelOwner);
-        LocalReturnsNormalizer localReturnsNormalizer = LocalReturnsNormalizer.createFor(node, labelOwner, sources);
+
+        Frame<SourceValue>[] sources = analyzeMethodNodeWithoutMandatoryTransformations(node);
 
         Set<AbstractInsnNode> toDelete = SmartSet.create();
         InsnList instructions = node.instructions;
@@ -551,11 +552,39 @@ public class MethodInliner {
             }
         }
 
-        localReturnsNormalizer.transform(node);
-
         return node;
     }
 
+    private void normalizeLocalReturns(@NotNull MethodNode node, @NotNull LabelOwner labelOwner) {
+        Frame<SourceValue>[] frames = analyzeMethodNodeBeforeInline(node);
+
+        LocalReturnsNormalizer localReturnsNormalizer = new LocalReturnsNormalizer();
+
+        AbstractInsnNode[] instructions = node.instructions.toArray();
+
+        for (int i = 0; i < instructions.length; ++i) {
+            Frame<SourceValue> frame = frames[i];
+            // Don't care about dead code, it will be eliminated
+            if (frame == null) continue;
+
+            AbstractInsnNode insnNode = instructions[i];
+            if (!InlineCodegenUtil.isReturnOpcode(insnNode.getOpcode())) continue;
+
+            AbstractInsnNode insertBeforeInsn = insnNode;
+
+            // TODO extract isLocalReturn / isNonLocalReturn, see processReturns
+            String labelName = getMarkedReturnLabelOrNull(insnNode);
+            if (labelName != null) {
+                if (!labelOwner.isMyLabel(labelName)) continue;
+                insertBeforeInsn = insnNode.getPrevious();
+            }
+
+            localReturnsNormalizer.addLocalReturnToTransform(insnNode, insertBeforeInsn, frame);
+        }
+
+        localReturnsNormalizer.transform(node);
+    }
+
     private boolean isAnonymousClassThatMustBeRegenerated(@Nullable Type type) {
         if (type == null || type.getSort() != Type.OBJECT) return false;
         AnonymousObjectTransformationInfo info = inliningContext.findAnonymousObjectTransformationInfo(type.getInternalName());
@@ -571,6 +600,10 @@ public class MethodInliner {
             throw wrapException(e, node, "couldn't inline method call");
         }
 
+        return analyzeMethodNodeWithoutMandatoryTransformations(node);
+    }
+
+    private static Frame<SourceValue>[] analyzeMethodNodeWithoutMandatoryTransformations(@NotNull MethodNode node) {
         Analyzer<SourceValue> analyzer = new Analyzer<SourceValue>(new SourceInterpreter()) {
             @NotNull
             @Override
@@ -590,7 +623,7 @@ public class MethodInliner {
             return analyzer.analyze("fake", node);
         }
         catch (AnalyzerException e) {
-            throw wrapException(e, node, "couldn't inline method call");
+            throw new RuntimeException(e);
         }
     }
 
@@ -872,7 +905,7 @@ public class MethodInliner {
 
         private final List<LocalReturn> localReturns = new SmartList<LocalReturn>();
 
-        private boolean needsReturnVariable = false;
+        private int returnVariableSize = 0;
         private int returnOpcode = -1;
 
         private void addLocalReturnToTransform(
@@ -887,55 +920,27 @@ public class MethodInliner {
 
             localReturns.add(new LocalReturn(returnInsn, insertBeforeInsn, sourceValueFrame));
 
-            if (returnInsn.getOpcode() != Opcodes.RETURN && sourceValueFrame.getStackSize() > 1) {
+            if (returnInsn.getOpcode() != Opcodes.RETURN) {
-                needsReturnVariable = true;
+                if (returnInsn.getOpcode() == Opcodes.LRETURN || returnInsn.getOpcode() == Opcodes.DRETURN) {
+                    returnVariableSize = 2;
+                }
+                else {
+                    returnVariableSize = 1;
+                }
             }
         }
 
         public void transform(@NotNull MethodNode methodNode) {
             int returnVariableIndex = -1;
-            if (needsReturnVariable) {
+            if (returnVariableSize > 0) {
                 returnVariableIndex = methodNode.maxLocals;
-                methodNode.maxLocals++;
+                methodNode.maxLocals += returnVariableSize;
             }
 
             for (LocalReturn localReturn : localReturns) {
                 localReturn.transform(methodNode.instructions, returnVariableIndex);
             }
         }
-
-        @NotNull
-        public static LocalReturnsNormalizer createFor(
-                @NotNull MethodNode methodNode,
-                @NotNull LabelOwner owner,
-                @NotNull Frame<SourceValue>[] frames
-        ) {
-            LocalReturnsNormalizer result = new LocalReturnsNormalizer();
-
-            AbstractInsnNode[] instructions = methodNode.instructions.toArray();
-
-            for (int i = 0; i < instructions.length; ++i) {
-                Frame<SourceValue> frame = frames[i];
-                // Don't care about dead code, it will be eliminated
-                if (frame == null) continue;
-
-                AbstractInsnNode insnNode = instructions[i];
-                if (!InlineCodegenUtil.isReturnOpcode(insnNode.getOpcode())) continue;
-
-                AbstractInsnNode insertBeforeInsn = insnNode;
-
-                // TODO extract isLocalReturn / isNonLocalReturn, see processReturns
-                String labelName = getMarkedReturnLabelOrNull(insnNode);
-                if (labelName != null) {
-                    if (!owner.isMyLabel(labelName)) continue;
-                    insertBeforeInsn = insnNode.getPrevious();
-                }
-
-                result.addLocalReturnToTransform(insnNode, insertBeforeInsn, frame);
-            }
-
-            return result;
-        }
     }
 
     @NotNull

compiler/testData/codegen/box/controlStructures/kt17590.kt

@@ -0,0 +1,7 @@
+fun zap(s: String): String? = s
+
+inline fun tryZap(s: String, fn: (String) -> String): String {
+    return fn(zap(s) ?: return "null")
+}
+
+fun box() = tryZap("OK") { it }

compiler/testData/codegen/box/controlStructures/kt17590_long.kt

@@ -0,0 +1,10 @@
+fun foo(x: Any?, y: Any?) = 0L
+
+inline fun test(value: Any?): Long {
+    return foo(null, value ?: return 1L)
+}
+
+fun box(): String {
+    val t = test(null)
+    return if (t == 1L) "OK" else "fail: t=$t"
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -4364,6 +4364,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("kt17590.kt")
+        public void testKt17590() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("kt17590_long.kt")
+        public void testKt17590_long() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590_long.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("kt1899.kt")
         public void testKt1899() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt1899.kt");

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

@@ -4364,6 +4364,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("kt17590.kt")
+        public void testKt17590() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("kt17590_long.kt")
+        public void testKt17590_long() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590_long.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("kt1899.kt")
         public void testKt1899() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt1899.kt");

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

@@ -4364,6 +4364,18 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("kt17590.kt")
+        public void testKt17590() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("kt17590_long.kt")
+        public void testKt17590_long() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590_long.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("kt1899.kt")
         public void testKt1899() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt1899.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -5037,6 +5037,18 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("kt17590.kt")
+        public void testKt17590() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("kt17590_long.kt")
+        public void testKt17590_long() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt17590_long.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("kt1899.kt")
         public void testKt1899() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/kt1899.kt");