Support bound function reference inlining

compiler/backend/src/org/jetbrains/kotlin/codegen/FunctionReferenceGenerationStrategy.java

@@ -185,8 +185,10 @@ public class FunctionReferenceGenerationStrategy extends FunctionGenerationStrat
         if (receiverValue != null) return receiverValue;
 
         if (receiverType != null) {
+            ClassDescriptor classDescriptor = (ClassDescriptor) codegen.getContext().getParentContext().getContextDescriptor();
+            Type asmType = codegen.getState().getTypeMapper().mapClass(classDescriptor);
             return StackValue.field(
-                    receiverType, Type.getObjectType(codegen.getParentCodegen().getClassName()), AsmUtil.CAPTURED_RECEIVER_FIELD,
+                    receiverType, asmType, AsmUtil.CAPTURED_RECEIVER_FIELD,
                     /* isStatic = */ false, StackValue.LOCAL_0
             );
         }

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

@@ -25,6 +25,9 @@ public class CapturedParamInfo extends ParameterInfo {
     private final String newFieldName;
     private final boolean skipInConstructor;
 
+    //Now used only for bound function reference receiver
+    private boolean synthetic;
+
     public CapturedParamInfo(@NotNull CapturedParamDesc desc, @NotNull String newFieldName, boolean skipped, int index, int remapIndex) {
         super(desc.getType(), skipped, index, remapIndex, -1);
         this.desc = desc;
@@ -63,6 +66,7 @@ public class CapturedParamInfo extends ParameterInfo {
                 desc, newFieldName, isSkipped, getIndex(), getRemapValue(), skipInConstructor, newDeclarationIndex
         );
         result.setLambda(getLambda());
+        result.setSynthetic(synthetic);
         return result;
     }
 
@@ -74,4 +78,16 @@ public class CapturedParamInfo extends ParameterInfo {
     public boolean isSkipInConstructor() {
         return skipInConstructor;
     }
+
+    public boolean isSynthetic() {
+        return synthetic;
+    }
+
+    public void setSynthetic(boolean synthetic) {
+        this.synthetic = synthetic;
+    }
+
+    public static boolean isSynthetic(@NotNull ParameterInfo info) {
+        return info instanceof CapturedParamInfo && ((CapturedParamInfo) info).isSynthetic();
+    }
 }

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

@@ -544,17 +544,17 @@ public class InlineCodegen extends CallGenerator {
         if (expression instanceof KtCallableReferenceExpression) {
             KtCallableReferenceExpression callableReferenceExpression = (KtCallableReferenceExpression) expression;
             KtExpression receiverExpression = callableReferenceExpression.getReceiverExpression();
-            StackValue receiverValue =
+            Type receiverValue =
                     receiverExpression != null && codegen.getBindingContext().getType(receiverExpression) != null
-                    ? codegen.gen(receiverExpression)
+                    ? codegen.getState().getTypeMapper().mapType(codegen.getBindingContext().getType(receiverExpression))
                     : null;
 
             strategy = new FunctionReferenceGenerationStrategy(
                     state,
                     descriptor,
                     CallUtilKt.getResolvedCallWithAssert(callableReferenceExpression.getCallableReference(), codegen.getBindingContext()),
-                    receiverValue != null ? receiverValue.type : null,
+                    receiverValue,
-                    receiverValue
+                    null
             );
         }
         else {
@@ -640,16 +640,16 @@ public class InlineCodegen extends CallGenerator {
         if (!asFunctionInline && Type.VOID_TYPE != type) {
             //TODO remap only inlinable closure => otherwise we could get a lot of problem
             boolean couldBeRemapped = !shouldPutValue(type, stackValue);
-            StackValue remappedIndex = couldBeRemapped ? stackValue : null;
+            StackValue remappedValue = couldBeRemapped ? stackValue : null;
 
             ParameterInfo info;
             if (capturedParamIndex >= 0) {
                 CapturedParamDesc capturedParamInfoInLambda = activeLambda.getCapturedVars().get(capturedParamIndex);
                 info = invocationParamBuilder.addCapturedParam(capturedParamInfoInLambda, capturedParamInfoInLambda.getFieldName(), false);
-                info.setRemapValue(remappedIndex);
+                info.setRemapValue(remappedValue);
             }
             else {
-                info = invocationParamBuilder.addNextValueParameter(type, false, remappedIndex, parameterIndex);
+                info = invocationParamBuilder.addNextValueParameter(type, false, remappedValue, parameterIndex);
             }
 
             recordParameterValueInLocalVal(info);
@@ -707,7 +707,12 @@ public class InlineCodegen extends CallGenerator {
             ParameterInfo info = infos[i];
             if (!info.isSkippedOrRemapped()) {
                 Type type = info.type;
-                StackValue.local(index[i], type).store(StackValue.onStack(type), codegen.v);
+                StackValue.Local local = StackValue.local(index[i], type);
+                local.store(StackValue.onStack(type), codegen.v);
+                if (info instanceof CapturedParamInfo) {
+                    info.setRemapValue(local);
+                    ((CapturedParamInfo) info).setSynthetic(true);
+                }
             }
         }
     }
@@ -734,7 +739,7 @@ public class InlineCodegen extends CallGenerator {
         List<ParameterInfo> infos = invocationParamBuilder.listAllParams();
         for (ListIterator<? extends ParameterInfo> iterator = infos.listIterator(infos.size()); iterator.hasPrevious(); ) {
             ParameterInfo param = iterator.previous();
-            if (!param.isSkippedOrRemapped()) {
+            if (!param.isSkippedOrRemapped() || CapturedParamInfo.isSynthetic(param)) {
                 codegen.getFrameMap().leaveTemp(param.type);
             }
         }
@@ -749,13 +754,6 @@ public class InlineCodegen extends CallGenerator {
             // TODO: support inline of property references passed to inlinable function parameters
             SimpleFunctionDescriptor functionReference = state.getBindingContext().get(BindingContext.FUNCTION, deparenthesized);
             if (functionReference == null) return false;
-
-            KtExpression receiverExpression = ((KtCallableReferenceExpression) deparenthesized).getReceiverExpression();
-            if (receiverExpression != null) {
-                // TODO (!): support inline for bound function references
-                DoubleColonLHS lhs = state.getBindingContext().get(BindingContext.DOUBLE_COLON_LHS, receiverExpression);
-                if (lhs instanceof DoubleColonLHS.Expression) return false;
-            }
         }
 
         return InlineUtil.isInlineLambdaParameter(valueParameterDescriptor) &&
@@ -768,7 +766,7 @@ public class InlineCodegen extends CallGenerator {
                deparenthesized instanceof KtCallableReferenceExpression;
     }
 
-    private void rememberClosure(@NotNull KtExpression expression, @NotNull Type type, @NotNull ValueParameterDescriptor parameter) {
+    private LambdaInfo rememberClosure(@NotNull KtExpression expression, @NotNull Type type, @NotNull ValueParameterDescriptor parameter) {
         KtExpression lambda = KtPsiUtil.deparenthesize(expression);
         assert isInlinableParameterExpression(lambda) : "Couldn't find inline expression in " + expression.getText();
 
@@ -777,6 +775,8 @@ public class InlineCodegen extends CallGenerator {
         ParameterInfo closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.getIndex());
         closureInfo.setLambda(info);
         expressionMap.put(closureInfo.getIndex(), info);
+        info.setBoundCallableReference(getBoundCallableReferenceReceiver(expression) != null);
+        return info;
     }
 
     @NotNull
@@ -801,9 +801,15 @@ public class InlineCodegen extends CallGenerator {
 
     private void putClosureParametersOnStack() {
         for (LambdaInfo next : expressionMap.values()) {
-            activeLambda = next;
+            //closure parameters for bounded callable references are generated inplace
-            codegen.pushClosureOnStack(next.getClassDescriptor(), true, this, /* functionReferenceReceiver = */ null);
+            if (next.isBoundCallableReference()) continue;
+            putClosureParametersOnStack(next, null);
         }
+    }
+
+    private void putClosureParametersOnStack(@NotNull LambdaInfo next, @Nullable StackValue functionReferenceReceiver) {
+        activeLambda = next;
+        codegen.pushClosureOnStack(next.getClassDescriptor(), true, this, functionReferenceReceiver);
         activeLambda = null;
     }
 
@@ -846,7 +852,12 @@ public class InlineCodegen extends CallGenerator {
             int parameterIndex
     ) {
         if (isInliningParameter(argumentExpression, valueParameterDescriptor)) {
-            rememberClosure(argumentExpression, parameterType, valueParameterDescriptor);
+            LambdaInfo lambdaInfo = rememberClosure(argumentExpression, parameterType, valueParameterDescriptor);
+
+            KtExpression receiver = getBoundCallableReferenceReceiver(argumentExpression);
+            if (receiver != null) {
+                putClosureParametersOnStack(lambdaInfo, codegen.gen(receiver));
+            }
         }
         else {
             StackValue value = codegen.gen(argumentExpression);
@@ -854,6 +865,20 @@ public class InlineCodegen extends CallGenerator {
         }
     }
 
+    private KtExpression getBoundCallableReferenceReceiver(
+            @NotNull KtExpression argumentExpression
+    ) {
+        KtExpression deparenthesized = KtPsiUtil.deparenthesize(argumentExpression);
+        if (deparenthesized instanceof KtCallableReferenceExpression) {
+            KtExpression receiverExpression = ((KtCallableReferenceExpression) deparenthesized).getReceiverExpression();
+            if (receiverExpression != null) {
+                DoubleColonLHS lhs = state.getBindingContext().get(BindingContext.DOUBLE_COLON_LHS, receiverExpression);
+                if (lhs instanceof DoubleColonLHS.Expression) return receiverExpression;
+            }
+        }
+        return null;
+    }
+
     @Override
     public void putValueIfNeeded(@NotNull Type parameterType, @NotNull StackValue value) {
         putValueIfNeeded(parameterType, value, -1);

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

@@ -51,6 +51,7 @@ public class LambdaInfo implements LabelOwner {
 
     private SMAPAndMethodNode node;
     private List<CapturedParamDesc> capturedVars;
+    private boolean isBoundCallableReference;
 
     public LambdaInfo(@NotNull KtExpression expression, @NotNull KotlinTypeMapper typeMapper, boolean isCrossInline) {
         this.isCrossInline = isCrossInline;
@@ -165,4 +166,12 @@ public class LambdaInfo implements LabelOwner {
     public boolean isMyLabel(@NotNull String name) {
         return labels.contains(name);
     }
+
+    public boolean isBoundCallableReference() {
+        return isBoundCallableReference;
+    }
+
+    public void setBoundCallableReference(boolean boundCallableReference) {
+        isBoundCallableReference = boundCallableReference;
+    }
 }

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

@@ -48,6 +48,9 @@ public class LocalVarRemapper {
             }
             else {
                 remapValues[shift] = info.isRemapped() ? info.getRemapValue() : null;
+                if (CapturedParamInfo.isSynthetic(info)) {
+                    realSize += info.getType().getSize();
+                }
             }
         }
 

compiler/testData/codegen/boxInline/argumentOrder/boundFunctionReference.kt

@@ -0,0 +1,40 @@
+// FILE: 1.kt
+
+package test
+
+inline fun test(a: String, b: String, c: () -> String): String {
+    return a + b + c();
+}
+
+// FILE: 2.kt
+
+import test.*
+
+var res = ""
+
+fun String.id() = this
+
+val receiver: String
+    get() {
+        res += "L"
+        return "L"
+    }
+
+
+fun box(): String {
+    res = ""
+    var call = test(b = { res += "K"; "K" }(), a = { res += "O"; "O" }(), c = receiver::id)
+    if (res != "KOL" || call != "OKL") return "fail 1: $res != KOL or $call != OKL"
+
+    res = ""
+    call = test(b = { res += "K"; "K" }(), c = receiver::id, a = { res += "O"; "O" }())
+    if (res != "KLO" || call != "OKL") return "fail 2: $res != KLO or $call != OKL"
+
+
+    res = ""
+    call = test(c = receiver::id, b = { res += "K"; "K" }(), a = { res += "O"; "O" }())
+    if (res != "LKO" || call != "OKL") return "fail 3: $res != LKO or $call != OKL"
+
+    return "OK"
+}
+

compiler/testData/codegen/boxInline/argumentOrder/boundFunctionReference2.kt

@@ -0,0 +1,40 @@
+// FILE: 1.kt
+
+package test
+
+inline fun test(a: String, b: Long, c: () -> String): String {
+    return a + b + c();
+}
+
+// FILE: 2.kt
+
+import test.*
+
+var res = ""
+
+fun String.id() = this
+
+val receiver: String
+    get() {
+        res += "L"
+        return "L"
+    }
+
+
+fun box(): String {
+    res = ""
+    var call = test(b = { res += "K"; 1L }(), a = { res += "O"; "O" }(), c = receiver::id)
+    if (res != "KOL" || call != "O1L") return "fail 1: $res != KOL or $call != O1L"
+
+    res = ""
+    call = test(b = { res += "K"; 1L }(), c = receiver::id, a = { res += "O"; "O" }())
+    if (res != "KLO" || call != "O1L") return "fail 2: $res != KLO or $call != O1L"
+
+
+    res = ""
+    call = test(c = receiver::id, b = { res += "K"; 1L }(), a = { res += "O"; "O" }())
+    if (res != "LKO" || call != "O1L") return "fail 3: $res != LKO or $call != O1L"
+
+    return "OK"
+}
+

compiler/testData/codegen/boxInline/callableReference/bound/expression.kt

@@ -0,0 +1,17 @@
+// FILE: 1.kt
+
+package test
+
+inline fun go(f: () -> String) = f()
+
+fun String.id(): String = this
+
+fun foo(x: String, y: String): String {
+    return go((x + y)::id)
+}
+
+// FILE: 2.kt
+
+import test.*
+
+fun box() = foo("O", "K")

compiler/testData/codegen/boxInline/callableReference/bound/extensionReceiver.kt

@@ -0,0 +1,20 @@
+// FILE: 1.kt
+
+package test
+
+inline fun foo(x: () -> String, z: String) = x() + z
+
+fun String.id() = this
+
+// FILE: 2.kt
+
+import test.*
+
+fun String.test() : String {
+    return foo(this::id, "K")
+}
+
+
+fun box() : String {
+    return "O".test()
+}

compiler/testData/codegen/boxInline/callableReference/bound/filter.kt

@@ -0,0 +1,21 @@
+// WITH_RUNTIME
+// FILE: 1.kt
+
+package test
+
+inline fun stub(f: () -> String): String = f()
+
+// FILE: 2.kt
+
+import test.*
+
+class A(val z: String) {
+    fun filter(s: String) = z == s
+}
+
+
+fun box(): String {
+    val a = A("OK")
+    val s = arrayOf("OK")
+    return s.filter(a::filter).first()
+}

compiler/testData/codegen/boxInline/callableReference/bound/intrinsic.kt

@@ -0,0 +1,17 @@
+// FILE: 1.kt
+
+package test
+
+inline fun foo(x: (String) -> String, z: String) = x(z)
+
+fun String.id() = this
+
+// FILE: 2.kt
+
+import test.*
+
+fun box() : String {
+    var zeroSlot = "fail";
+    val z = "O"
+    return foo(z::plus, "K")
+}

compiler/testData/codegen/boxInline/callableReference/bound/map.kt

@@ -0,0 +1,21 @@
+// WITH_RUNTIME
+// FILE: 1.kt
+
+package test
+
+inline fun stub(f: () -> String): String = f()
+
+// FILE: 2.kt
+
+import test.*
+
+class A(val z: String) {
+    fun map(s: String) = z + s
+}
+
+
+fun box(): String {
+    val a = A("O")
+    val s = arrayOf("K")
+    return s.map(a::map).first()
+}

compiler/testData/codegen/boxInline/callableReference/bound/simple.kt

@@ -0,0 +1,17 @@
+// FILE: 1.kt
+
+package test
+
+inline fun foo(x: () -> String, z: String) = x() + z
+
+fun String.id() = this
+
+// FILE: 2.kt
+
+import test.*
+
+fun box() : String {
+    var zeroSlot = "fail";
+    val z = "O"
+    return foo(z::id, "K")
+}

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

@@ -406,6 +406,18 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/argumentOrder"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("boundFunctionReference.kt")
+        public void testBoundFunctionReference() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/argumentOrder/boundFunctionReference.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("boundFunctionReference2.kt")
+        public void testBoundFunctionReference2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/argumentOrder/boundFunctionReference2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("captured.kt")
         public void testCaptured() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/argumentOrder/captured.kt");
@@ -564,6 +576,51 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/topLevelExtension.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("compiler/testData/codegen/boxInline/callableReference/bound")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class Bound extends AbstractBlackBoxInlineCodegenTest {
+            public void testAllFilesPresentInBound() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/callableReference/bound"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+
+            @TestMetadata("expression.kt")
+            public void testExpression() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/expression.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("extensionReceiver.kt")
+            public void testExtensionReceiver() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/extensionReceiver.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("filter.kt")
+            public void testFilter() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/filter.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("intrinsic.kt")
+            public void testIntrinsic() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/intrinsic.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("map.kt")
+            public void testMap() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/map.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("simple.kt")
+            public void testSimple() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/simple.kt");
+                doTest(fileName);
+            }
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/boxInline/capture")

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

@@ -406,6 +406,18 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/argumentOrder"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("boundFunctionReference.kt")
+        public void testBoundFunctionReference() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/argumentOrder/boundFunctionReference.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("boundFunctionReference2.kt")
+        public void testBoundFunctionReference2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/argumentOrder/boundFunctionReference2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("captured.kt")
         public void testCaptured() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/argumentOrder/captured.kt");
@@ -564,6 +576,51 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/topLevelExtension.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("compiler/testData/codegen/boxInline/callableReference/bound")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class Bound extends AbstractCompileKotlinAgainstInlineKotlinTest {
+            public void testAllFilesPresentInBound() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/callableReference/bound"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+
+            @TestMetadata("expression.kt")
+            public void testExpression() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/expression.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("extensionReceiver.kt")
+            public void testExtensionReceiver() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/extensionReceiver.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("filter.kt")
+            public void testFilter() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/filter.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("intrinsic.kt")
+            public void testIntrinsic() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/intrinsic.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("map.kt")
+            public void testMap() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/map.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("simple.kt")
+            public void testSimple() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/callableReference/bound/simple.kt");
+                doTest(fileName);
+            }
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/boxInline/capture")