Handle local version variables in UpgradePluginVersion for kotlin plugins (#6746)

Author: MBoegers

rewrite-gradle/src/main/java/org/openrewrite/gradle/plugins/UpgradePluginVersion.java

@@ -133,15 +133,30 @@ public TreeVisitor<?, ExecutionContext> getScanner(DependencyVersionState acc) {
             @Nullable
             private GradleSettings gradleSettings;
 
+            private final Map<String, String> localVariableValues = new HashMap<>();
+
             @Override
             public @Nullable J visit(@Nullable Tree tree, ExecutionContext ctx) {
                 if (tree instanceof SourceFile) {
                     gradleProject = tree.getMarkers().findFirst(GradleProject.class).orElse(null);
                     gradleSettings = tree.getMarkers().findFirst(GradleSettings.class).orElse(null);
+                    localVariableValues.clear();
                 }
                 return super.visit(tree, ctx);
             }
 
+            @Override
+            public J visitVariable(J.VariableDeclarations.NamedVariable variable, ExecutionContext ctx) {
+                J.VariableDeclarations.NamedVariable v = (J.VariableDeclarations.NamedVariable) super.visitVariable(variable, ctx);
+                if (v.getInitializer() instanceof J.Literal) {
+                    String value = literalValue(v.getInitializer());
+                    if (value != null) {
+                        localVariableValues.put(v.getSimpleName(), value);
+                    }
+                }
+                return v;
+            }
+
             @Override
             public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
                 J.MethodInvocation m = (J.MethodInvocation) super.visitMethodInvocation(method, ctx);
@@ -185,20 +200,34 @@ public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx)
 
                         G.GString.Value gStringValue = (G.GString.Value) gString.getStrings().get(0);
                         String versionVariableName = gStringValue.getTree().toString();
+                        String resolverPluginId = pluginId;
+                        if ("kotlin".equals(pluginId)) {
+                            resolverPluginId = String.format("org.jetbrains.%s.%s", pluginId, literalValue(pluginArgs.get(0)));
+                        }
                         String resolvedPluginVersion = new DependencyVersionSelector(metadataFailures, gradleProject, gradleSettings)
-                                .select(new GroupArtifact(pluginId, pluginId + ".gradle.plugin"), "classpath", newVersion, versionPattern, ctx);
+                                .select(new GroupArtifact(resolverPluginId, resolverPluginId + ".gradle.plugin"), "classpath", newVersion, versionPattern, ctx);
 
                         acc.versionPropNameToPluginId.put(versionVariableName, pluginId);
                         assert resolvedPluginVersion != null;
                         acc.pluginIdToNewVersion.put(pluginId, resolvedPluginVersion);
                     } else if (versionArgs.get(0) instanceof J.Identifier) {
                         J.Identifier identifier = (J.Identifier) versionArgs.get(0);
                         String versionVariableName = identifier.getSimpleName();
-                        String resolvedPluginVersion = new DependencyVersionSelector(metadataFailures, gradleProject, gradleSettings)
-                                .select(new GroupArtifact(pluginId, pluginId + ".gradle.plugin"), "classpath", newVersion, versionPattern, ctx);
+                        String resolverPluginId = pluginId;
+                        if ("kotlin".equals(pluginId)) {
+                            resolverPluginId = String.format("org.jetbrains.%s.%s", pluginId, literalValue(pluginArgs.get(0)));
+                        }
+                        String localCurrentVersion = localVariableValues.get(versionVariableName);
+                        String resolvedPluginVersion;
+                        if (localCurrentVersion != null) {
+                            resolvedPluginVersion = new DependencyVersionSelector(metadataFailures, gradleProject, gradleSettings)
+                                    .select(new GroupArtifactVersion(resolverPluginId, resolverPluginId + ".gradle.plugin", localCurrentVersion), "classpath", newVersion, versionPattern, ctx);
+                        } else {
+                            resolvedPluginVersion = new DependencyVersionSelector(metadataFailures, gradleProject, gradleSettings)
+                                    .select(new GroupArtifact(resolverPluginId, resolverPluginId + ".gradle.plugin"), "classpath", newVersion, versionPattern, ctx);
+                        }
 
                         acc.versionPropNameToPluginId.put(versionVariableName, pluginId);
-                        assert resolvedPluginVersion != null;
                         acc.pluginIdToNewVersion.put(pluginId, resolvedPluginVersion);
                     }
                 } catch (MavenDownloadingException e) {

rewrite-gradle/src/test/java/org/openrewrite/gradle/plugins/UpgradePluginVersionTest.java

@@ -98,6 +98,66 @@ void upgradeKotlinPlugin() {
         );
     }
 
+    @Test
+    void upgradeKotlinPluginLocalVariable() {
+        rewriteRun(
+          spec -> spec.recipe(new UpgradePluginVersion("kotlin", "2.1.0", null)),
+          buildGradleKts(
+            """
+              plugins {
+                  val kotlinVersion = "1.9.25"
+                  kotlin("jvm") version kotlinVersion
+                  kotlin("plugin.allopen") version kotlinVersion
+                  kotlin("plugin.spring") version kotlinVersion
+              }
+              """,
+            """
+              plugins {
+                  val kotlinVersion = "2.1.0"
+                  kotlin("jvm") version kotlinVersion
+                  kotlin("plugin.allopen") version kotlinVersion
+                  kotlin("plugin.spring") version kotlinVersion
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void upgradeKotlinPluginLocalVariableWithSemverSelector() {
+        rewriteRun(
+          spec -> spec.recipe(new UpgradePluginVersion("kotlin", "latest.minor", null)),
+          buildGradleKts(
+            """
+              plugins {
+                  val kotlinVersion = "2.0.0"
+                  kotlin("jvm") version kotlinVersion
+              }
+              """,
+            spec -> spec.after(s -> {
+                assertThat(s).doesNotContain("2.0.0");
+                assertThat(s).containsPattern("val kotlinVersion = \"2.\\d+.\\d+\"");
+                return s;
+            })
+          )
+        );
+    }
+
+    @Test
+    void dontDowngradeKotlinPluginLocalVariable() {
+        rewriteRun(
+          spec -> spec.recipe(new UpgradePluginVersion("kotlin", "1.9.0", null)),
+          buildGradleKts(
+            """
+              plugins {
+                  val kotlinVersion = "1.9.25"
+                  kotlin("jvm") version kotlinVersion
+              }
+              """
+          )
+        );
+    }
+
     @Test
     void upgradeGradleSettingsPlugin() {
         rewriteRun(