Pseudocode: Generate additional returns for implicit return values (lambdas, functions with expression bodies)

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilder.java

@@ -74,7 +74,7 @@ public interface JetControlFlowBuilder {
     void nondeterministicJump(@NotNull List<Label> label, @NotNull JetElement element);
     void jumpToError(@NotNull JetElement element);
 
-    void returnValue(@NotNull JetReturnExpression returnExpression, @NotNull PseudoValue returnValue, @NotNull JetElement subroutine);
+    void returnValue(@NotNull JetExpression returnExpression, @NotNull PseudoValue returnValue, @NotNull JetElement subroutine);
     void returnNoValue(@NotNull JetReturnExpression returnExpression, @NotNull JetElement subroutine);
 
     void throwException(@NotNull JetThrowExpression throwExpression, @NotNull PseudoValue thrownValue);

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilderAdapter.java

@@ -233,7 +233,7 @@ public abstract class JetControlFlowBuilderAdapter implements JetControlFlowBuil
     }
 
     @Override
-    public void returnValue(@NotNull JetReturnExpression returnExpression, @NotNull PseudoValue returnValue, @NotNull JetElement subroutine) {
+    public void returnValue(@NotNull JetExpression returnExpression, @NotNull PseudoValue returnValue, @NotNull JetElement subroutine) {
         getDelegateBuilder().returnValue(returnExpression, returnValue, subroutine);
     }
 

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowProcessor.java

@@ -87,6 +87,9 @@ public class JetControlFlowProcessor {
             JetExpression bodyExpression = declarationWithBody.getBodyExpression();
             if (bodyExpression != null) {
                 cfpVisitor.generateInstructions(bodyExpression);
+                if (!declarationWithBody.hasBlockBody()) {
+                    generateImplicitReturnValue(bodyExpression, subroutine);
+                }
             }
         } else {
             cfpVisitor.generateInstructions(subroutine);
@@ -94,6 +97,13 @@ public class JetControlFlowProcessor {
         return builder.exitSubroutine(subroutine);
     }
 
+    private void generateImplicitReturnValue(@NotNull JetExpression bodyExpression, @NotNull JetElement subroutine) {
+        PseudoValue returnValue = builder.getBoundValue(bodyExpression);
+        if (returnValue == null) return;
+
+        builder.returnValue(bodyExpression, returnValue, subroutine);
+    }
+
     private void processLocalDeclaration(@NotNull JetDeclaration subroutine) {
         JetElement parent = PsiTreeUtil.getParentOfType(subroutine, JetElement.class);
         assert parent != null;

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/JetControlFlowInstructionsGenerator.java

@@ -284,9 +284,7 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
         }
 
         @Override
-        public void returnValue(
+        public void returnValue(@NotNull JetExpression returnExpression, @NotNull PseudoValue returnValue, @NotNull JetElement subroutine) {
-                @NotNull JetReturnExpression returnExpression, @NotNull PseudoValue returnValue, @NotNull JetElement subroutine
-        ) {
             Label exitPoint = getExitPoint(subroutine);
             handleJumpInsideTryFinally(exitPoint);
             add(new ReturnValueInstruction(returnExpression, getCurrentScope(), exitPoint, returnValue));

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/instructions/jumps/ReturnValueInstruction.kt

@@ -23,6 +23,7 @@ import java.util.Collections
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.LexicalScope
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.InstructionVisitor
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.InstructionVisitorWithResult
+import org.jetbrains.jet.lang.psi.JetReturnExpression
 
 public class ReturnValueInstruction(
         returnExpression: JetExpression,
@@ -47,4 +48,8 @@ public class ReturnValueInstruction(
     override fun createCopy(newLabel: Label, lexicalScope: LexicalScope): AbstractJumpInstruction {
         return ReturnValueInstruction((element as JetExpression), lexicalScope, newLabel, returnedValue)
     }
+
+    public val resultExpression: JetExpression =
+            element.let{ if (it is JetReturnExpression) it.getReturnedExpression()!! else element as JetExpression }
+    public val returnExpressionIfAny: JetReturnExpression? = element as? JetReturnExpression
 }

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/pseudocodeUtil.kt

@@ -32,37 +32,8 @@ import org.jetbrains.jet.lang.resolve.OverridingUtil
 import org.jetbrains.jet.lang.types.TypeUtils
 import org.jetbrains.jet.lang.types.JetType
 
-fun JetExpression.isStatement(pseudocode: Pseudocode): Boolean {
+fun JetExpression.isStatement(pseudocode: Pseudocode): Boolean =
-    val value = pseudocode.getElementValue(this);
+        pseudocode.getUsages(pseudocode.getElementValue(this)).isEmpty()
-    if (value == null) return true
-
-    fun considerUsedIfCreatedBeforeExit(): Boolean {
-        return when {
-            (getParent() as? JetFunction)?.getBodyExpression() == this ->
-                true
-            value.implicitReturnValue ->
-                true
-            else ->
-                false
-        }
-    }
-
-    val instruction = value.createdAt
-    if (considerUsedIfCreatedBeforeExit() && instruction.nextInstructions.any { it == pseudocode.getExitInstruction() }) return false
-    return pseudocode.getUsages(pseudocode.getElementValue(this)).isEmpty()
-}
-
-val PseudoValue.implicitReturnValue: Boolean
-    get() {
-        val pseudocode = createdAt.owner
-
-        val function = pseudocode.getCorrespondingElement() as? JetDeclarationWithBody
-
-        if (function is JetFunctionLiteral || (function != null && !function.hasBlockBody())) {
-            return pseudocode.getElementValue(function.getBodyExpression()) == this
-        }
-        return false
-    }
 
 fun getReceiverTypePredicate(resolvedCall: ResolvedCall<*>, receiverValue: ReceiverValue): TypePredicate? {
     val callableDescriptor = resolvedCall.getResultingDescriptor()
@@ -91,17 +62,15 @@ fun getExpectedTypePredicate(value: PseudoValue, bindingContext: BindingContext)
     fun addSubtypesOf(jetType: JetType?) = typePredicates.add(jetType?.getSubtypesPredicate())
 
     fun addTypePredicates(value: PseudoValue) {
-        if (value.implicitReturnValue) {
-            val function = value.createdAt.owner.getCorrespondingElement() as? JetDeclarationWithBody
-            val functionDescriptor = bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, function] as? FunctionDescriptor
-            addSubtypesOf(functionDescriptor?.getReturnType())
-        }
-
         pseudocode.getUsages(value).forEach {
             when (it) {
                 is ReturnValueInstruction -> {
-                    val functionDescriptor = (it.element as JetReturnExpression).getTargetFunctionDescriptor(bindingContext)
+                    val returnElement = it.element
-                    addSubtypesOf(functionDescriptor?.getReturnType())
+                    val functionDescriptor = when(returnElement) {
+                        is JetReturnExpression -> returnElement.getTargetFunctionDescriptor(bindingContext)
+                        else -> bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, value.createdAt.owner.getCorrespondingElement()]
+                    }
+                    addSubtypesOf((functionDescriptor as? CallableDescriptor)?.getReturnType())
                 }
 
                 is ConditionalJumpInstruction ->

compiler/testData/cfg-variables/basic/IfWithUninitialized.instructions

@@ -45,6 +45,7 @@ L0:
     magic(vararg a: Any?) -> <v0>    INIT: in: {a=D} out: {a=D}
     w(a|<v0>)                        INIT: in: {a=D} out: {a=ID}     USE: in: {a=READ} out: {a=READ}
     r(a) -> <v1>                     INIT: in: {a=ID} out: {a=ID}    USE: in: {} out: {a=READ}
+    ret(*|<v1>) L1
 L1:
     <END>
 error:

compiler/testData/cfg-variables/basic/UsageInFunctionLiteral.instructions

@@ -48,8 +48,9 @@ L3:
     r(a) -> <v4>                            INIT: in: {x=ID, y=ID} out: {x=ID, y=ID}    USE: in: {} out: {a=READ}
     mark(use(a))
     call(use(a), use|<v4>) -> <v5>
+  3 ret(*|<v5>) L4                          INIT: in: {x=ID} out: {x=ID}
 L4:
-  3 <END>                                   INIT: in: {x=ID} out: {x=ID}
+    <END>
 error:
     <ERROR>                                 INIT: in: {} out: {}
 sink:
@@ -64,6 +65,7 @@ L0:
     magic(vararg a: Any?) -> <v0>    INIT: in: {a=D} out: {a=D}
     w(a|<v0>)                        INIT: in: {a=D} out: {a=ID}     USE: in: {a=READ} out: {a=READ}
     r(a) -> <v1>                     INIT: in: {a=ID} out: {a=ID}    USE: in: {} out: {a=READ}
+    ret(*|<v1>) L1
 L1:
     <END>
 error:

compiler/testData/cfg-variables/basic/VariablesUsage.instructions

@@ -55,6 +55,7 @@ L0:
     magic(a: Int) -> <v0>    INIT: in: {a=D} out: {a=D}
     w(a|<v0>)                INIT: in: {a=D} out: {a=ID}     USE: in: {a=READ} out: {a=READ}
     r(a) -> <v1>             INIT: in: {a=ID} out: {a=ID}    USE: in: {} out: {a=READ}
+    ret(*|<v1>) L1
 L1:
     <END>
 error:

compiler/testData/cfg-variables/bugs/referenceToPropertyInitializer.instructions

@@ -43,8 +43,9 @@ L3:
     r(x) -> <v7>                                                                     USE: in: {} out: {x=READ}
     mark(sum(x - 1) + x)
     call(sum(x - 1) + x, plus|<v6>, <v7>) -> <v8>
+  2 ret(*|<v8>) L4
 L4:
-  2 <END>
+    <END>
 error:
     <ERROR>                                          INIT: in: {} out: {}
 sink:

compiler/testData/cfg-variables/bugs/varInitializationInIfInCycle.instructions

@@ -68,6 +68,7 @@ L0:
     magic(vararg a: Any?) -> <v0>    INIT: in: {a=D} out: {a=D}
     w(a|<v0>)                        INIT: in: {a=D} out: {a=ID}     USE: in: {a=READ} out: {a=READ}
     r(a) -> <v1>                     INIT: in: {a=ID} out: {a=ID}    USE: in: {} out: {a=READ}
+    ret(*|<v1>) L1
 L1:
     <END>
 error:

compiler/testData/cfg-variables/lexicalScopes/localFunctionScopeWithoutBody.instructions

@@ -38,6 +38,7 @@ L3:
     r(b) -> <v2>                                                            USE: in: {} out: {b=READ}
     mark(x + b)
     call(x + b, plus|<v1>, <v2>) -> <v3>
+    ret(*|<v3>) L4
 L4:
     <END>
 error:

compiler/testData/cfg/basic/Basic.instructions

@@ -79,12 +79,13 @@ L3:
   3 <START>
   4 mark(1)
     r(1) -> <v0>
+  3 ret(*|<v0>) L4
 L4:
-  3 <END>           NEXT:[<SINK>]
+    <END>             NEXT:[<SINK>]
 error:
-    <ERROR>         PREV:[]
+    <ERROR>           PREV:[]
 sink:
-    <SINK>          PREV:[<ERROR>, <END>]
+    <SINK>            PREV:[<ERROR>, <END>]
 =====================
 == foo ==
 fun foo(a : Boolean, b : Int) : Unit {}

compiler/testData/cfg/basic/ShortFunction.instructions

@@ -4,10 +4,11 @@ fun short() = 1
 L0:
   1 <START>
     r(1) -> <v0>
+    ret(*|<v0>) L1
 L1:
-    <END>           NEXT:[<SINK>]
+    <END>             NEXT:[<SINK>]
 error:
-    <ERROR>         PREV:[]
+    <ERROR>           PREV:[]
 sink:
-    <SINK>          PREV:[<ERROR>, <END>]
+    <SINK>            PREV:[<ERROR>, <END>]
 =====================

compiler/testData/cfg/declarations/local/ObjectExpression.instructions

@@ -45,12 +45,13 @@ override fun foo() = 10
 L0:
   1 <START>
     r(10) -> <v0>
+    ret(*|<v0>) L1
 L1:
-    <END>            NEXT:[<SINK>]
+    <END>             NEXT:[<SINK>]
 error:
-    <ERROR>          PREV:[]
+    <ERROR>           PREV:[]
 sink:
-    <SINK>           PREV:[<ERROR>, <END>]
+    <SINK>            PREV:[<ERROR>, <END>]
 =====================
 == foo ==
 fun foo(b: B) : Int {

compiler/testData/cfg/declarations/multiDeclaration/MultiDecl.instructions

@@ -19,12 +19,13 @@ fun component1() = 1
 L0:
   1 <START>
     r(1) -> <v0>
+    ret(*|<v0>) L1
 L1:
-    <END>           NEXT:[<SINK>]
+    <END>             NEXT:[<SINK>]
 error:
-    <ERROR>         PREV:[]
+    <ERROR>           PREV:[]
 sink:
-    <SINK>          PREV:[<ERROR>, <END>]
+    <SINK>            PREV:[<ERROR>, <END>]
 =====================
 == component2 ==
 fun component2() = 2
@@ -32,12 +33,13 @@ fun component2() = 2
 L0:
   1 <START>
     r(2) -> <v0>
+    ret(*|<v0>) L1
 L1:
-    <END>           NEXT:[<SINK>]
+    <END>             NEXT:[<SINK>]
 error:
-    <ERROR>         PREV:[]
+    <ERROR>           PREV:[]
 sink:
-    <SINK>          PREV:[<ERROR>, <END>]
+    <SINK>            PREV:[<ERROR>, <END>]
 =====================
 == test ==
 fun test(c: C) {

compiler/testData/cfg/declarations/properties/DelegatedProperty.instructions

@@ -24,6 +24,7 @@ L0:
     magic(p: PropertyMetadata) -> <v1>
     w(p|<v1>)
     r(0) -> <v2>
+    ret(*|<v2>) L1
 L1:
     <END>                                 NEXT:[<SINK>]
 error:

compiler/testData/cfg/declarations/properties/backingFieldQualifiedWithThis.instructions

@@ -26,6 +26,7 @@ L0:
     r($bar|<v1>) -> <v2>
     mark("foo" + this.$bar)
     call("foo" + this.$bar, plus|<v0>, <v2>) -> <v3>
+    ret(*|<v3>) L1
 L1:
     <END>                                               NEXT:[<SINK>]
 error:

compiler/testData/cfg/functions/DefaultValuesForArguments.instructions

@@ -17,6 +17,7 @@ L2 [after default value for parameter i]:
     r(j) -> <v5>
     mark(i + j)
     call(i + j, plus|<v4>, <v5>) -> <v6>
+    ret(*|<v6>) L1
 L1:
     <END>                                             NEXT:[<SINK>]
 error:

compiler/tests/org/jetbrains/jet/cfg/PseudoValueTestGenerated.java

@@ -16,14 +16,17 @@
 
 package org.jetbrains.jet.cfg;
 
+import junit.framework.Assert;
 import junit.framework.Test;
 import junit.framework.TestSuite;
+
+import java.io.File;
+import java.util.regex.Pattern;
 import org.jetbrains.jet.JetTestUtils;
 import org.jetbrains.jet.test.InnerTestClasses;
 import org.jetbrains.jet.test.TestMetadata;
 
-import java.io.File;
+import org.jetbrains.jet.cfg.AbstractPseudoValueTest;
-import java.util.regex.Pattern;
 
 /** This class is generated by {@link org.jetbrains.jet.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
 @SuppressWarnings("all")

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractFunctionUtils.kt

@@ -90,7 +90,7 @@ private fun List<Instruction>.getResultType(
     fun instructionToType(instruction: Instruction): JetType? {
         val expression = when (instruction) {
             is ReturnValueInstruction ->
-                (instruction.element as JetReturnExpression).getReturnedExpression()
+                instruction.resultExpression
             is InstructionWithValue ->
                 instruction.outputValue?.element as? JetExpression
             else -> null
@@ -145,10 +145,15 @@ private fun List<Instruction>.analyzeControlFlow(
                 }
 
         when (insn) {
-            is ReturnValueInstruction ->
+            is ReturnValueInstruction -> {
-                    if (isCurrentFunctionReturn(insn.element as JetReturnExpression)) {
+                val returnExpression = insn.returnExpressionIfAny
-                        valuedReturnExits.add(insn)
+                if (returnExpression == null) {
-                    }
+                    defaultExits.add(insn)
+                }
+                else if (isCurrentFunctionReturn(returnExpression)) {
+                    valuedReturnExits.add(insn)
+                }
+            }
 
             is AbstractJumpInstruction -> {
                 val element = insn.element