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>()
 

compiler/frontend/src/org/jetbrains/kotlin/extensions/AnnotationBasedExtension.kt

@@ -24,9 +24,9 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.types.TypeUtils
 
 interface AnnotationBasedExtension {
-    fun getAnnotationFqNames(modifierListOwner: KtModifierListOwner): List<String>
+    fun getAnnotationFqNames(modifierListOwner: KtModifierListOwner?): List<String>
 
-    fun DeclarationDescriptor.hasSpecialAnnotation(modifierListOwner: KtModifierListOwner): Boolean {
+    fun DeclarationDescriptor.hasSpecialAnnotation(modifierListOwner: KtModifierListOwner?): Boolean {
         if (annotations.any { it.isASpecialAnnotation(modifierListOwner) }) return true
 
         if (this is ClassDescriptor) {
@@ -40,7 +40,7 @@ interface AnnotationBasedExtension {
     }
 
     private fun AnnotationDescriptor.isASpecialAnnotation(
-            modifierListOwner: KtModifierListOwner,
+            modifierListOwner: KtModifierListOwner?,
             allowMetaAnnotations: Boolean = true
     ): Boolean {
         val annotationType = type.constructor.declarationDescriptor ?: return false

compiler/frontend/src/org/jetbrains/kotlin/extensions/DeclarationAttributeAltererExtension.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.extensions
 
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -37,4 +38,6 @@ interface DeclarationAttributeAltererExtension {
             currentModality: Modality,
             bindingContext: BindingContext
     ): Modality? = null
+
+    fun shouldConvertFirstSAMParameterToReceiver(function: FunctionDescriptor) : Boolean = false
 }

compiler/frontend/src/org/jetbrains/kotlin/extensions/StorageComponentContainerContributor.kt

@@ -16,6 +16,8 @@
 
 package org.jetbrains.kotlin.extensions
 
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.container.ComponentProvider
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.resolve.TargetPlatform
 
@@ -24,5 +26,7 @@ interface StorageComponentContainerContributor {
             "org.jetbrains.kotlin.storageComponentContainerContributor", StorageComponentContainerContributor::class.java
     )
 
-    fun addDeclarations(container: StorageComponentContainer, platform: TargetPlatform)
+    fun addDeclarations(container: StorageComponentContainer, platform: TargetPlatform) {}
+
+    fun onContainerComposed(container: ComponentProvider, moduleInfo: ModuleInfo?) {}
 }

compiler/frontend/src/org/jetbrains/kotlin/script/KotlinScriptDefinition.kt

@@ -34,6 +34,9 @@ open class KotlinScriptDefinition(val template: KClass<out Any>) {
     // TODO: consider creating separate type (subtype? for kotlin scripts)
     open val fileType: LanguageFileType = KotlinFileType.INSTANCE
 
+    open val annotationsForSamWithReceivers: List<String>
+        get() = emptyList()
+
     open fun <TF> isScript(file: TF): Boolean =
             getFileName(file).endsWith(KotlinParserDefinition.STD_SCRIPT_EXT)
 

compiler/tests-common/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTest.java

@@ -23,6 +23,7 @@ import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.search.GlobalSearchScope;
+import junit.framework.TestCase;
 import kotlin.TuplesKt;
 import kotlin.collections.CollectionsKt;
 import kotlin.collections.MapsKt;
@@ -189,7 +190,7 @@ public abstract class AbstractDiagnosticsTest extends BaseDiagnosticsTest {
         }
         else {
             File lazyLogFile = getLazyLogFile(testDataFile);
-            assertFalse("No lazy log expected, but found: " + lazyLogFile.getAbsolutePath(), lazyLogFile.exists());
+            TestCase.assertFalse("No lazy log expected, but found: " + lazyLogFile.getAbsolutePath(), lazyLogFile.exists());
         }
 
         Throwable exceptionFromDescriptorValidation = null;
@@ -225,7 +226,7 @@ public abstract class AbstractDiagnosticsTest extends BaseDiagnosticsTest {
 
         KotlinTestUtils.assertEqualsToFile(testDataFile, actualText.toString());
 
-        assertTrue("Diagnostics mismatch. See the output above", ok);
+        TestCase.assertTrue("Diagnostics mismatch. See the output above", ok);
 
         // now we throw a previously found error, if any
         if (exceptionFromDescriptorValidation != null) {
@@ -440,7 +441,7 @@ public abstract class AbstractDiagnosticsTest extends BaseDiagnosticsTest {
                 return InTextDirectivesUtils.isDirectiveDefined(file.expectedText, "// SKIP_TXT");
             }
         })) {
-            assertFalse(".txt file should not exist if SKIP_TXT directive is used: " + expectedFile, expectedFile.exists());
+            TestCase.assertFalse(".txt file should not exist if SKIP_TXT directive is used: " + expectedFile, expectedFile.exists());
             return;
         }
 
@@ -452,7 +453,7 @@ public abstract class AbstractDiagnosticsTest extends BaseDiagnosticsTest {
         for (Iterator<TestModule> module = CollectionsKt.sorted(modules.keySet()).iterator(); module.hasNext(); ) {
             ModuleDescriptorImpl moduleDescriptor = modules.get(module.next());
             PackageViewDescriptor aPackage = moduleDescriptor.getPackage(FqName.ROOT);
-            assertFalse(aPackage.isEmpty());
+            TestCase.assertFalse(aPackage.isEmpty());
 
             if (isMultiModuleTest) {
                 rootPackageText.append(String.format("// -- Module: %s --\n", moduleDescriptor.getName()));
@@ -587,8 +588,8 @@ public abstract class AbstractDiagnosticsTest extends BaseDiagnosticsTest {
             DiagnosticUtils.LineAndColumn lineAndColumn =
                     DiagnosticUtils.getLineAndColumnInPsiFile(element.getContainingFile(), element.getTextRange());
 
-            assertTrue("Resolved call for '" + element.getText() + "'" + lineAndColumn + " is not completed",
-                       ((MutableResolvedCall<?>) resolvedCall).isCompleted());
+            TestCase.assertTrue("Resolved call for '" + element.getText() + "'" + lineAndColumn + " is not completed",
+                                ((MutableResolvedCall<?>) resolvedCall).isCompleted());
         }
 
         checkResolvedCallsInDiagnostics(bindingContext);
@@ -633,8 +634,8 @@ public abstract class AbstractDiagnosticsTest extends BaseDiagnosticsTest {
         DiagnosticUtils.LineAndColumn lineAndColumn =
                 DiagnosticUtils.getLineAndColumnInPsiFile(element.getContainingFile(), element.getTextRange());
 
-        assertTrue("Resolved calls stored in " + diagnostic.getFactory().getName() + "\n" +
-                   "for '" + element.getText() + "'" + lineAndColumn + " are not completed",
-                   allCallsAreCompleted);
+        TestCase.assertTrue("Resolved calls stored in " + diagnostic.getFactory().getName() + "\n" +
+                            "for '" + element.getText() + "'" + lineAndColumn + " are not completed",
+                            allCallsAreCompleted);
     }
 }