[KTIJ-24365] KotlinMppGradleProjectResolverExtension: Add extension points for...

[KTIJ-24365] KotlinMppGradleProjectResolverExtension: Add extension points for populating dependencies

(cherry picked from commit 9184ab71e49358c005e222c67d9044cad66309d9)

GitOrigin-RevId: 36b284154e3ee71c232080c72f6a76459163f655

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/KotlinMppGradleProjectResolverExtension.kt

@@ -4,8 +4,10 @@ package org.jetbrains.kotlin.idea.gradleJava.configuration.mpp
 
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.externalSystem.model.DataNode
+import com.intellij.openapi.externalSystem.model.project.AbstractDependencyData
 import com.intellij.openapi.externalSystem.model.project.ModuleData
 import org.gradle.tooling.model.idea.IdeaModule
+import org.jetbrains.kotlin.gradle.idea.tcs.IdeaKotlinDependency
 import org.jetbrains.kotlin.idea.gradleTooling.KotlinMPPGradleModel
 import org.jetbrains.kotlin.idea.projectModel.KotlinComponent
 import org.jetbrains.kotlin.idea.projectModel.KotlinSourceSet
@@ -46,5 +48,20 @@ interface KotlinMppGradleProjectResolverExtension {
         context: Context, sourceSetDataNode: DataNode<GradleSourceSetData>, sourceSet: KotlinSourceSet
     ) = Unit
 
-    fun provideAdditionalProjectArtifactDependencyResolver(): KotlinProjectArtifactDependencyResolver? = null
+    fun beforePopulateSourceSetDependencies(
+        context: Context,
+        sourceSetDataNode: DataNode<GradleSourceSetData>,
+        sourceSet: KotlinSourceSet,
+        dependencies: Set<IdeaKotlinDependency>
+    ): Result = Result.Proceed
+
+    fun afterPopulateSourceSetDependencies(
+        context: Context,
+        sourceSetDataNode: DataNode<GradleSourceSetData>,
+        sourceSet: KotlinSourceSet,
+        dependencies: Set<IdeaKotlinDependency>,
+        dependencyNodes: List<DataNode<out AbstractDependencyData<*>>>
+    ) = Unit
+
+    fun provideAdditionalProjectArtifactDependencyResolvers(): List<KotlinProjectArtifactDependencyResolver> = emptyList()
 }

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/KotlinMppGradleProjectResolverExtensionContextImpl.kt

@@ -7,9 +7,18 @@ import org.gradle.tooling.model.idea.IdeaModule
 import org.jetbrains.kotlin.idea.gradleTooling.KotlinMPPGradleModel
 import org.jetbrains.plugins.gradle.service.project.ProjectResolverContext
 
-internal class KotlinMppGradleProjectResolverExtensionContextImpl(
-  override val model: KotlinMPPGradleModel,
-  override val resolverCtx: ProjectResolverContext,
-  override val gradleModule: IdeaModule,
-  override val moduleDataNode: DataNode<ModuleData>
-) : KotlinMppGradleProjectResolverExtension.Context
\ No newline at end of file
+fun KotlinMppGradleProjectResolverExtension.Companion.Context(
+    model: KotlinMPPGradleModel,
+    resolverCtx: ProjectResolverContext,
+    gradleModule: IdeaModule,
+    moduleDataNode: DataNode<ModuleData>
+): KotlinMppGradleProjectResolverExtension.Context = KotlinMppGradleProjectResolverExtensionContextImpl(
+    model, resolverCtx, gradleModule, moduleDataNode
+)
+
+private class KotlinMppGradleProjectResolverExtensionContextImpl(
+    override val model: KotlinMPPGradleModel,
+    override val resolverCtx: ProjectResolverContext,
+    override val gradleModule: IdeaModule,
+    override val moduleDataNode: DataNode<ModuleData>
+) : KotlinMppGradleProjectResolverExtension.Context

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/KotlinMppGradleProjectResolverExtensionInstance.kt

+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.kotlin.idea.gradleJava.configuration.mpp
+
+import com.intellij.openapi.externalSystem.model.DataNode
+import com.intellij.openapi.externalSystem.model.project.AbstractDependencyData
+import org.jetbrains.kotlin.gradle.idea.tcs.IdeaKotlinDependency
+import org.jetbrains.kotlin.idea.gradleJava.configuration.mpp.KotlinMppGradleProjectResolverExtension.Result
+import org.jetbrains.kotlin.idea.projectModel.KotlinComponent
+import org.jetbrains.kotlin.idea.projectModel.KotlinSourceSet
+import org.jetbrains.plugins.gradle.model.data.GradleSourceSetData
+
+internal val KotlinMppGradleProjectResolverExtension.Companion.instance: KotlinMppGradleProjectResolverExtension
+    get() = KotlinMppGradleProjectResolverExtensionInstance(EP_NAME.extensionList)
+
+private class KotlinMppGradleProjectResolverExtensionInstance(
+    private val extensions: List<KotlinMppGradleProjectResolverExtension>
+) : KotlinMppGradleProjectResolverExtension {
+
+    override fun beforeMppGradleSourceSetDataNodeCreation(
+        context: KotlinMppGradleProjectResolverExtension.Context,
+        component: KotlinComponent
+    ): Result {
+        return extensions.map { extension -> extension.beforeMppGradleSourceSetDataNodeCreation(context, component) }.reduce()
+    }
+
+    override fun afterMppGradleSourceSetDataNodeCreated(
+        context: KotlinMppGradleProjectResolverExtension.Context,
+        component: KotlinComponent,
+        sourceSetDataNode: DataNode<GradleSourceSetData>
+    ) {
+        extensions.forEach { extension -> extension.afterMppGradleSourceSetDataNodeCreated(context, component, sourceSetDataNode) }
+    }
+
+    override fun beforePopulateContentRoots(
+        context: KotlinMppGradleProjectResolverExtension.Context,
+        sourceSetDataNode: DataNode<GradleSourceSetData>,
+        sourceSet: KotlinSourceSet
+    ): Result {
+        return extensions.map { extension -> extension.beforePopulateContentRoots(context, sourceSetDataNode, sourceSet) }.reduce()
+    }
+
+    override fun afterPopulateContentRoots(
+        context: KotlinMppGradleProjectResolverExtension.Context,
+        sourceSetDataNode: DataNode<GradleSourceSetData>,
+        sourceSet: KotlinSourceSet
+    ) {
+        extensions.forEach { extension -> extension.afterPopulateContentRoots(context, sourceSetDataNode, sourceSet) }
+    }
+
+    override fun beforePopulateSourceSetDependencies(
+        context: KotlinMppGradleProjectResolverExtension.Context,
+        sourceSetDataNode: DataNode<GradleSourceSetData>,
+        sourceSet: KotlinSourceSet,
+        dependencies: Set<IdeaKotlinDependency>
+    ): Result = extensions
+        .map { extension -> extension.beforePopulateSourceSetDependencies(context, sourceSetDataNode, sourceSet, dependencies) }.reduce()
+
+    override fun afterPopulateSourceSetDependencies(
+        context: KotlinMppGradleProjectResolverExtension.Context,
+        sourceSetDataNode: DataNode<GradleSourceSetData>,
+        sourceSet: KotlinSourceSet,
+        dependencies: Set<IdeaKotlinDependency>,
+        dependencyNodes: List<DataNode<out AbstractDependencyData<*>>>
+    ) {
+        extensions.forEach { extension ->
+            extension.afterPopulateSourceSetDependencies(context, sourceSetDataNode, sourceSet, dependencies, dependencyNodes)
+        }
+    }
+
+    override fun provideAdditionalProjectArtifactDependencyResolvers(): List<KotlinProjectArtifactDependencyResolver> {
+        return extensions.flatMap { extension -> extension.provideAdditionalProjectArtifactDependencyResolvers() }
+    }
+}
+
+private operator fun Result.plus(other: Result) = if (this == Result.Skip) this else other
+
+private fun Iterable<Result>.reduce() = reduceOrNull { acc, result -> acc + result } ?: Result.Proceed
\ No newline at end of file

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/KotlinProjectArtifactDependencyResolver.kt

@@ -43,7 +43,7 @@ private class KotlinProjectArtifactDependencyResolverImpl : KotlinProjectArtifac
         dependency: IdeaKotlinProjectArtifactDependency
     ): Set<IdeaKotlinSourceDependency> {
         val resolvedByExtensions = projectNode.getUserData(KotlinProjectArtifactDependencyResolver.key).orEmpty()
-            .plus(KotlinMppGradleProjectResolverExtension.provideAdditionalProjectArtifactDependencyResolvers())
+            .plus(KotlinMppGradleProjectResolverExtension.instance.provideAdditionalProjectArtifactDependencyResolvers())
             .flatMap { resolver -> resolver.resolve(projectNode, sourceSetNode, dependency) }
             .toSet()
 

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/kotlinMppGradleProjectResolverExtensionUtils.kt

-// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-
-package org.jetbrains.kotlin.idea.gradleJava.configuration.mpp
-
-import com.intellij.openapi.externalSystem.model.DataNode
-import com.intellij.openapi.externalSystem.model.project.ModuleData
-import org.gradle.tooling.model.idea.IdeaModule
-import org.jetbrains.kotlin.idea.gradleJava.configuration.mpp.KotlinMppGradleProjectResolverExtension.Result
-import org.jetbrains.kotlin.idea.gradleTooling.KotlinMPPGradleModel
-import org.jetbrains.kotlin.idea.projectModel.KotlinComponent
-import org.jetbrains.kotlin.idea.projectModel.KotlinSourceSet
-import org.jetbrains.plugins.gradle.model.data.GradleSourceSetData
-import org.jetbrains.plugins.gradle.service.project.ProjectResolverContext
-
-internal fun KotlinMppGradleProjectResolverExtension.Companion.beforeMppGradleSourceSetDataNodeCreation(
-    model: KotlinMPPGradleModel,
-    resolverCtx: ProjectResolverContext,
-    gradleModule: IdeaModule,
-    moduleDataNode: DataNode<ModuleData>,
-    component: KotlinComponent
-): Result {
-    val context = KotlinMppGradleProjectResolverExtensionContextImpl(model, resolverCtx, gradleModule, moduleDataNode)
-    return EP_NAME.extensionList.fold(Result.Proceed) { result, extension ->
-        val nextResult = extension.beforeMppGradleSourceSetDataNodeCreation(context, component)
-        if (result == Result.Skip || nextResult == Result.Skip) Result.Skip
-        else Result.Proceed
-    }
-}
-
-internal fun KotlinMppGradleProjectResolverExtension.Companion.afterMppGradleSourceSetDataNodeCreated(
-    model: KotlinMPPGradleModel,
-    resolverCtx: ProjectResolverContext,
-    gradleModule: IdeaModule,
-    moduleDataNode: DataNode<ModuleData>,
-    component: KotlinComponent,
-    sourceSetDataNode: DataNode<GradleSourceSetData>
-) {
-    val context = KotlinMppGradleProjectResolverExtensionContextImpl(model, resolverCtx, gradleModule, moduleDataNode)
-    return EP_NAME.extensionList.forEach { extension ->
-        extension.afterMppGradleSourceSetDataNodeCreated(context, component, sourceSetDataNode)
-    }
-}
-
-internal fun KotlinMppGradleProjectResolverExtension.Companion.beforePopulateContentRoots(
-    model: KotlinMPPGradleModel,
-    resolverCtx: ProjectResolverContext,
-    gradleModule: IdeaModule,
-    moduleDataNode: DataNode<ModuleData>,
-    sourceSetDataNode: DataNode<GradleSourceSetData>,
-    sourceSet: KotlinSourceSet
-): Result {
-    val context = KotlinMppGradleProjectResolverExtensionContextImpl(model, resolverCtx, gradleModule, moduleDataNode)
-
-    return EP_NAME.extensionList.fold(Result.Proceed) { result, extension ->
-        val nextResult = extension.beforePopulateContentRoots(context, sourceSetDataNode, sourceSet)
-        if (result == Result.Skip || nextResult == Result.Skip) Result.Skip
-        else Result.Proceed
-    }
-}
-
-internal fun KotlinMppGradleProjectResolverExtension.Companion.afterPopulateContentRoots(
-    model: KotlinMPPGradleModel,
-    resolverCtx: ProjectResolverContext,
-    gradleModule: IdeaModule,
-    moduleDataNode: DataNode<ModuleData>,
-    sourceSetDataNode: DataNode<GradleSourceSetData>,
-    sourceSet: KotlinSourceSet
-) {
-    val context = KotlinMppGradleProjectResolverExtensionContextImpl(model, resolverCtx, gradleModule, moduleDataNode)
-    EP_NAME.extensionList.forEach { extension ->
-        extension.afterPopulateContentRoots(context, sourceSetDataNode, sourceSet)
-    }
-}
-
-internal fun KotlinMppGradleProjectResolverExtension.Companion.provideAdditionalProjectArtifactDependencyResolvers()
-        : List<KotlinProjectArtifactDependencyResolver> = EP_NAME.extensionList.mapNotNull { extension ->
-    extension.provideAdditionalProjectArtifactDependencyResolver()
-}

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/moduleContentRoots.kt

@@ -31,6 +31,10 @@ internal fun populateContentRoots(
     resolverCtx: ProjectResolverContext
 ) {
     val mppModel = resolverCtx.getMppModel(gradleModule) ?: return
+
+    val extensionContext = KotlinMppGradleProjectResolverExtension.Context(mppModel, resolverCtx, gradleModule, ideModule)
+    val extensionInstance = KotlinMppGradleProjectResolverExtension.instance
+
     val sourceSetToPackagePrefix = mppModel.targets.flatMap { it.compilations }
         .flatMap { compilation ->
             compilation.declaredSourceSets.map { sourceSet -> sourceSet.name to compilation.kotlinTaskProperties.packagePrefix }
@@ -41,10 +45,9 @@ internal fun populateContentRoots(
         if (dataNode == null || shouldDelegateToOtherPlugin(sourceSet)) return@processSourceSets
 
         /* Execute all registered extension points and skip population of content roots if instructed by extensions */
-        if (KotlinMppGradleProjectResolverExtension.beforePopulateContentRoots(
-                mppModel, resolverCtx, gradleModule, ideModule, dataNode, sourceSet
-            ) == Skip
-        ) return@processSourceSets
+        if (extensionInstance.beforePopulateContentRoots(extensionContext, dataNode, sourceSet) == Skip) {
+            return@processSourceSets
+        }
 
         createContentRootData(
             sourceSet.sourceDirs,
@@ -59,9 +62,7 @@ internal fun populateContentRoots(
             dataNode
         )
 
-        KotlinMppGradleProjectResolverExtension.afterPopulateContentRoots(
-            mppModel, resolverCtx, gradleModule, ideModule, dataNode, sourceSet
-        )
+        extensionInstance.afterPopulateContentRoots(extensionContext, dataNode, sourceSet)
     }
 
     for (gradleContentRoot in gradleModule.contentRoots ?: emptySet<IdeaContentRoot?>()) {

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/populateModuleDependencies.kt

@@ -45,29 +45,43 @@ internal fun populateModuleDependenciesWithDependenciesContainer(
 ) {
     mppModel.dependencyMap.values.modifyDependenciesOnMppModules(ideProject)
 
+    val extensionContext = KotlinMppGradleProjectResolverExtension.Context(mppModel, resolverCtx, gradleModule, ideModule)
+    val extension = KotlinMppGradleProjectResolverExtension.instance
+
     mppModel.sourceSetsByName.values.forEach { sourceSet ->
         val sourceSetModuleIde = KotlinSourceSetModuleId(resolverCtx, gradleModule, sourceSet)
         val sourceSetDataNode = ideModule.findSourceSetNode(sourceSetModuleIde) ?: return@forEach
+        val sourceSetDependencies = dependencies[sourceSet.name]
+
+        /* Call into extension points, skipping dependency population of source set if instructed */
+        if (
+            extension.beforePopulateSourceSetDependencies(
+                extensionContext, sourceSetDataNode, sourceSet, sourceSetDependencies
+            ) == KotlinMppGradleProjectResolverExtension.Result.Skip
+        ) return@forEach
 
         /*
         Some dependencies are represented as IdeaKotlinProjectArtifactDependency.
         Such dependencies can be resolved to the actual source sets that built this artifact.
          */
         val projectArtifactDependencyResolver = KotlinProjectArtifactDependencyResolver()
-        val resolvedDependencies = dependencies[sourceSet.name].flatMap { dependency ->
+        val resolvedSourceSetDependencies = sourceSetDependencies.flatMap { dependency ->
             if (dependency is IdeaKotlinProjectArtifactDependency)
                 projectArtifactDependencyResolver.resolve(ideProject, sourceSetDataNode, dependency)
             else listOf(dependency)
-        }
+        }.toSet()
 
-        /*
-        Add each resolved dependency
-         */
-        resolvedDependencies.forEachIndexed { index, dependency ->
+        /* Add each resolved dependency */
+        val createdDependencyNodes = resolvedSourceSetDependencies.flatMapIndexed { index, dependency ->
             sourceSetDataNode.addDependency(dependency)
                 /* The classpath order of the dependencies is given by the order they were sent by the Kotlin Gradle Plugin */
-                .forEach { it.data.setOrder(index) }
+                .onEach { it.data.setOrder(index) }
         }
+
+        /* Calling into extensions, notifying them about all populated dependencies */
+        extension.afterPopulateSourceSetDependencies(
+            extensionContext, sourceSetDataNode, sourceSet, resolvedSourceSetDependencies, createdDependencyNodes
+        )
     }
 }
 

plugins/kotlin/gradle/gradle-java/src/org/jetbrains/kotlin/idea/gradleJava/configuration/mpp/populateMppModuleDataNode.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.idea.gradle.configuration.utils.predictedProductionS
 import org.jetbrains.kotlin.idea.gradleJava.configuration.*
 import org.jetbrains.kotlin.idea.gradleJava.configuration.KotlinMPPGradleProjectResolver.Companion.resourceType
 import org.jetbrains.kotlin.idea.gradleJava.configuration.KotlinMPPGradleProjectResolver.Companion.sourceType
+import org.jetbrains.kotlin.idea.gradleJava.configuration.mpp.KotlinMppGradleProjectResolverExtension.Result.Skip
 import org.jetbrains.kotlin.idea.gradleJava.configuration.utils.KotlinModuleUtils
 import org.jetbrains.kotlin.idea.gradleJava.configuration.utils.KotlinModuleUtils.fullName
 import org.jetbrains.kotlin.idea.gradleTooling.KotlinCompilationImpl
@@ -261,6 +262,9 @@ private fun createMppGradleSourceSetDataNodes(
     val mainModuleConfigPath = mainModuleData.linkedExternalProjectPath
     val mainModuleFileDirectoryPath = mainModuleData.moduleFileDirectoryPath
 
+    val extensionContext = KotlinMppGradleProjectResolverExtension.Context(mppModel, resolverCtx, gradleModule, mainModuleNode)
+    val extensionInstance = KotlinMppGradleProjectResolverExtension.instance
+
     val externalProject = resolverCtx.getExtraProject(gradleModule, ExternalProject::class.java) ?: return
     val projectDataNode = ExternalSystemApiUtil.findParent(mainModuleNode, ProjectKeys.PROJECT) ?: return
 
@@ -293,10 +297,9 @@ private fun createMppGradleSourceSetDataNodes(
             if (existingSourceSetDataNode?.kotlinSourceSetData?.sourceSetInfo != null) continue
 
             /* Execute extensions and do not create any GradleSourceSetData node if any extension wants us to Skip */
-            if (KotlinMppGradleProjectResolverExtension.beforeMppGradleSourceSetDataNodeCreation(
-                    mppModel, resolverCtx, gradleModule, mainModuleNode, compilation
-                ) == KotlinMppGradleProjectResolverExtension.Result.Skip
-            ) continue
+            if (extensionInstance.beforeMppGradleSourceSetDataNodeCreation(extensionContext, compilation) == Skip) {
+                continue
+            }
 
             compilationIds.add(moduleId)
 
@@ -355,9 +358,7 @@ private fun createMppGradleSourceSetDataNodes(
             }
 
             /* Execution all extensions after we freshly created a GradleSourceSetData node for the given compilation */
-            KotlinMppGradleProjectResolverExtension.afterMppGradleSourceSetDataNodeCreated(
-                mppModel, resolverCtx, gradleModule, mainModuleNode, compilation, compilationDataNode
-            )
+            extensionInstance.afterMppGradleSourceSetDataNodeCreated(extensionContext, compilation, compilationDataNode)
         }
 
         targetData.moduleIds = compilationIds
@@ -374,10 +375,9 @@ private fun createMppGradleSourceSetDataNodes(
         if (existingSourceSetDataNode?.kotlinSourceSetData != null) continue
 
         /* Execute extensions and do not create any GradleSourceSetData node if any extension wants us to Skip */
-        if (KotlinMppGradleProjectResolverExtension.beforeMppGradleSourceSetDataNodeCreation(
-                mppModel, resolverCtx, gradleModule, mainModuleNode, sourceSet
-            ) == KotlinMppGradleProjectResolverExtension.Result.Skip
-        ) continue
+        if (extensionInstance.beforeMppGradleSourceSetDataNodeCreation(extensionContext, sourceSet) == Skip) {
+            continue
+        }
 
         val sourceSetData = existingSourceSetDataNode?.data ?: createGradleSourceSetData(
             sourceSet, gradleModule, mainModuleNode, resolverCtx
@@ -428,9 +428,7 @@ private fun createMppGradleSourceSetDataNodes(
         }
 
         /* Execution all extensions after we freshly created a GradleSourceSetData node for the given compilation */
-        KotlinMppGradleProjectResolverExtension.afterMppGradleSourceSetDataNodeCreated(
-            mppModel, resolverCtx, gradleModule, mainModuleNode, sourceSet, sourceSetDataNode
-        )
+        extensionInstance.afterMppGradleSourceSetDataNodeCreated(extensionContext, sourceSet, sourceSetDataNode)
     }
 }