Intentions: Convert function type parameter to receiver

#KT-14246 In Progress

compiler/frontend/src/org/jetbrains/kotlin/psi/KtFunctionType.java

@@ -75,9 +75,14 @@ public class KtFunctionType extends KtElementImplStub<KotlinPlaceHolderStub<KtFu
         return list != null ? list.getParameters() : Collections.<KtParameter>emptyList();
     }
 
+    @Nullable
+    public KtFunctionTypeReceiver getReceiver() {
+        return getStubOrPsiChild(KtStubElementTypes.FUNCTION_TYPE_RECEIVER);
+    }
+
     @Nullable
     public KtTypeReference getReceiverTypeReference() {
-        KtFunctionTypeReceiver receiverDeclaration = getStubOrPsiChild(KtStubElementTypes.FUNCTION_TYPE_RECEIVER);
+        KtFunctionTypeReceiver receiverDeclaration = getReceiver();
         if (receiverDeclaration == null) {
             return null;
         }

compiler/frontend/src/org/jetbrains/kotlin/psi/KtParameterList.java

@@ -64,6 +64,10 @@ public class KtParameterList extends KtElementImplStub<KotlinPlaceHolderStub<KtP
         EditCommaSeparatedListHelper.INSTANCE.removeItem(parameter);
     }
 
+    public void removeParameter(int index) {
+        removeParameter(getParameters().get(index));
+    }
+
     public KtFunction getOwnerFunction() {
         PsiElement parent = getParentByStub();
         if (!(parent instanceof KtFunction)) return null;

compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt

@@ -108,6 +108,10 @@ class KtPsiFactory(private val project: Project) {
         return if (typeReference?.text == type) typeReference else null
     }
 
+    fun createFunctionTypeReceiver(typeReference: KtTypeReference): KtFunctionTypeReceiver {
+        return (createType("A.() -> B").typeElement as KtFunctionType).receiver!!.apply { this.typeReference.replace(typeReference) }
+    }
+
     fun createTypeAlias(name: String, typeParameters: List<String>, typeElement: KtTypeElement): KtTypeAlias {
         return createTypeAlias(name, typeParameters, "X").apply { getTypeReference()!!.replace(createType(typeElement)) }
     }

compiler/frontend/src/org/jetbrains/kotlin/psi/KtValueArgumentList.java

@@ -69,4 +69,8 @@ public class KtValueArgumentList extends KtElementImpl {
         assert argument.getParent() == this;
         EditCommaSeparatedListHelper.INSTANCE.removeItem(argument);
     }
+
+    public void removeArgument(int index) {
+        removeArgument(getArguments().get(index));
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/TypeRefHelpers.kt

@@ -20,10 +20,8 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiErrorElement
 import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtCallableDeclaration
-import org.jetbrains.kotlin.psi.KtFunctionType
-import org.jetbrains.kotlin.psi.KtPsiFactory
-import org.jetbrains.kotlin.psi.KtTypeReference
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
@@ -56,18 +54,21 @@ fun setTypeReference(declaration: KtCallableDeclaration, addAfter: PsiElement?,
     }
 }
 
-fun KtCallableDeclaration.setReceiverTypeReference(typeRef: KtTypeReference?): KtTypeReference? {
+private inline fun <T : KtElement> T.doSetReceiverTypeReference(
+        typeRef: KtTypeReference?,
+        getReceiverTypeReference: T.() -> KtTypeReference?,
+        addReceiverTypeReference: T.(typeRef: KtTypeReference) -> KtTypeReference
+): KtTypeReference? {
     val needParentheses = typeRef != null && typeRef.typeElement is KtFunctionType && !typeRef.hasParentheses()
-    val oldTypeRef = receiverTypeReference
+    val oldTypeRef = getReceiverTypeReference()
     if (typeRef != null) {
         val newTypeRef =
                 if (oldTypeRef != null) {
                     oldTypeRef.replace(typeRef) as KtTypeReference
                 }
                 else {
-                    val anchor = nameIdentifier ?: valueParameterList
-                    val newTypeRef = addBefore(typeRef, anchor) as KtTypeReference
-                    addAfter(KtPsiFactory(project).createDot(), newTypeRef)
+                    val newTypeRef = addReceiverTypeReference(typeRef)
+                    addAfter(KtPsiFactory(project).createDot(), newTypeRef.parentsWithSelf.first { it.parent == this })
                     newTypeRef
                 }
         if (needParentheses) {
@@ -79,9 +80,27 @@ fun KtCallableDeclaration.setReceiverTypeReference(typeRef: KtTypeReference?): K
     }
     else {
         if (oldTypeRef != null) {
-            val dot = oldTypeRef.siblings(forward = true).firstOrNull { it.node.elementType == KtTokens.DOT }
-            deleteChildRange(oldTypeRef, dot ?: oldTypeRef)
+            val dotSibling = oldTypeRef.parent as? KtFunctionTypeReceiver ?: oldTypeRef
+            val dot = dotSibling.siblings(forward = true).firstOrNull { it.node.elementType == KtTokens.DOT }
+            deleteChildRange(dotSibling, dot ?: dotSibling)
         }
         return null
     }
-}
+}
+
+fun KtCallableDeclaration.setReceiverTypeReference(typeRef: KtTypeReference?) =
+        doSetReceiverTypeReference(
+                typeRef,
+                { receiverTypeReference },
+                { this.addBefore(it, nameIdentifier ?: valueParameterList) as KtTypeReference }
+        )
+
+fun KtFunctionType.setReceiverTypeReference(typeRef: KtTypeReference?) =
+        doSetReceiverTypeReference(
+                typeRef,
+                { receiverTypeReference },
+                {
+                    (addBefore(KtPsiFactory(project).createFunctionTypeReceiver(it),
+                               parameterList ?: firstChild) as KtFunctionTypeReceiver).typeReference
+                }
+        )

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/resolvedCallUtil.kt

@@ -17,9 +17,9 @@
 package org.jetbrains.kotlin.resolve.calls.resolvedCallUtil
 
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.psi.KtPsiUtil
-import org.jetbrains.kotlin.psi.KtThisExpression
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.callUtil.isSafeCall
 import org.jetbrains.kotlin.resolve.calls.context.CallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
@@ -112,3 +112,9 @@ fun ResolvedCall<*>.hasBothReceivers() = dispatchReceiver != null && extensionRe
 
 fun ResolvedCall<*>.getDispatchReceiverWithSmartCast(): ReceiverValue?
         = getReceiverValueWithSmartCast(dispatchReceiver, smartCastDispatchReceiverType)
+
+fun KtCallElement.getArgumentByParameterIndex(index: Int, context: BindingContext): List<ValueArgument> {
+    val resolvedCall = getResolvedCall(context) ?: return emptyList()
+    val parameterToProcess = resolvedCall.resultingDescriptor.valueParameters.getOrNull(index) ?: return emptyList()
+    return resolvedCall.valueArguments[parameterToProcess]?.arguments ?: emptyList()
+}