IDEA-CR-51057: [Gradle] populate build level models API cleanup

(cherry picked from commit d632a0eda809f61ab3bb99bf4bcedf42a7ca794e)

GitOrigin-RevId: b24e49e2f8c6a4ef2ca0d90b6e21901a1276c4bc

plugins/gradle/testSources/org/jetbrains/plugins/gradle/importing/GradleActionWithImportTest.kt

@@ -7,6 +7,7 @@ import com.intellij.util.ThrowableRunnable
 import org.assertj.core.api.Assertions.assertThat
 import org.gradle.tooling.BuildController
 import org.gradle.tooling.GradleConnectionException
+import org.gradle.tooling.model.BuildModel
 import org.gradle.tooling.model.Model
 import org.jetbrains.plugins.gradle.GradleManager
 import org.jetbrains.plugins.gradle.model.ProjectImportModelProvider
@@ -94,10 +95,9 @@ class GradleActionWithImportTest: GradleImportingTestCase() {
 }
 
 class TestModelProvider : ProjectImportModelProvider {
-
-  override fun populateBuildModels(controller: BuildController,
-                                   project: Model,
-                                   consumer: ProjectImportModelProvider.BuildModelConsumer) {
+  override fun <T> populateBuildModels(controller: BuildController,
+                                       buildModel: T,
+                                       consumer: ProjectImportModelProvider.BuildModelConsumer) where T : Model?, T : BuildModel? {
     controller.findModel(Object::class.java)
   }
 

plugins/gradle/tooling-extension-api/src/org/jetbrains/plugins/gradle/model/ProjectImportModelProvider.java

@@ -4,9 +4,9 @@ package org.jetbrains.plugins.gradle.model;
 
 import org.gradle.tooling.BuildController;
 import org.gradle.tooling.model.BuildModel;
+import org.gradle.tooling.model.GradleProject;
 import org.gradle.tooling.model.Model;
-import org.gradle.tooling.model.idea.IdeaModule;
-import org.gradle.tooling.model.idea.IdeaProject;
+import org.gradle.tooling.model.gradle.GradleBuild;
 import org.jetbrains.annotations.NotNull;
 
 import java.io.Serializable;
@@ -14,10 +14,10 @@ import java.io.Serializable;
 /**
  * Allows the {@link ProjectImportAction} to be extended to allow extra flexibility to extensions when requesting the models.
  *
- * {@link #populateProjectModels(BuildController, Model, ProjectModelConsumer)} is called once for each {@link IdeaModule} obtained
- * from the Gradle Tooling API (this includes modules from included builds).
+ * {@link #populateProjectModels(BuildController, Model, ProjectModelConsumer)} is called once for each {@link GradleProject} obtained
+ * from the Gradle Tooling API (this includes projects from included builds).
  *
- * {@link #populateBuildModels(BuildController, Model, BuildModelConsumer)} is called once for each {@link IdeaProject} that is
+ * {@link #populateBuildModels(BuildController, Model, BuildModelConsumer)} is called once for each {@link GradleBuild} that is
  * obtained from the Gradle Tooling API, for none-composite builds this will be called exactly once, for composite builds this will be
  * called once for each included build and once for the name build. This will always be called after
  * {@link #populateProjectModels(BuildController, Model, ProjectModelConsumer)}.
@@ -31,7 +31,9 @@ public interface ProjectImportModelProvider extends Serializable {
     void consume(@NotNull BuildModel buildModel, @NotNull Object object, @NotNull Class clazz);
   }
 
-  void populateBuildModels(@NotNull BuildController controller, @NotNull Model buildModel, @NotNull BuildModelConsumer consumer);
+  <T extends Model & BuildModel> void populateBuildModels(@NotNull BuildController controller,
+                                                          @NotNull T buildModel,
+                                                          @NotNull BuildModelConsumer consumer);
 
   void populateProjectModels(@NotNull BuildController controller,
                              @NotNull Model projectModel,

plugins/gradle/tooling-extension-impl/src/org/jetbrains/plugins/gradle/model/ClassSetBuildImportModelProvider.java

+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package org.jetbrains.plugins.gradle.model;
+
+import org.gradle.tooling.BuildController;
+import org.gradle.tooling.model.BuildModel;
+import org.gradle.tooling.model.Model;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+public final class ClassSetBuildImportModelProvider implements ProjectImportModelProvider {
+  @NotNull private final Set<Class> classSet;
+
+  public ClassSetBuildImportModelProvider(@NotNull Collection<Class> classes) {
+    classSet = new LinkedHashSet<Class>(classes);
+  }
+
+  @Override
+  public <T extends Model & BuildModel> void populateBuildModels(@NotNull BuildController controller,
+                                                                 @NotNull T buildModel,
+                                                                 @NotNull BuildModelConsumer consumer) {
+    for (Class<?> aClass : classSet) {
+      Object instance = controller.findModel(buildModel, aClass);
+      if (instance != null) {
+        consumer.consume(buildModel, instance, aClass);
+      }
+    }
+  }
+
+  @Override
+  public void populateProjectModels(@NotNull BuildController controller,
+                                    @Nullable Model module,
+                                    @NotNull ProjectModelConsumer modelConsumer) {
+    // Do nothing, this provider only works on the build model level
+  }
+}