Transform variable-as-function calls for extension functions

In the following code example
  fun test(f: Any.() -> Unit) = 42.f()
front-end resolves variable-as-function call for 'f' as 'invoke'
with signature 'Function1<Any, Unit>#Any.() -> Unit'.
However, Function1<Any, Unit> has a single 'invoke' method
with signature 'Function1<Any, Unit>#(Any) -> Unit'.
This didn't cause any problems with loosely typed JVM and JS back-ends.
However, in IR with symbols this means a reference to non-existing
declaration.

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ArgumentsGenerationUtils.kt

@@ -16,12 +16,11 @@
 
 package org.jetbrains.kotlin.psi2ir.generators
 
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
+import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.SyntheticFieldDescriptor
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.expressions.IrDeclarationReference
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrExpressionWithCopy
@@ -38,6 +37,7 @@ import org.jetbrains.kotlin.resolve.calls.callUtil.isSafeCall
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.scopes.receivers.*
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeSubstitutor
 import java.lang.AssertionError
 
 fun StatementGenerator.generateReceiverOrNull(ktDefaultElement: KtElement, receiver: ReceiverValue?): IntermediateValue? =
@@ -46,7 +46,7 @@ fun StatementGenerator.generateReceiverOrNull(ktDefaultElement: KtElement, recei
 fun StatementGenerator.generateReceiver(ktDefaultElement: KtElement, receiver: ReceiverValue): IntermediateValue =
         generateReceiver(ktDefaultElement.startOffset, ktDefaultElement.endOffset, receiver)
 
-fun StatementGenerator.generateReceiver(startOffset: Int, endOffset: Int, receiver: ReceiverValue): IntermediateValue =
+fun StatementGenerator.generateReceiver(defaultStartOffset: Int, defaultEndOffset: Int, receiver: ReceiverValue): IntermediateValue =
         if (receiver is TransientReceiver)
             TransientReceiverValue(receiver.type)
         else generateDelegatedValue(receiver.type) {
@@ -54,9 +54,9 @@ fun StatementGenerator.generateReceiver(startOffset: Int, endOffset: Int, receiv
                 is ImplicitClassReceiver -> {
                     val receiverClassDescriptor = receiver.classDescriptor
                     if (shouldGenerateReceiverAsSingletonReference(receiverClassDescriptor))
-                        generateSingletonReference(receiverClassDescriptor, startOffset, endOffset, receiver.type)
+                        generateSingletonReference(receiverClassDescriptor, defaultStartOffset, defaultEndOffset, receiver.type)
                     else
-                        IrGetValueImpl(startOffset, endOffset,
+                        IrGetValueImpl(defaultStartOffset, defaultEndOffset,
                                        context.symbolTable.referenceValueParameter(receiverClassDescriptor.thisAsReceiverParameter))
                 }
                 is ThisClassReceiver ->
@@ -69,7 +69,7 @@ fun StatementGenerator.generateReceiver(startOffset: Int, endOffset: Int, receiv
                     IrGetObjectValueImpl(receiver.expression.startOffset, receiver.expression.endOffset, receiver.type,
                                          context.symbolTable.referenceClass(receiver.classQualifier.descriptor as ClassDescriptor))
                 is ExtensionReceiver ->
-                    IrGetValueImpl(startOffset, startOffset,
+                    IrGetValueImpl(defaultStartOffset, defaultStartOffset,
                                    context.symbolTable.referenceValueParameter(receiver.declarationDescriptor.extensionReceiverParameter!!))
                 else ->
                     TODO("Receiver: ${receiver::class.java.simpleName}")
@@ -228,11 +228,72 @@ fun Generator.getSuperQualifier(resolvedCall: ResolvedCall<*>): ClassDescriptor?
 }
 
 fun StatementGenerator.pregenerateCall(resolvedCall: ResolvedCall<*>): CallBuilder {
+    if (resolvedCall.isExtensionInvokeCall()) {
+        return pregenerateExtensionInvokeCall(resolvedCall)
+    }
+
     val call = pregenerateCallWithReceivers(resolvedCall)
     pregenerateValueArguments(call, resolvedCall)
     return call
 }
 
+fun StatementGenerator.pregenerateExtensionInvokeCall(resolvedCall: ResolvedCall<*>): CallBuilder {
+    val extensionInvoke = resolvedCall.resultingDescriptor
+    val functionNClass = extensionInvoke.containingDeclaration as? ClassDescriptor ?:
+                         throw AssertionError("'invoke' should be a class member: $extensionInvoke")
+    val unsubstitutedPlainInvokes = functionNClass.unsubstitutedMemberScope.getContributedFunctions(extensionInvoke.name, NoLookupLocation.FROM_BACKEND)
+    val unsubstitutedPlainInvoke = unsubstitutedPlainInvokes.singleOrNull() ?:
+                                   throw AssertionError("There should be a single 'invoke' in FunctionN class: $unsubstitutedPlainInvokes")
+
+    val expectedValueParametersCount = extensionInvoke.valueParameters.size + 1
+    assert(unsubstitutedPlainInvoke.valueParameters.size == expectedValueParametersCount) {
+        "Plain 'invoke' should have $expectedValueParametersCount value parameters, got ${unsubstitutedPlainInvoke.valueParameters}"
+    }
+
+    val functionNType = extensionInvoke.dispatchReceiverParameter!!.type
+    val plainInvoke = unsubstitutedPlainInvoke.substitute(TypeSubstitutor.create(functionNType)) ?:
+                      throw AssertionError("Substitution failed for $unsubstitutedPlainInvoke, type=$functionNType")
+
+    val ktCallElement = resolvedCall.call.callElement
+
+    val call = CallBuilder(resolvedCall, plainInvoke, isExtensionInvokeCall = true)
+
+    val functionReceiverValue = run {
+        val dispatchReceiver = resolvedCall.dispatchReceiver ?:
+                               throw AssertionError("Extension 'invoke' call should have a dispatch receiver")
+        generateReceiver(ktCallElement, dispatchReceiver)
+    }
+
+    val extensionInvokeReceiverValue = run {
+        val extensionReceiver = resolvedCall.extensionReceiver ?:
+                                throw AssertionError("Extension 'invoke' call should have an extension receiver")
+        generateReceiver(ktCallElement, extensionReceiver)
+    }
+
+    call.callReceiver =
+            if (resolvedCall.call.isSafeCall())
+                SafeExtensionInvokeCallReceiver(this, ktCallElement.startOffset, ktCallElement.endOffset,
+                                                call, functionReceiverValue, extensionInvokeReceiverValue)
+            else
+                ExtensionInvokeCallReceiver(call, functionReceiverValue, extensionInvokeReceiverValue)
+
+    call.irValueArgumentsByIndex[0] = null
+    resolvedCall.valueArgumentsByIndex!!.forEachIndexed { index, valueArgument ->
+        val valueParameter = call.descriptor.valueParameters[index]
+        call.irValueArgumentsByIndex[index + 1] = generateValueArgument(valueArgument, valueParameter)
+    }
+
+    return call
+}
+
+private fun ResolvedCall<*>.isExtensionInvokeCall(): Boolean {
+    val callee = resultingDescriptor as? SimpleFunctionDescriptor ?: return false
+    if (callee.name.asString() != "invoke") return false
+    val dispatchReceiverType = callee.dispatchReceiverParameter?.type ?: return false
+    if (!dispatchReceiverType.isBuiltinFunctionalType) return false
+    return extensionReceiver != null
+}
+
 fun getTypeArguments(resolvedCall: ResolvedCall<*>?): Map<TypeParameterDescriptor, KotlinType>? {
     if (resolvedCall == null) return null
 

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/intermediate/CallBuilder.kt

@@ -27,16 +27,19 @@ import org.jetbrains.kotlin.types.KotlinType
 
 class CallBuilder(
         val original: ResolvedCall<*>,
-        val descriptor: CallableDescriptor
+        val descriptor: CallableDescriptor,
+        val isExtensionInvokeCall: Boolean = false
 ) {
     var superQualifier: ClassDescriptor? = null
 
     lateinit var callReceiver: CallReceiver
 
+    private val parametersOffset = if (isExtensionInvokeCall) 1 else 0
+
     val irValueArgumentsByIndex = arrayOfNulls<IrExpression>(descriptor.valueParameters.size)
 
     fun getValueArgument(valueParameterDescriptor: ValueParameterDescriptor) =
-            irValueArgumentsByIndex[valueParameterDescriptor.index]
+            irValueArgumentsByIndex[valueParameterDescriptor.index + parametersOffset]
 }
 
 val CallBuilder.argumentsCount: Int get() =

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/intermediate/ExtensionInvokeCallReceiver.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.psi2ir.intermediate
+
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+
+class ExtensionInvokeCallReceiver(
+        val callBuilder: CallBuilder,
+        val functionReceiver: IntermediateValue,
+        val extensionInvokeReceiver: IntermediateValue
+) : CallReceiver {
+    override fun call(withDispatchAndExtensionReceivers: (IntermediateValue?, IntermediateValue?) -> IrExpression): IrExpression {
+        // extensionInvokeReceiver is actually a first argument:
+        //      receiver.extFun(p1, ..., pN)
+        //      =>
+        //      extFun.invoke(receiver, p1, ..., pN)
+
+        assert(callBuilder.irValueArgumentsByIndex[0] == null) {
+            "Extension 'invoke' call should have null as its 1st value argument, got: ${callBuilder.irValueArgumentsByIndex[0]}"
+        }
+        callBuilder.irValueArgumentsByIndex[0] = extensionInvokeReceiver.load()
+        return withDispatchAndExtensionReceivers(functionReceiver, null)
+    }
+}

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/intermediate/SafeExtensionInvokeCallReceiver.kt

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.psi2ir.intermediate
+
+import org.jetbrains.kotlin.ir.builders.constNull
+import org.jetbrains.kotlin.ir.builders.equalsNull
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrIfThenElseImpl
+import org.jetbrains.kotlin.psi2ir.generators.GeneratorWithScope
+import org.jetbrains.kotlin.types.typeUtil.makeNullable
+
+class SafeExtensionInvokeCallReceiver(
+        val generator: GeneratorWithScope,
+        val startOffset: Int,
+        val endOffset: Int,
+        val callBuilder: CallBuilder,
+        val functionReceiver: IntermediateValue,
+        val extensionInvokeReceiver: IntermediateValue
+) : CallReceiver {
+    override fun call(withDispatchAndExtensionReceivers: (IntermediateValue?, IntermediateValue?) -> IrExpression): IrExpression {
+        // extensionInvokeReceiver is actually a first argument:
+        //      receiver?.extFun(p1, ..., pN)
+        //      =>
+        //      val tmp = [receiver]
+        //      if (tmp == null) null else extFun.invoke(tmp, p1, ..., pN)
+
+        val irTmp = generator.scope.createTemporaryVariable(extensionInvokeReceiver.load(), "safe_receiver")
+
+        val safeReceiverValue = VariableLValue(irTmp)
+
+        // Patch call and generate it
+        assert(callBuilder.irValueArgumentsByIndex[0] == null) {
+            "Safe extension 'invoke' call should have null as its 1st value argument, got: ${callBuilder.irValueArgumentsByIndex[0]}"
+        }
+        callBuilder.irValueArgumentsByIndex[0] = safeReceiverValue.load()
+        val irResult = withDispatchAndExtensionReceivers(functionReceiver, null)
+
+        val resultType = irResult.type.makeNullable()
+
+        return IrBlockImpl(
+                startOffset, endOffset, resultType, IrStatementOrigin.SAFE_CALL,
+                arrayListOf(
+                        irTmp,
+                        IrIfThenElseImpl(
+                                startOffset, endOffset, resultType,
+                                generator.context.equalsNull(startOffset, endOffset, safeReceiverValue.load()),
+                                generator.context.constNull(startOffset, endOffset),
+                                irResult,
+                                IrStatementOrigin.SAFE_CALL
+                        )
+                )
+        )
+    }
+}
+

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/transformations/InsertImplicitCasts.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.psi2ir.transformations
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.isBuiltinExtensionFunctionalType
+import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrField

compiler/testData/ir/irText/expressions/extFunInvokeAsFun.txt

@@ -12,6 +12,6 @@ FILE /extFunInvokeAsFun.kt
     VALUE_PARAMETER value-parameter block: kotlin.Any?.() -> kotlin.Unit
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='with2(Any?, Any?.() -> Unit): Unit'
-        CALL 'invoke() on Any?: Unit' type=kotlin.Unit origin=INVOKE
+        CALL 'invoke(Any?): Unit' type=kotlin.Unit origin=INVOKE
           $this: GET_VAR 'value-parameter block: Any?.() -> Unit' type=kotlin.Any?.() -> kotlin.Unit origin=VARIABLE_AS_FUNCTION
-          $receiver: GET_VAR 'value-parameter receiver: Any?' type=kotlin.Any? origin=null
+          p1: GET_VAR 'value-parameter receiver: Any?' type=kotlin.Any? origin=null

compiler/testData/ir/irText/expressions/extFunSafeInvoke.kt

@@ -0,0 +1,2 @@
+fun test(receiver: Any?, fn: Any.(Int, String) -> Unit) =
+        receiver?.fn(42, "Hello")

compiler/testData/ir/irText/expressions/extFunSafeInvoke.txt

@@ -0,0 +1,22 @@
+FILE /extFunSafeInvoke.kt
+  FUN public fun test(receiver: kotlin.Any?, fn: kotlin.Any.(kotlin.Int, kotlin.String) -> kotlin.Unit): kotlin.Unit?
+    VALUE_PARAMETER value-parameter receiver: kotlin.Any?
+    VALUE_PARAMETER value-parameter fn: kotlin.Any.(kotlin.Int, kotlin.String) -> kotlin.Unit
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='test(Any?, Any.(Int, String) -> Unit): Unit?'
+        BLOCK type=kotlin.Unit? origin=SAFE_CALL
+          VAR IR_TEMPORARY_VARIABLE val tmp0_safe_receiver: kotlin.Any?
+            GET_VAR 'value-parameter receiver: Any?' type=kotlin.Any? origin=null
+          WHEN type=kotlin.Unit? origin=SAFE_CALL
+            BRANCH
+              if: CALL 'EQEQ(Any?, Any?): Boolean' type=kotlin.Boolean origin=EQEQ
+                arg0: GET_VAR 'tmp0_safe_receiver: Any?' type=kotlin.Any? origin=null
+                arg1: CONST Null type=kotlin.Nothing? value='null'
+              then: CONST Null type=kotlin.Nothing? value='null'
+            BRANCH
+              if: CONST Boolean type=kotlin.Boolean value='true'
+              then: CALL 'invoke(Any, Int, String): Unit' type=kotlin.Unit origin=INVOKE
+                $this: GET_VAR 'value-parameter fn: Any.(Int, String) -> Unit' type=kotlin.Any.(kotlin.Int, kotlin.String) -> kotlin.Unit origin=VARIABLE_AS_FUNCTION
+                p1: GET_VAR 'tmp0_safe_receiver: Any?' type=kotlin.Any? origin=null
+                p2: CONST Int type=kotlin.Int value='42'
+                p3: CONST String type=kotlin.String value='Hello'

compiler/testData/ir/irText/expressions/variableAsFunctionCall.txt

@@ -19,9 +19,9 @@ FILE /variableAsFunctionCall.kt
     VALUE_PARAMETER value-parameter f: kotlin.String.() -> kotlin.Unit
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='test2(String.() -> Unit): Unit'
-        CALL 'invoke() on String: Unit' type=kotlin.Unit origin=INVOKE
+        CALL 'invoke(String): Unit' type=kotlin.Unit origin=INVOKE
           $this: GET_VAR 'value-parameter f: String.() -> Unit' type=kotlin.String.() -> kotlin.Unit origin=VARIABLE_AS_FUNCTION
-          $receiver: CONST String type=kotlin.String value='hello'
+          p1: CONST String type=kotlin.String value='hello'
   FUN public fun test3(): kotlin.String
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='test3(): String'

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -635,6 +635,12 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             doTest(fileName);
         }
 
+        @TestMetadata("extFunSafeInvoke.kt")
+        public void testExtFunSafeInvoke() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/ir/irText/expressions/extFunSafeInvoke.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("extensionPropertyGetterCall.kt")
         public void testExtensionPropertyGetterCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/ir/irText/expressions/extensionPropertyGetterCall.kt");