Allow to turn the first parameter of a SAM-converted lambda into the receiver (KT-12848)

compiler/frontend.java/src/org/jetbrains/kotlin/frontend/java/di/injection.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
 import org.jetbrains.kotlin.load.java.components.*
 import org.jetbrains.kotlin.load.java.lazy.ModuleClassResolver
 import org.jetbrains.kotlin.load.java.sam.SamConversionResolverImpl
+import org.jetbrains.kotlin.load.java.sam.SamWithReceiverResolver
 import org.jetbrains.kotlin.load.kotlin.DeserializationComponentsForJava
 import org.jetbrains.kotlin.load.kotlin.JvmVirtualFileFinderFactory
 import org.jetbrains.kotlin.platform.JvmBuiltIns
@@ -68,7 +69,8 @@ fun StorageComponentContainer.configureJavaTopDownAnalysis(
     useImpl<TraceBasedErrorReporter>()
     useImpl<PsiBasedExternalAnnotationResolver>()
     useImpl<JavaPropertyInitializerEvaluatorImpl>()
-    useInstance(SamConversionResolverImpl)
+    useImpl<SamWithReceiverResolver>()
+    useImpl<SamConversionResolverImpl>()
     useImpl<JavaSourceElementFactoryImpl>()
     useInstance(InternalFlexibleTypeTransformer)
 

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/sam/SamConversionResolverImpl.kt

@@ -25,9 +25,10 @@ import org.jetbrains.kotlin.load.java.descriptors.JavaClassConstructorDescriptor
 import org.jetbrains.kotlin.load.java.descriptors.JavaMethodDescriptor
 import org.jetbrains.kotlin.load.java.descriptors.SamConstructorDescriptor
 import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaClassDescriptor
+import org.jetbrains.kotlin.storage.StorageManager
 import org.jetbrains.kotlin.types.SimpleType
 
-object SamConversionResolverImpl : SamConversionResolver {
+class SamConversionResolverImpl(val storageManager: StorageManager, val samWithReceiverResolver: SamWithReceiverResolver): SamConversionResolver {
     override fun resolveSamConstructor(constructorOwner: DeclarationDescriptor, classifier: () -> ClassifierDescriptor?): SamConstructorDescriptor? {
         val classifierDescriptor = classifier()
         if (classifierDescriptor !is LazyJavaClassDescriptor || classifierDescriptor.functionTypeForSamInterface == null) return null
@@ -44,8 +45,13 @@ object SamConversionResolverImpl : SamConversionResolver {
         }
     }
 
+    private val functionTypesForSamInterfaces = storageManager.createCacheWithNullableValues<JavaClassDescriptor, SimpleType>()
+
     override fun resolveFunctionTypeIfSamInterface(classDescriptor: JavaClassDescriptor): SimpleType? {
-        val abstractMethod = SingleAbstractMethodUtils.getSingleAbstractMethodOrNull(classDescriptor) ?: return null
-        return SingleAbstractMethodUtils.getFunctionTypeForAbstractMethod(abstractMethod)
+        return functionTypesForSamInterfaces.computeIfAbsent(classDescriptor) {
+            val abstractMethod = SingleAbstractMethodUtils.getSingleAbstractMethodOrNull(classDescriptor) ?: return@computeIfAbsent null
+            val shouldConvertFirstParameterToDescriptor = samWithReceiverResolver.shouldConvertFirstSamParameterToReceiver(abstractMethod)
+            SingleAbstractMethodUtils.getFunctionTypeForAbstractMethod(abstractMethod, shouldConvertFirstParameterToDescriptor)
+        }
     }
 }

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/sam/SamWithReceiverResolver.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2016 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.load.java.sam
+
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+
+class SamWithReceiverResolver {
+    private val extensions = mutableListOf<Extension>()
+
+    fun registerExtension(extension: Extension) {
+        extensions += extension
+    }
+
+    fun shouldConvertFirstSamParameterToReceiver(function: FunctionDescriptor): Boolean {
+        return extensions.any { it.shouldConvertFirstSamParameterToReceiver(function) }
+    }
+
+    interface Extension {
+        fun shouldConvertFirstSamParameterToReceiver(function: FunctionDescriptor): Boolean
+    }
+}

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/sam/SingleAbstractMethodUtils.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.load.java.sam;
 
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.*;
@@ -23,19 +24,19 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl;
 import org.jetbrains.kotlin.descriptors.impl.TypeParameterDescriptorImpl;
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl;
+import org.jetbrains.kotlin.extensions.DeclarationAttributeAltererExtension;
 import org.jetbrains.kotlin.load.java.descriptors.*;
+import org.jetbrains.kotlin.load.java.sources.JavaSourceElement;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.name.SpecialNames;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.calls.util.FunctionTypeResolveUtilsKt;
 import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.jvm.JavaResolverUtils;
+import org.jetbrains.kotlin.resolve.source.PsiSourceElement;
 import org.jetbrains.kotlin.types.*;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 import static org.jetbrains.kotlin.types.Variance.IN_VARIANCE;
 
@@ -97,18 +98,33 @@ public class SingleAbstractMethodUtils {
     }
 
     @NotNull
-    public static SimpleType getFunctionTypeForAbstractMethod(@NotNull FunctionDescriptor function) {
+    public static SimpleType getFunctionTypeForAbstractMethod(
+            @NotNull FunctionDescriptor function,
+            boolean shouldConvertFirstParameterToDescriptor
+    ) {
         KotlinType returnType = function.getReturnType();
         assert returnType != null : "function is not initialized: " + function;
         List<ValueParameterDescriptor> valueParameters = function.getValueParameters();
         List<KotlinType> parameterTypes = new ArrayList<KotlinType>(valueParameters.size());
         List<Name> parameterNames = new ArrayList<Name>(valueParameters.size());
-        for (ValueParameterDescriptor parameter : valueParameters) {
+
+        int startIndex = 0;
+        KotlinType receiverType = null;
+
+        if (shouldConvertFirstParameterToDescriptor && !function.getValueParameters().isEmpty()) {
+            receiverType = valueParameters.get(0).getType();
+            startIndex = 1;
+        }
+
+        for (int i = startIndex; i < valueParameters.size(); ++i) {
+            ValueParameterDescriptor parameter = valueParameters.get(i);
             parameterTypes.add(parameter.getType());
             parameterNames.add(function.hasSynthesizedParameterNames() ? SpecialNames.NO_NAME_PROVIDED : parameter.getName());
         }
+
         return FunctionTypeResolveUtilsKt.createFunctionType(
-                DescriptorUtilsKt.getBuiltIns(function), Annotations.Companion.getEMPTY(), null, parameterTypes, parameterNames, returnType
+                DescriptorUtilsKt.getBuiltIns(function), Annotations.Companion.getEMPTY(),
+                receiverType, parameterTypes, parameterNames, returnType
         );
     }
 

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/JvmAnalyzerFacade.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.context.ModuleContext
 import org.jetbrains.kotlin.descriptors.PackagePartProvider
 import org.jetbrains.kotlin.descriptors.impl.CompositePackageFragmentProvider
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.frontend.java.di.createContainerForLazyResolveWithJava
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.load.java.lazy.ModuleClassResolverImpl
@@ -91,6 +92,9 @@ object JvmAnalyzerFacade : AnalyzerFacade<JvmPlatformParameters>() {
                 useBuiltInsProvider = false, // TODO: load built-ins from module dependencies in IDE
                 useLazyResolve = true
         )
+
+        StorageComponentContainerContributor.getInstances(project).forEach { it.onContainerComposed(container, moduleInfo) }
+
         val resolveSession = container.get<ResolveSession>()
         val javaDescriptorResolver = container.get<JavaDescriptorResolver>()
 

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/TopDownAnalyzerFacadeForJVM.kt

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.descriptors.PackageFragmentProvider
 import org.jetbrains.kotlin.descriptors.PackagePartProvider
 import org.jetbrains.kotlin.descriptors.impl.CompositePackageFragmentProvider
 import org.jetbrains.kotlin.descriptors.impl.ModuleDependenciesImpl
+import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.frontend.java.di.createContainerForTopDownAnalyzerForJvm
 import org.jetbrains.kotlin.frontend.java.di.initJvmBuiltInsForTopDownAnalysis
 import org.jetbrains.kotlin.frontend.java.di.initialize
@@ -80,6 +81,8 @@ object TopDownAnalyzerFacadeForJVM {
                 project, files, trace, configuration, packagePartProvider, declarationProviderFactory, sourceModuleSearchScope
         )
 
+        StorageComponentContainerContributor.getInstances(project).forEach { it.onContainerComposed(container, null) }
+
         val module = container.get<ModuleDescriptor>()
         val moduleContext = container.get<ModuleContext>()