Remove StorageComponentContainerContributor::onContainerComposed

Rename addDeclarations -> registerModuleComponents
Use it to provide SamWithReceiverResolver extensions instead

Post construction on container composition can be achieved
    but manually inserting injections where it seems appropriate
    is bug prone
This fixes a bug where SamWithReceiverPlugin extension was not registered
    for some containers in IDE which led to incorrect highlighting in IDE
Add IDE test for applying SamWithReceiver plugin

 #KT-18062 Fixed

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/TopDownAnalyzerFacadeForJVM.kt

@@ -174,8 +174,6 @@ object TopDownAnalyzerFacadeForJVM {
                     packagePartProvider(dependencyScope), moduleClassResolver, jvmTarget, languageVersionSettings, configureJavaClassFinder
             )
 
-            StorageComponentContainerContributor.getInstances(project).forEach { it.onContainerComposed(dependenciesContainer, null) }
-
             moduleClassResolver.compiledCodeResolver = dependenciesContainer.get<JavaDescriptorResolver>()
 
             dependenciesContext.setDependencies(listOfNotNull(dependenciesContext.module, optionalBuiltInsModule))
@@ -202,8 +200,6 @@ object TopDownAnalyzerFacadeForJVM {
         ).apply {
             initJvmBuiltInsForTopDownAnalysis()
             (partProvider as? IncrementalPackagePartProvider)?.deserializationConfiguration = get<DeserializationConfiguration>()
-
-            StorageComponentContainerContributor.getInstances(project).forEach { it.onContainerComposed(this, null) }
         }
 
         moduleClassResolver.sourceCodeResolver = container.get<JavaDescriptorResolver>()

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

@@ -29,11 +29,11 @@ import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.PackagePartProvider
 import org.jetbrains.kotlin.frontend.di.configureModule
 import org.jetbrains.kotlin.incremental.components.LookupTracker
-import org.jetbrains.kotlin.load.java.*
+import org.jetbrains.kotlin.load.java.AbstractJavaClassFinder
+import org.jetbrains.kotlin.load.java.InternalFlexibleTypeTransformer
+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.VirtualFileFinderFactory
 import org.jetbrains.kotlin.platform.JvmBuiltIns
@@ -64,8 +64,6 @@ private fun StorageComponentContainer.configureJavaTopDownAnalysis(
     useImpl<SignaturePropagatorImpl>()
     useImpl<TraceBasedErrorReporter>()
     useImpl<PsiBasedExternalAnnotationResolver>()
-    useInstance(SamWithReceiverResolver())
-    useImpl<SamConversionResolverImpl>()
     useInstance(InternalFlexibleTypeTransformer)
 
     useImpl<CompilerDeserializationConfiguration>()

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

@@ -21,13 +21,16 @@ import org.jetbrains.kotlin.load.java.descriptors.JavaClassDescriptor
 import org.jetbrains.kotlin.storage.StorageManager
 import org.jetbrains.kotlin.types.SimpleType
 
-class SamConversionResolverImpl(val storageManager: StorageManager, val samWithReceiverResolver: SamWithReceiverResolver): SamConversionResolver {
+class SamConversionResolverImpl(
+        storageManager: StorageManager,
+        private val samWithReceiverResolvers: Iterable<SamWithReceiverResolver>
+): SamConversionResolver {
     private val functionTypesForSamInterfaces = storageManager.createCacheWithNullableValues<JavaClassDescriptor, SimpleType>()
 
     override fun resolveFunctionTypeIfSamInterface(classDescriptor: JavaClassDescriptor): SimpleType? {
         return functionTypesForSamInterfaces.computeIfAbsent(classDescriptor) {
             val abstractMethod = SingleAbstractMethodUtils.getSingleAbstractMethodOrNull(classDescriptor) ?: return@computeIfAbsent null
-            val shouldConvertFirstParameterToDescriptor = samWithReceiverResolver.shouldConvertFirstSamParameterToReceiver(abstractMethod)
+            val shouldConvertFirstParameterToDescriptor = samWithReceiverResolvers.any { it.shouldConvertFirstSamParameterToReceiver(abstractMethod) }
             SingleAbstractMethodUtils.getFunctionTypeForAbstractMethod(abstractMethod, shouldConvertFirstParameterToDescriptor)
         }
     }

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

@@ -18,18 +18,6 @@ package org.jetbrains.kotlin.load.java.sam
 
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 
-class SamWithReceiverResolver {
+interface SamWithReceiverResolver {
-    private val extensions = mutableListOf<Extension>()
+    fun shouldConvertFirstSamParameterToReceiver(function: FunctionDescriptor): Boolean
-
-    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/resolve/jvm/JvmAnalyzerFacade.kt

@@ -100,8 +100,6 @@ object JvmAnalyzerFacade : AnalyzerFacade() {
                 useBuiltInsProvider = false // TODO: load built-ins from module dependencies in IDE
         )
 
-        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/platform/JvmPlatformConfigurator.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.resolve.jvm.platform
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.useImpl
 import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.load.java.sam.SamConversionResolverImpl
 import org.jetbrains.kotlin.platform.JavaToKotlinClassMap
 import org.jetbrains.kotlin.resolve.PlatformConfigurator
 import org.jetbrains.kotlin.resolve.calls.checkers.ReifiedTypeParameterSubstitutionChecker
@@ -83,6 +84,7 @@ object JvmPlatformConfigurator : PlatformConfigurator(
     override fun configureModuleComponents(container: StorageComponentContainer) {
         container.useImpl<JvmReflectionAPICallChecker>()
         container.useImpl<JavaSyntheticScopes>()
+        container.useImpl<SamConversionResolverImpl>()
         container.useImpl<InterfaceDefaultMethodCallChecker>()
         container.useImpl<InlinePlatformCompatibilityChecker>()
         container.useImpl<JvmModuleAccessibilityChecker>()

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

@@ -16,9 +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.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.resolve.TargetPlatform
 
 interface StorageComponentContainerContributor {
@@ -26,7 +25,5 @@ interface StorageComponentContainerContributor {
             "org.jetbrains.kotlin.storageComponentContainerContributor", StorageComponentContainerContributor::class.java
     )
 
-    fun addDeclarations(container: StorageComponentContainer, platform: TargetPlatform) {}
+    fun registerModuleComponents(container: StorageComponentContainer, platform: TargetPlatform, moduleDescriptor: ModuleDescriptor) {}
-
-    fun onContainerComposed(container: ComponentProvider, moduleInfo: ModuleInfo?) {}
 }

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

@@ -53,7 +53,7 @@ fun StorageComponentContainer.configureModule(
     platform.platformConfigurator.configureModuleComponents(this)
 
     for (extension in StorageComponentContainerContributor.getInstances(moduleContext.project)) {
-        extension.addDeclarations(this, platform)
+        extension.registerModuleComponents(this, platform,  moduleContext.module)
     }
 
     configurePlatformIndependentComponents()

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/components/SamConversionResolver.kt

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.load.java.descriptors.JavaClassDescriptor
 import org.jetbrains.kotlin.types.SimpleType
 
 interface SamConversionResolver {
-    companion object EMPTY : SamConversionResolver {
+    object Empty : SamConversionResolver {
         override fun resolveFunctionTypeIfSamInterface(classDescriptor: JavaClassDescriptor): SimpleType? = null
     }
 

core/descriptors.runtime/src/org/jetbrains/kotlin/load/kotlin/reflect/RuntimeModuleData.kt

@@ -63,7 +63,7 @@ class RuntimeModuleData private constructor(
             val globalJavaResolverContext = JavaResolverComponents(
                     storageManager, ReflectJavaClassFinder(classLoader), reflectKotlinClassFinder, deserializedDescriptorResolver,
                     ExternalAnnotationResolver.EMPTY, SignaturePropagator.DO_NOTHING, RuntimeErrorReporter, javaResolverCache,
-                    JavaPropertyInitializerEvaluator.DoNothing, SamConversionResolver, RuntimeSourceElementFactory, singleModuleClassResolver,
+                    JavaPropertyInitializerEvaluator.DoNothing, SamConversionResolver.Empty, RuntimeSourceElementFactory, singleModuleClassResolver,
                     runtimePackagePartProvider, SupertypeLoopChecker.EMPTY, LookupTracker.DO_NOTHING, module,
                     ReflectionTypes(module, notFoundClasses),
                     annotationTypeQualifierResolver,

idea/idea.iml

@@ -70,5 +70,6 @@
     <orderEntry type="module" module-name="backend.jvm" />
     <orderEntry type="library" name="kotlin-reflect" level="project" />
     <orderEntry type="module" module-name="idea-gradle" scope="TEST" />
+    <orderEntry type="module" module-name="sam-with-receiver-cli" scope="TEST" />
   </component>
 </module>

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m1/annotations.kt

@@ -0,0 +1,4 @@
+package anno
+
+annotation class A
+annotation class B

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m1/javaInterface/API.java

@@ -0,0 +1,8 @@
+package javaInterface;
+
+public class API {
+    public <T> void useM1A(InterfaceM1A<T> m1A);
+    public <G> void useM1B(InterfaceM1B<G> m1B);
+    public <T> void useM2A(InterfaceM2A<T> m1A);
+    public <G> void useM2B(InterfaceM2B<G> m1B);
+}

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m1/javaInterface/InterfaceM1A.java

@@ -0,0 +1,6 @@
+package javaInterface;
+
+@anno.A
+interface InterfaceM1A<T> {
+    void foo(T t);
+}

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m1/javaInterface/InterfaceM1B.java

@@ -0,0 +1,6 @@
+package javaInterface;
+
+@anno.B
+interface InterfaceM1B<T> {
+    void foo(T t);
+}

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m1/m1.kt

@@ -0,0 +1,31 @@
+package m1
+
+public fun testUseAsReceiver(api: javaInterface.API) {
+    api.useM1A<String> {
+        this.length
+    }
+    api.useM1B<String> {
+        <error>this</error>.length
+    }
+    api.useM2A<String> {
+        this.length
+    }
+    api.useM2B<String> {
+        <error>this</error>.length
+    }
+}
+
+public fun testUseAsParameter(api: javaInterface.API) {
+    api.useM1A<String> {
+        <error>it</error>.length
+    }
+    api.useM1B<String> {
+        it.length
+    }
+    api.useM2A<String> {
+        <error>it</error>.length
+    }
+    api.useM2B<String> {
+        it.length
+    }
+}

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m2/javaInterface/InterfaceM2A.java

@@ -0,0 +1,6 @@
+package javaInterface;
+
+@anno.A
+interface InterfaceM2A<T> {
+    void foo(T t);
+}

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m2/javaInterface/InterfaceM2B.java

@@ -0,0 +1,6 @@
+package javaInterface;
+
+@anno.B
+interface InterfaceM2B<T> {
+    void foo(T t);
+}

idea/testData/multiModuleHighlighting/samWithReceiverExtension/m2/m2.kt

@@ -0,0 +1,31 @@
+package m2
+
+public fun testUseAsReceiver(api: javaInterface.API) {
+    api.useM1A<String> {
+        <error>this</error>.length
+    }
+    api.useM1B<String> {
+        this.length
+    }
+    api.useM2A<String> {
+        <error>this</error>.length
+    }
+    api.useM2B<String> {
+        this.length
+    }
+}
+
+public fun testUseAsParameter(api: javaInterface.API) {
+    api.useM1A<String> {
+        it.length
+    }
+    api.useM1B<String> {
+        <error>it</error>.length
+    }
+    api.useM2A<String> {
+        it.length
+    }
+    api.useM2B<String> {
+        <error>it</error>.length
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/caches/resolve/MultiModuleHighlightingTest.kt

@@ -44,6 +44,8 @@ import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
 import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.idea.util.projectStructure.sdk
+import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.ANNOTATION_OPTION
+import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.PLUGIN_ID
 import org.jetbrains.kotlin.test.TestJdkKind.FULL_JDK
 
 open class MultiModuleHighlightingTest : AbstractMultiModuleHighlightingTest() {
@@ -193,6 +195,25 @@ open class MultiModuleHighlightingTest : AbstractMultiModuleHighlightingTest() {
         checkHighlightingInAllFiles()
     }
 
+    fun testSamWithReceiverExtension() {
+        val module1 = module("m1").setupKotlinFacet {
+            settings.compilerArguments!!.pluginOptions =
+                    arrayOf("plugin:${PLUGIN_ID}:${ANNOTATION_OPTION.name}=anno.A")
+        }
+
+        val module2 = module("m2").setupKotlinFacet {
+            settings.compilerArguments!!.pluginOptions =
+                    arrayOf("plugin:${PLUGIN_ID}:${ANNOTATION_OPTION.name}=anno.B")
+        }
+
+
+        module1.addDependency(module2)
+        module2.addDependency(module1)
+
+        checkHighlightingInAllFiles()
+    }
+
+
     private fun Module.setupKotlinFacet(configure: KotlinFacetConfiguration.() -> Unit) = apply {
         runWriteAction {
             val facet = FacetManager.getInstance(this).addFacet(KotlinFacetType.INSTANCE, KotlinFacetType.NAME, null)

plugins/android-extensions/android-extensions-compiler/src/org/jetbrains/kotlin/android/synthetic/AndroidComponentRegistrar.kt

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.android.synthetic
 
 import com.intellij.mock.MockProject
-import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.project.Project
 import kotlinx.android.extensions.CacheImplementation
 import org.jetbrains.kotlin.android.parcel.ParcelableCodegenExtension
@@ -25,8 +24,8 @@ import org.jetbrains.kotlin.android.parcel.ParcelableDeclarationChecker
 import org.jetbrains.kotlin.android.parcel.ParcelableResolveExtension
 import org.jetbrains.kotlin.android.synthetic.codegen.CliAndroidExtensionsExpressionCodegenExtension
 import org.jetbrains.kotlin.android.synthetic.codegen.CliAndroidOnDestroyClassBuilderInterceptorExtension
+import org.jetbrains.kotlin.android.synthetic.codegen.ParcelableClinitClassBuilderInterceptorExtension
 import org.jetbrains.kotlin.android.synthetic.diagnostic.AndroidExtensionPropertiesCallChecker
-import org.jetbrains.kotlin.android.synthetic.diagnostic.DefaultErrorMessagesAndroid
 import org.jetbrains.kotlin.android.synthetic.res.AndroidLayoutXmlFileManager
 import org.jetbrains.kotlin.android.synthetic.res.AndroidVariant
 import org.jetbrains.kotlin.android.synthetic.res.CliAndroidLayoutXmlFileManager
@@ -41,13 +40,12 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.CompilerConfigurationKey
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.useInstance
-import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.resolve.TargetPlatform
+import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 import org.jetbrains.kotlin.resolve.jvm.extensions.PackageFragmentProviderExtension
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatform
-import org.jetbrains.kotlin.android.synthetic.codegen.ParcelableClinitClassBuilderInterceptorExtension
-import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 
 object AndroidConfigurationKeys {
     val VARIANT = CompilerConfigurationKey.create<List<String>>("Android build variant")
@@ -130,10 +128,12 @@ class AndroidComponentRegistrar : ComponentRegistrar {
 }
 
 class AndroidExtensionPropertiesComponentContainerContributor : StorageComponentContainerContributor {
-    override fun addDeclarations(container: StorageComponentContainer, platform: TargetPlatform) {
+    override fun registerModuleComponents(
-        if (platform is JvmPlatform) {
+            container: StorageComponentContainer, platform: TargetPlatform, moduleDescriptor: ModuleDescriptor
-            container.useInstance(AndroidExtensionPropertiesCallChecker())
+    ) {
-            container.useInstance(ParcelableDeclarationChecker())
+        if (platform != JvmPlatform) return
-        }
+
+        container.useInstance(AndroidExtensionPropertiesCallChecker())
+        container.useInstance(ParcelableDeclarationChecker())
     }
 }

plugins/noarg/noarg-cli/src/NoArgPlugin.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.CompilerConfigurationKey
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.noarg.NoArgCommandLineProcessor.Companion.SUPPORTED_PRESETS
 import org.jetbrains.kotlin.noarg.NoArgConfigurationKeys.ANNOTATION
@@ -89,9 +90,11 @@ class NoArgComponentRegistrar : ComponentRegistrar {
 }
 
 class CliNoArgComponentContainerContributor(val annotations: List<String>) : StorageComponentContainerContributor {
-    override fun addDeclarations(container: StorageComponentContainer, platform: TargetPlatform) {
+    override fun registerModuleComponents(
-        if (platform is JvmPlatform) {
+            container: StorageComponentContainer, platform: TargetPlatform, moduleDescriptor: ModuleDescriptor
-            container.useInstance(CliNoArgDeclarationChecker(annotations))
+    ) {
-        }
+        if (platform != JvmPlatform) return
+
+        container.useInstance(CliNoArgDeclarationChecker(annotations))
     }
 }

plugins/noarg/noarg-ide/src/IdeStorageComponentContainerContributor.kt

@@ -19,14 +19,17 @@ package org.jetbrains.kotlin.noarg.ide
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.resolve.TargetPlatform
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatform
 
 class IdeNoArgComponentContainerContributor(val project: Project) : StorageComponentContainerContributor {
-    override fun addDeclarations(container: StorageComponentContainer, platform: TargetPlatform) {
+    override fun registerModuleComponents(
-        if (platform is JvmPlatform) {
+            container: StorageComponentContainer, platform: TargetPlatform, moduleDescriptor: ModuleDescriptor
-            container.useInstance(IdeNoArgDeclarationChecker(project))
+    ) {
-        }
+        if (platform != JvmPlatform) return
+
+        container.useInstance(IdeNoArgDeclarationChecker(project))
     }
 }

plugins/sam-with-receiver/sam-with-receiver-cli/src/SamWithReceiverPlugin.kt

@@ -17,20 +17,21 @@
 package org.jetbrains.kotlin.samWithReceiver
 
 import com.intellij.mock.MockProject
-import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.compiler.plugin.CliOption
 import org.jetbrains.kotlin.compiler.plugin.CliOptionProcessingException
 import org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.CompilerConfigurationKey
-import org.jetbrains.kotlin.container.ComponentProvider
+import org.jetbrains.kotlin.container.StorageComponentContainer
-import org.jetbrains.kotlin.container.get
+import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
-import org.jetbrains.kotlin.load.java.sam.SamWithReceiverResolver
+import org.jetbrains.kotlin.resolve.TargetPlatform
+import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatform
+import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.SUPPORTED_PRESETS
 import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverConfigurationKeys.ANNOTATION
 import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverConfigurationKeys.PRESET
-import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.SUPPORTED_PRESETS
 
 object SamWithReceiverConfigurationKeys {
     val ANNOTATION: CompilerConfigurationKey<List<String>> =
@@ -75,7 +76,9 @@ class SamWithReceiverComponentRegistrar : ComponentRegistrar {
 }
 
 class CliSamWithReceiverComponentContributor(val annotations: List<String>): StorageComponentContainerContributor {
-    override fun onContainerComposed(container: ComponentProvider, moduleInfo: ModuleInfo?) {
+    override fun registerModuleComponents(container: StorageComponentContainer, platform: TargetPlatform, moduleDescriptor: ModuleDescriptor) {
-        container.get<SamWithReceiverResolver>().registerExtension(SamWithReceiverResolverExtension(annotations))
+        if (platform != JvmPlatform) return
+
+        container.useInstance(SamWithReceiverResolverExtension(annotations))
     }
 }

plugins/sam-with-receiver/sam-with-receiver-cli/src/SamWithReceiverResolverExtension.kt

@@ -22,7 +22,9 @@ import org.jetbrains.kotlin.extensions.AnnotationBasedExtension
 import org.jetbrains.kotlin.load.java.sam.SamWithReceiverResolver
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 
-class SamWithReceiverResolverExtension(val annotations: List<String>) : SamWithReceiverResolver.Extension, AnnotationBasedExtension {
+class SamWithReceiverResolverExtension(
+        private val annotations: List<String>
+) : SamWithReceiverResolver, AnnotationBasedExtension {
     override fun getAnnotationFqNames(modifierListOwner: KtModifierListOwner?) = annotations
 
     override fun shouldConvertFirstSamParameterToReceiver(function: FunctionDescriptor): Boolean {

plugins/sam-with-receiver/sam-with-receiver-ide/src/IdeSamWithReceiverComponentContributor.kt

@@ -19,20 +19,22 @@ package org.jetbrains.kotlin.samWithReceiver.ide
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootModificationTracker
-import com.intellij.psi.util.CachedValueProvider.*
+import com.intellij.psi.util.CachedValueProvider.Result
 import com.intellij.psi.util.CachedValuesManager
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.annotation.plugin.ide.getSpecialAnnotations
-import org.jetbrains.kotlin.container.ComponentProvider
+import org.jetbrains.kotlin.container.StorageComponentContainer
-import org.jetbrains.kotlin.container.get
+import org.jetbrains.kotlin.container.useInstance
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.idea.caches.resolve.ModuleProductionSourceInfo
 import org.jetbrains.kotlin.idea.caches.resolve.ScriptDependenciesModuleInfo
 import org.jetbrains.kotlin.idea.caches.resolve.ScriptModuleInfo
-import org.jetbrains.kotlin.load.java.sam.SamWithReceiverResolver
+import org.jetbrains.kotlin.resolve.TargetPlatform
-import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverResolverExtension
+import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatform
-import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.PLUGIN_ID
 import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.ANNOTATION_OPTION
+import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverCommandLineProcessor.Companion.PLUGIN_ID
+import org.jetbrains.kotlin.samWithReceiver.SamWithReceiverResolverExtension
 import java.util.*
 
 class IdeSamWithReceiverComponentContributor(val project: Project) : StorageComponentContainerContributor {
@@ -48,14 +50,23 @@ class IdeSamWithReceiverComponentContributor(val project: Project) : StorageComp
         return cache.value.getOrPut(module) { module.getSpecialAnnotations(ANNOTATION_OPTION_PREFIX) }
     }
 
-    override fun onContainerComposed(container: ComponentProvider, moduleInfo: ModuleInfo?) {
-        val annotations = when (moduleInfo) {
-            is ScriptModuleInfo -> moduleInfo.scriptDefinition.annotationsForSamWithReceivers
-            is ScriptDependenciesModuleInfo -> moduleInfo.scriptModuleInfo?.scriptDefinition?.annotationsForSamWithReceivers
-            is ModuleProductionSourceInfo -> getAnnotationsForModule(moduleInfo.module)
-            else -> null
-        } ?: return
 
-        container.get<SamWithReceiverResolver>().registerExtension(SamWithReceiverResolverExtension(annotations))
+    override fun registerModuleComponents(
+            container: StorageComponentContainer,
+            platform: TargetPlatform,
+            moduleDescriptor: ModuleDescriptor
+    ) {
+        if (platform != JvmPlatform) return
+
+        val moduleInfo = moduleDescriptor.getCapability(ModuleInfo.Capability)
+        val annotations =
+                when (moduleInfo) {
+                    is ScriptModuleInfo -> moduleInfo.scriptDefinition.annotationsForSamWithReceivers
+                    is ScriptDependenciesModuleInfo -> moduleInfo.scriptModuleInfo?.scriptDefinition?.annotationsForSamWithReceivers
+                    is ModuleProductionSourceInfo -> getAnnotationsForModule(moduleInfo.module)
+                    else -> null
+                } ?: return
+
+        container.useInstance(SamWithReceiverResolverExtension(annotations))
     }
 }